Residual Connections Harm Generative Representation Learning

Thank you for your feedback.

• W1: Systematical analysis of the model configurations and selection of optimal $\alpha_{\text{min}}$. (also addressed in the general reply)

We provide extensive analysis across multiple model configurations and provide an empirical criterion for selecting the optimal $\alpha_{\text{min}}$.

Specifically, we evaluate different choices of $\alpha_{\text{min}}$ across two axes: feature dimension and encoder depth, and report the results in Table R1. Linear probing accuracy is measured using features extracted from the encoder's final layer.

In Table R2, we extend the analysis by reporting the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses). Across various configurations, including ViT-Base and ViT-Large, our method consistently outperforms the baseline that employs a full residual connection ($\alpha_{\text{min}} = 1$).

Our findings reveal that the number of encoder layers $L$, rather than the embedding dimension, primarily determines the optimal $\alpha_{\text{min}}$. We attribute this behavior to the scaling effect of the input data on the encoder's final layer, quantified as: $\alpha_L^{\text{eff}} = \prod_{l=1}^L(1 - \frac{(1 - \alpha_{\text{min}})l}{L})$. Deeper models require larger $\alpha_{\text{min}}$ to maintain a consistent cumulative decay effect represented by $\alpha_L^{\text{eff}}$.

Table R3 presents the values of $\alpha_L^{\text{eff}}$ for the best-performing layer under various configurations. From these results, we see that selecting $\alpha_{\text{min}}$ to make $\alpha_L^{\text{eff}} \in [0.001, 0.01)$ consistently yields a better result than the baseline model with a full residual connection. Particularly, along with Table R2, we observe that the improvements are robust when $\alpha_L^{\text{eff}}$ falls in this preferred range. Choosing $\alpha_L^{\text{eff}}$, and using the above formula, we can compute $\alpha_{\text{min}}$ for a particular network.

Table R1: Linear probing accuracy of the final layer of the encoder. Across multiple model configurations, we report the probing score when applying our method. For a family of architectures, by properly choosing $\alpha_{\text{min}}$, our method leads to a significantly better result than the baseline with the full residual connection.

Table R2: Linear probing accuracy of the best performing layer of the encoder. Across multiple model configurations, we report the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses). Comparing to Table R1 shows that the best probing score appears at shallower layers for deeper models.

Table R3: Best linear probing accuracy for various $\alpha_L^{\text{eff}}$. Results show that for different model configurations, setting $\alpha_{\text{min}}$ such that $\alpha_{L}^{\text{eff}} \in [1\text{e-3}, 1\text{e-2})$ consistently achieves the best performance.

---

• W2: Comparison to baseline methods

Our original submission contains extensive comparisons to baseline methods in Tables 2,3,4 and Figure 1. Please note that $\alpha_{\text{min}} = 1$ corresponds to the default residual connections (baseline network architecture) in each setting. The requested comparisons are all present in the original submission, and we see in every setting that our approach outperforms the baseline by using smaller values of $\alpha_{\rm min}$.

Thank you for your feedback.

• W1: The experiments on representation learning are only on MAEs.

Experiments on Diffusion models. We demonstrate that our method enhances both representation learning and image generation in diffusion models.

Specifically, Figure 1 and Table 4 of our submission illustrate improvements in representation learning with diffusion models using ViTs and CNNs: (a) For diffusion models with ViT, we increase linear probing accuracy from 72.8% to 76.1% and improve FID from 6.93 to 4.98 (conditional generation) and 44.40 to 40.96 (unconditional generation) on ImageNet-100; (b) For diffusion models with CNN, we increase linear probing accuracy from 62.47% to 66.86% with FID improving from 14.34 to 8.99.

Fine-tuning performance. (also answered in responses to Reviewers uy11 and Hwrn)

Our primary focus is on representation learning. End-to-end fine-tuning primarily evaluates the utility of features as initialization for downstream tasks, which is influenced by various factors such as augmentation strategy and layer-wise learning rates in the fine-tuning stage, rather than directly measuring the quality of the learned representations. In contrast, linear probing offers a more direct assessment of representation quality by utilizing a simple linear projection to adapt the features for downstream applications.

Moreover, the results from end-to-end fine-tuning do not accurately reflect the quality of unsupervised representations, as the fine-tuning process significantly alters the original features. For example, ADDP [1] achieves only 23.8% accuracy in linear probing, highlighting poor representation learning. However, its accuracy improves significantly to 85.9% after end-to-end fine-tuning, illustrating the substantial transformation of representations during fine-tuning.

Table R9: Comparison of fine-tuning and linear probing for MAE models on ImageNet-1k.

---

• W2: Linear probing of the inner state of different layers

Thank you for the suggestions! Here we show the probing performance of $\alpha_{\text{min}} = 0.6$ for ViT-B/16 and $\alpha_{\text{min}} = 0.7$ (the best configuration for each model) for ViT-L/16 for ImageNet-100 experiments. In the standard model, such as ViT-B, the top probing scores arise from the last layer; In a much deeper model like ViT-L, the optimal features are found in earlier layers (specifically, layer 19).

As depth increases, the optimal features are likely to reside in the earlier layers because the effective decayed weight, calculated as $\prod_{i=1}^l \alpha_i$, decreases exponentially with greater layer depth. In response to the next question, we show with additional experiments (Table R2) that the optimal rate consistently remains within a narrow range.

Table R10: Linear probing accuracy for inner layers

---

• W3: Hyper-parameter selection (also addressed in our general reply)

Our method relies on a single hyperparameter $\alpha_{\text{min}}$, resulting in a minimum modification to the generative representation learning framework, while achieving significant enhancement compared to the baseline and the related work (e.g., I-JEPA).

As shown in Table 2 of our paper, the optimal $\alpha_{\text{min}}$ is about 0.6 for both ViT-B and ViT-S models, highlighting the robustness of our technique. We provide additional results for various model configurations and an empirical guideline for optimal $\alpha_{\text{min}}$ selection. Even when increasing the encoder from 12 to 24 layers and the feature dimension from 768 to 1024, the performance of the optimal layers (see Table R2) is largely unaffected by the choice of $\alpha_{\rm{min}}$.

Specifically, we evaluate different choices of $\alpha_{\text{min}}$ across two axes: feature dimension and encoder depth, and report the results in Table R1. Linear probing accuracy is measured using features extracted from the encoder's final layer.

In Table R2, we extend the analysis by reporting the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses).

Across various configurations, including ViT-Base and ViT-Large, our method consistently outperforms the baseline that employs a full residual connection ($\alpha_{\text{min}} = 1$).

We thank you for your comment.

it's reasonable to show linear probing results but the choice of base model is questionable and undermines the point

We respectfully disagree with the statement that ``MAE is a questionable model for showing linear probing results.'' While generative-based frameworks like MAE do not currently achieve state-of-the-art linear probing accuracy compared to contrastive-based methods, they offer a unified framework for both representation learning and data generation. Improving the linear probing performance of generative-based models is a key research focus within the representation learning community. For example, methods such as [1] I-JEPA, [2] Latent-MIM, and [3] SODA target improvements in masked models or diffusion models, while we demonstrate enhancements in both frameworks.

In addition to linear probing, we also evaluate feature quality using K-NN on the ImageNet-1K dataset:

Table R12. K-Nearest Neighbor (K=20) accuracy on ImageNet-1K. Our method significantly improves upon the MAE baseline and approaches the performance of MoCo-V3.

I was hoping to see models other than MAE in this setting.

We agree! In our draft, we already show scalability across (1) MAE, (2) diffusion models (both ViT and CNN architectures), and (3) BEiT (presented in our rebuttal). Across three key metrics—(1) linear probing accuracy, (2) FID for image generation quality, and (3) K-NN accuracy—we demonstrate consistent and significant improvements.

Specifically, compared to the baseline models:

Although our fine-tuning experiments have not yet shown improvements (with a slight 0.7% decrease, likely due to suboptimal hyperparameter tuning and augmentation strategies at the fine-tuning stage), our results demonstrate significant advancements across three metrics and multiple frameworks. This highlights the scalability, robustness, and effectiveness of our method. Therefore, we would greatly appreciate it if you could evaluate our contributions more comprehensively and reconsider your rating. We are happy to address any additional concerns. Thank you!

[1] Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture, Mahmoud Assran et al., 2023

[2] Towards Latent Masked Image Modeling for Self-Supervised Visual Representation Learning, Yibing Wei, Abhinav Gupta, Pedro Morgado, 2024.

[3] SODA: Bottleneck Diffusion Models for Representation Learning, Drew A. Hudson et al., 2023

• W2: fine-tuning experiments (also answered in responses to Reviewers uy11 and Kv2k)

Our primary focus is on representation learning. End-to-end fine-tuning primarily evaluates the utility of features as initialization for downstream tasks, which is influenced by various factors such as augmentation strategy and layer-wise learning rates in the fine-tuning stage, rather than directly measuring the quality of the learned representations. In contrast, linear probing offers a more direct assessment of representation quality by utilizing a simple linear projection to adapt the features for downstream applications.

Moreover, the results from end-to-end fine-tuning do not accurately reflect the quality of unsupervised representations, as the fine-tuning process significantly alters the original features. For example, ADDP [1] achieves only 23.8% accuracy in linear probing, highlighting poor representation learning. However, its accuracy improves significantly to 85.9% after end-to-end fine-tuning, illustrating the substantial transformation of representations during fine-tuning.

Table R9: Comparison of fine-tuning and linear probing for MAE models on ImageNet-1k.

• W2: experiments with other generative models

We appreciate your recommendation to extend our methods to other generative models. We investigate this by presenting our results for BEiT. Unfortunately, MaskGiT's official training script is not publicly available, and I-JEPA's training objective could lead to collapsed solutions. Their technique of using exponential moving average (EMA) to prevent a trivial solution is inadequate for our proposed low-rank bias. To address this, additional regularization is necessary. Our method is applied to BEiT by adhering to the original setup, incorporating skip connections between the encoder and decoder, and modifying the encoder's residual connection decay as per our approach. We run experiments on ImageNet-100 and report the highest linear probing accuracy across all layers, demonstrating enhanced performance over baseline.

Table R8: Results with BEiT. Our method improves the probing accuracy compared to the baseline with the full residual connection.

---

• W3: Discussion of encoder-decoder skip connections

We discuss the UNet setup in Lines 256-262 and show the comparison for our method w/ and w/o the UNet setup in Table 3(a). For a fair comparison, competing methods (including the baseline method, equivalent to $\alpha_{\text{min}} = 1.0$) also have encoder-decoder skip connections and Table 3(b) shows our method is substantially better than the baseline.

---

• W4: Typo

Thank you for pointing it out; we will correct it in the final version.

---

• Q1: Experiments with Stochastic Depth (also covered in the general reply)

Training with stochastic depth achieves the opposite of our desired effect. Stochastic depth randomly drops a parameterized layer, keeping the residual connection that runs parallel to it. This increases the relative contribution of the residual link to downstream layer inputs, thus amplifying the importance of that residual connection. We want to decay, not amplify, residual connection strength.

To verify that dropping layers, as in stochastic depth, is not a replacement for decaying identity shortcuts, we perform the following experiment. We keep the same training configurations for our ImageNet-100 MAE experiments and replace our proposed decayed identity shortcuts with randomly dropping network blocks as implemented in [2]. Following [2], we use the same linear decay scheduler, and layer $l$ is dropped with a chance of $\frac{l}{L}\cdot \text{Drop}_{\text{min}}$ during training. We report the linear probing of final representations as follows:

Table R7: Comparison to stochastic depth.

The more we increase the drop chance, the more stochastic depth harms generative representation learning.

Thank you for your feedback.

• W1: Scalable for deeper networks and the optimal choices of $\alpha_{\text{min}}$.

Our approach effectively scales with both increased model depth and channel size. The decay scheme we propose gradually adjusts the decay factor as depth increases, ensuring a smooth transition in even the deeper models.

To show this, we provide additional results (also in the general reply) across multiple model configurations and provide an empirical criterion for selecting the optimal $\alpha_{\text{min}}$. Specifically, we evaluate different choices of $\alpha_{\text{min}}$ across two axes: feature dimension and encoder depth, and report the results in Table R1. Linear probing accuracy is measured using features extracted from the encoder's final layer.

In Table R2, we extend the analysis by reporting the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses). Across various configurations, including ViT-Base and ViT-Large, our method consistently outperforms the baseline that employs a full residual connection ($\alpha_{\text{min}} = 1$).

Our findings reveal that the number of encoder layers $L$, rather than the embedding dimension, primarily determines the optimal $\alpha_{\text{min}}$. We attribute this behavior to the scaling effect of the input data on the encoder's final layer, quantified as: $\alpha_L^{\text{eff}} = \prod_{l=1}^L(1 - \frac{(1 - \alpha_{\text{min}})l}{L})$. Deeper models require larger $\alpha_{\text{min}}$ to maintain a consistent cumulative decay effect represented by $\alpha_L^{\text{eff}}$.

Table R3 presents the values of $\alpha_L^{\text{eff}}$ for the best-performing layer under various configurations. From these results, we see that selecting $\alpha_{\text{min}}$ to make $\alpha_L^{\text{eff}} \in [0.001, 0.01)$ consistently yields a better result than the baseline model with a full residual connection. Particularly, along with Table R2, we observe that the improvements are robust when $\alpha_L^{\text{eff}}$ falls in this preferred range. Choosing $\alpha_L^{\text{eff}}$, and using the above formula, we can compute $\alpha_{\text{min}}$ for a particular network.

Table R1: Linear probing accuracy of the final layer of the encoder. Across multiple model configurations, we report the probing score when applying our method. For a family of architectures, by properly choosing $\alpha_{\text{min}}$, our method leads to a significantly better result than the baseline with the full residual connection.

Table R2: Linear probing accuracy of the best-performing layer of the encoder. Across multiple model configurations, we report the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses). Comparing to Table R1 shows that the best probing score appears at shallower layers for deeper models.

Table R3: Best linear probing accuracy for various $\alpha_L^{\text{eff}}$. Results show that for different model configurations, setting $\alpha_{\text{min}}$ such that $\alpha_{L}^{\text{eff}} \in [1\text{e-3}, 1\text{e-2})$ consistently achieves the best performance.

Thank you for your feedback.

• W1: Fine-tuning performance (also answered in responses to Reviewers Hwrn and Kv2k)

Our primary focus is on representation learning. End-to-end fine-tuning primarily evaluates the utility of features as initialization for downstream tasks, which is influenced by various factors such as augmentation strategy and layer-wise learning rates in the fine-tuning stage, rather than directly measuring the quality of the learned representations. In contrast, linear probing offers a more direct assessment of representation quality by utilizing a simple linear projection to adapt the features for downstream applications.

Moreover, the results from end-to-end fine-tuning do not accurately reflect the quality of unsupervised representations, as the fine-tuning process significantly alters the original features. For example, ADDP [1] achieves only 23.8% accuracy in linear probing, highlighting poor representation learning. However, its accuracy improves significantly to 85.9% after end-to-end fine-tuning, illustrating the substantial transformation of representations during fine-tuning.

---

• Q1-Q2: Applying our method to DINO.

Different representation learning approaches exhibit varying preferences for feature rank. Generative representation learning methods, such as MAE and Diffusion, benefit from low-rank regularization. In contrast, contrastive and discriminative methods, including DINO, favor high-rank features due to mechanisms like the repulsive term in the denominator or excessive cluster centers like DINO. RankMe [2] supports this hypothesis, demonstrating that the performance of contrastive models improves as they learn higher-rank features.

Our approach, which induces low-rank features, is therefore not well-suited for contrastive learning and could negatively impact its performance. As highlighted in our title, our method is specifically designed for generative representation learning, which imposes fewer assumptions on data distribution compared to contrastive learning, and is not intended for the contrastive learning framework.

[1] Tian et al., ADDP: Learning General Representations for Image Recognition and Generation with Alternating Denoising Diffusion Process. ICLR, 2024.

[2] Garrido et al., RankMe: Assessing the Downstream Performance of Pretrained Self-Supervised Representations by Their Rank. ICML, 2023.

1. Effect of $\alpha_{\text{min}}$ for Wider Model Configurations and Automatic Selection of Optimal $\alpha_{\text{min}}$

To assess the compatibility of our proposed method across a broader range of model configurations, we evaluate different choices of $\alpha_{\text{min}}$ across two axes: feature dimension and encoder depth, and report the results in Table R1. Linear probing accuracy is measured using features extracted from the encoder's final layer.

In Table R2, we extend the analysis by reporting the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses). Across various configurations, including ViT-Base and ViT-Large, our method consistently outperforms the baseline that employs a full residual connection ($\alpha_{\text{min}} = 1$).

These results reveal that the number of encoder layers $L$, rather than the embedding dimension, primarily determines the optimal $\alpha_{\text{min}}$. We attribute this behavior to the scaling effect of the input data on the encoder's final layer, quantified as: $\alpha_L^{\text{eff}} = \prod_{l=1}^L(1 - \frac{(1 - \alpha_{\text{min}})l}{L})$. Deeper models require larger $\alpha_{\text{min}}$ to maintain a consistent cumulative decay effect represented by $\alpha_L^{\text{eff}}$.

Table R3 presents the values of $\alpha_L^{\text{eff}}$ for the best-performing layer under various configurations. From these results, we see that selecting $\alpha_{\text{min}}$ to make $\alpha_L^{\text{eff}} \in [0.001, 0.01)$ consistently yields a better result than the baseline model with a full residual connection. Particularly, along with Table R2, we observe that the improvements are robust when $\alpha_L^{\text{eff}}$ falls in this preferred range. Choosing $\alpha_L^{\text{eff}}$, and using the above formula, we can compute $\alpha_{\text{min}}$ for a particular network.

Table R1: Linear probing accuracy of the final layer of the encoder. Across multiple model configurations, we report the probing score when applying our method. For a family of architectures, by properly choosing $\alpha_{\text{min}}$, our method leads to a significantly better result than the baseline with the full residual connection.

Table R2: Linear probing accuracy of the best-performing layer of the encoder. Across multiple model configurations, we report the highest linear probing accuracy across all layers, along with the index of the best-performing layer (displayed in parentheses). Comparing to Table R1 shows that the best probing score appears at shallower layers for deeper models.

Table R3: Best linear probing accuracy for various $\alpha_L^{\text{eff}}$. Results show that for different model configurations, setting $\alpha_{\text{min}}$ such that $\alpha_{L}^{\text{eff}} \in [1\text{e-3}, 1\text{e-2})$ consistently achieves the best performance.

2. Cosine Decay Schedule v.s. Linear Decay Schedule

We investigate decay schedule options and perform additional experiments on ImageNet-100, employing a cosine decay schedule for MAE.

Table R4: Ablation on cosine decay schedule.

The cosine decay schedule is less sensitive to the selection of $\alpha_{\text{min}}$ than the linear decay schedule, yet it achieves lower performance at the optimal $\alpha_{\text{min}}$ than for the linear schedule.

3. Learnable $\alpha_l$

We run experiments by making the $\alpha_l$ learnable parameters rather than a fixed constant during training. We run experiments with ViT-B/16 on ImageNet-100 dataset and we use the same training configurations for our ImageNet-100 experiment. We add sigmoid activations to constrain the value of $\alpha_l$ between [0,1] and even initialize the learnable parameters with a nearly optimal constant, $\sqrt{0.5}$, based on our prior experiments. No weight decay or penalty is applied for regularizing $\alpha_l$, since we found such regularizations tend to reduce $\alpha_l$ in deeper layers, adversely affecting performance. We show the final $\alpha_l$ for each layer in the following tables:

Table R5: The learned $\alpha_l$ at convergence of training.

Table R6: Comparison of the linear probing accuracy using learnable $\alpha_l$ and our constant linear decay schedule.

From the table, we find that, in deeper layers, $\alpha_l$ values decrease and become unstable and the feature quality is even worse than our simpler linear schedule. The unstable nature of learnable $\alpha_l$ might raise a concern for deeper models. This suggests that $\alpha$ is more appropriately regarded as a regularization hyperparameter rather than a learned component of the model.

4. Comparison to stochastic depth

Training with stochastic depth, as suggested by Reviewer Hwrn, achieves the opposite of our desired effect. Stochastic depth randomly drops a parameterized layer, keeping the residual connection that runs parallel to it. This increases the relative contribution of the residual link to downstream layer inputs, thus amplifying the importance of that residual connection. We want to decay, not amplify, residual connection strength.

To verify that dropping layers, as in stochastic depth, is not a replacement for decaying identity shortcuts, we perform the following experiment. We keep the same training configurations for our ImageNet-100 MAE experiments and replace our proposed decayed identity shortcuts with randomly dropping network blocks as implemented in [1]. Following [1], we use the same linear decay scheduler, and layer $l$ is dropped with a chance of $\frac{l}{L}\cdot \text{Drop}_{\text{min}}$ during training. We report the linear probing of final representations as follows:

Table R7: Comparison to stochastic depth.

The more we increase the drop chance, the more stochastic depth harms generative representation learning.

5. Experiments with BEiT

We apply our method to BEiT in adherence to the original configuration, with the addition of skip connections between encoder and decoder, and decaying the encoder's residual connections with our formulation. We run experiments on ImageNet-100 and report the highest linear probing accuracy across all layers, demonstrating enhanced performance over baseline.

Table R8: Results with BEiT. Our method improves the probing accuracy compared to the baseline with the full residual connection.

6. Qualitative Comparison of Generated Images

We have updated our supplementary material and Appendix Figure 9 provides visualization of generated images from the conditional diffusion models, comparing our method with $\alpha_{\text{min}} = 0.6$ to the baseline $\alpha_{\text{min}} = 1.0$.

[1] Huang et al., Deep Networks with Stochastic Depth. ECCV, 2016.