Residual Connections Harm Generative Representation Learning

Authors: Xiao Zhang, Ruoxi Jiang, Will Gao, Rebecca Willett, Michael Maire
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Decayed Residual Connections, Representation Learning

Abstract:
We show that introducing a weighting factor to reduce the influence of identity shortcuts in residual networks significantly enhances semantic feature learning in generative representation learning frameworks, such as masked autoencoders (MAEs) and diffusion models.  Our modification improves linear probing accuracy for both, notably increasing ImageNet accuracy from 67.8\% to 72.7\% for MAEs with a VIT-B/16 backbone, while also boosting generation quality for diffusion models.  This significant gap suggests that, while residual connection structure serves an essential role in facilitating gradient propagation, it may have a harmful side effect of reducing capacity for abstract learning by virtue of injecting an echo of shallower representations into deeper layers.  We ameliorate this downside via a fixed formula for monotonically decreasing the contribution of identity connections as layer depth increases.  Our design promotes the gradual development of feature abstractions, without impacting network trainability.  Analyzing the representations learned by our modified residual networks, we find correlation between low effective feature rank and downstream task performance.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730375130551

thread_type:
official_review

<thread_object>
creator: reviewer_sarc

<content>
summary:
This paper proposes adjusting the skip connection for generative models by introducing a weight decay. In particular, skip connections at later layers of the network are multiplied by a $\alpha \leq 1$ that decreases monotonically as a function of the layer index. By doing so, the experiments in this paper shows that the learnt representation in generative models improve.

strengths:
The main strengths of this work are:

1- Simplicity of the proposed method: a simple adjustment to the network architecture.

2- Experiments are carried out on different generative models and different settings.

3- The paper is generally well-written and easy to follow.

weaknesses:
Despite the strengths of this work, there are a few concerns that need to be resolved before accepting this work. The order presents the the importance of the concern based on which I rated this paper.

1- Despite the simplicity of the approach, there is no clear methodology on how to choose $\alpha_{\text{min}}$. For example,  in section 4.2, experiments are carried out with $\alpha_{\text{min}}=0.7$, while Table 2 shows that $\alpha_{\text{min}}=0.6$ works the best with ViT-B. This is a key concern as one can choose a suboptimal value of $\alpha_{\text{min}}$ that results in Even a worse performance than just setting $\alpha_{\text{min}}=1$ (i.e. naive residual connection). A discussion on how to properly set $\alpha_{\text{min}}$ should be a core part of the methodology section. Preferably, I suggest the authors to provide systematic analysis of how $\alpha_{\text{min}}$ affects performance across different model sizes and tasks.

2- While the experiments are carried out on both MAE and Diffusion models, the results in this work need to show experiment on  training a generative model with and without the proposed method. For example, a good experiment to show this is to revisit Table and show for each method the impact of the proposed method.

3- The third or of Table 3 (b) raise a critical question: Since weighing both the terms with a simple $\sqrt{0.5}$ results in a significant performance gain, does this mean that the baseline is trained sub-optimally? 

4- In section 4.2, the experiments shows a better separation in the learnt feature space when employing the proposed method. It would also make the argument stronger if we can quantify the improvement in that regard (e.g. reporting silhouette score).

questions:
After going through the rebuttal, I commend the authors on the efforts put in addressing my comments. Thus, I raise my score from 3 to 6.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701689126

thread_type:
official_review

<thread_object>
creator: reviewer_kv2k

<content>
summary:
The authors proposed decayed identity shortcuts in residual networks, promoting the gradual development of feature abstractions. The author hypothesized that injecting shallower representations into deep layers with residual connections can be harmful by reducing the capacity for abstract learning. Then the authors conducted experiments to show this improves MAE by 4% and also benefits generation with diffusion models.

strengths:
S1. The method is straightforward and easy to follow.

S2. The intuition is reasonable, aligns with some hierarchical representation learning and generation work.

S3. Dive into both representation learning and generative modelling.

weaknesses:
W1. The experiments on representation learning are only on MAEs. In the literature, MAE is not best performing with linear probing and showed to work better in a finetuned manner, more experiments on different models would be greatly beneficial for the paper. No evidence this works universally across models. And the finetuned MAE with decayed identity shortcuts is no better than finetune MAE.
 
W2. Probing the inner state of different layers could enhance explainability and align with the intuition that underpins the work, providing deeper insights into the model’s internal processing, e.g., using representations from previous layers rather than the final layer, etc.

W3. Introducing an extra hyperparameter that needs to be justified/selected when using.

questions:
Q1. What about a cosine decay schedule? Compared to linear.

Q2. How is this fixed hyperparameter alpha compared to a learnable parameter with regularization and constraints?

Q3. How to compare the idea of the hierarchical representation learning with Matryoshka representation learning? Instead of doing it hierarchically, matryoshka rep learning embeds different virtual concepts into different chunks of embeddings.


[C1] Matryoshka representation learning., Kusupati, Aditya, et al., NeurIPS 2022
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703345315

thread_type:
official_review

<thread_object>
creator: reviewer_hwrn

<content>
summary:
The paper shows that linearly decaying (over layers) the residual term in ViT-like models improves the training of generative models. With a few details, e.g., adding U-Net like skip connections between encoder and decoder, and initializing the final encoder output layer to be zero, the proposed method could successfully stabilize the training despite of the gradual removal of residual connections. Experiments are mainly done in two scenarios: training (a) Masked Autoencoders on ImageNet-1K, and (b) Diffusion models on CIFAR-100 and ImageNet-100. The paper reports that the proposed method improves the linear probing accuracy of both MAE and Diffusion training, as well as FID for Diffusion models at certain choice of decaying hyperparameter.

strengths:
- The paper is well-motivated and easy-to-follow.
- The proposed method is notably simple and easy-to-adapt, which can be a strength in large-scale training regimes where recent generative models are indeed at.
- Rethinking the design space of neural architecture, e.g., ViTs, for generative modeling is a timely research direction.
- Figure 4 effectively highlight the quality of the learned representations with quantitative results.

weaknesses:
- It is still questionable to me whether the proposed method is scalable for deeper networks; in my experience weakening the residual terms does introduce instability of overall training, whereas the paper did not much discuss about when the method fails. For example, Table 2 shows that the optimal hyperparameter of the method varies depending on the depth of the backbone networks, i.e., ViT-S vs. ViT-B. Also, Table 4 seems to show that the method is sensitive to hyperparameter choice depending on datasets. The paper may extend such analysis for wider depths, data, etc., possibly suggesting any guideline in using the method in practice.
- Table 1: The proposed method does not improve MAE on fine-tuning setup, while the performance of MAE was originally focused on after the fine-tuning. Although I agree with the paper’s comment that the fine-tuning performance is harder to compare with, I feel the current results in Table 1 does not fully confirm the main claim, i.e., residual connection harms generative representation learning. The paper may strengthen this point by showing that the method is effective for other generative modeling approaches, e.g., I-JEPA, BEiT, MaskGiT, etc., if applicable.
- The effectiveness of method seems to be dependent to an architectural modification of ViT, i.e., only after adding encoder-decoder skip connections as like U-Net (Table 3a). But such a skip connection may already improve MAE, although I don’t think this point is clearly discussed in the current manuscript.
- (minor) Section 4.2: “Embedding analysis. → Embedding analysis”

questions:
- One alternative of decaying the residual term can be stochastic depth; i.e., by randomly dropping the identity term during training. I am curious that the authors have tested its performance, given that stochastic depth is also well-known to improve generalization of residual network training.
- Table 3b: The comparison shows that an ablation $\sqrt{0.5}(x_l + f(x_l))$ could achieve a quite competitive result to the proposed method. If so, given that the form has been already used in previous works and that it is simpler (e.g., the form does not have hyperparameter), what would be the key benefits of the proposed method over that ablated form?
- Can this method also improve models other than ViTs?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710718531

thread_type:
official_review

<thread_object>
creator: reviewer_uy11

<content>
summary:
The paper proposes two architectural changes to ViT that can greatly improve MAE's linear probing performance: (1) decreasing shortcut branch weight in deeper layers and (2) skip connection from encoder layers to decoder layers. With both changes, the authors observe better representations, evaluated with qualitative visualizations and linear probing accuracy. The authors showed that similar modifications can also improve diffusion model performance, and used effective rank as a lens to experimentally analyze & intuitively explain the performance gain.

----

I have read the authors' response and decide to keep my original evaluation.

strengths:
+ Simple modification with effective performance gains.
+ Clear writing of the method & motivation.
+ Diverse experiments for MAE, diffusion models, and effective rank analysis.

weaknesses:
+ The usual way of using MAE is not linear probing, but fine-tuning (or maybe nonlinear probing). The proposed method, unfortunately, reduces fine-tuning performance (Table 1).

questions:
1. The proposed method reduces fine-tuning performance of MAE (which is its strongest suit) and doesn't improve linear probing to be on-par with SOTA (e.g., DINOv2). Given this, have the authors explored applying the changes to DINO? If it yields benefits, it could significantly improve the paper's value.

2. If the residual connection fundamentally limits effective rank, and thus hurts representation quality, why do objectives like DINO not have such issues?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732837086648

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Reply (1/3)

comment:
We thank the reviewers for their feedback. We appreciate the recognition of the **conciseness** (uy11, Hwrn, Kv2k, saRC) and **effectiveness** (uy11, HWrn, Kv2k) of our proposed method, along with the acknowledgment of our **clear presentation and motivation** (uy11, Hwrn, Kv2k) and **diverse experiments and analysis** (uy11, saRC). In response to reviewer questions, we provide both new experimental results and clarifications.

**New Results:** To demonstrate the scalability and robustness of our method, we present experiments over a broader range of model configurations. We show that the optimal value of our hyper-parameter $\alpha_{\text{min}}$ changes in a consistent manner with network depth: the effective decay rate (cumulative decay over the entire network) is key; once this is found for a model, we can use an empirical formula to select $\alpha_{\text{min}}$ for other model architectures in the same class.

Additionally, we provide ablation studies on different decay schedules, the use of learnable $\alpha_l$, and a comparison to stochastic depth. Our findings indicate that our method, combined with a simpler linear schedule, performs the best. We validate our method on BEiT and show improvement over the baseline. We also include visualizations of our generated images in the updated supplementary material (Appendix Figure 9), making it apparent that our method also improves generation quality.

**Clarifications:** In responses to individual reviewers, we address multiple misconceptions, including overlooked primary experiments on diffusion models and convolutional neural networks, discussions regarding encoder-decoder skip connections, and the baseline models used for comparison.

The results of our new experiments are as follows:
</content>
</thread_object>

</thread>
<thread>
creation date:
1733297346167

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Post Rebuttal Reply

comment:
Dear Area Chairs and Reviewers,

We want to thank you again for your feedback!

**1. Primary Results** 

As we approach the final day of the rebuttal phase, we would like to provide a summary of our primary results
(including the new ones added during rebuttal), to facilitate a comprehensive evaluation of our work.

|Method|Model|Dataset|Metric|Baseline|Ours|
|:---:|:---:|:---:|:---:|:---:|:---:|
|MAE|ViT-B|ImageNet-1K|KNN(K=20)|27.4|**63.9**|
|MAE|ViT-B|ImageNet-1K|Linear Probing|67.8|**72.7**|
|MAE|ViT-L|ImageNet-100|Linear Probing|81.4|**86.0**|
|MAE|ViT-B|ImageNet-100|Linear Probing|76.5|**83.6**|
|MAE|ViT-S|ImageNet-100|Linear Probing|69.2|**78.5**|
|Diffusion|ViT-B|ImageNet-100|Linear Probing|72.8|**76.1**|
|Diffusion|ViT-B|ImageNet-100|FID|6.93|**4.98**|
|Diffusion|CNN|CIFAR-100|Linear Probing|62.5|**66.9**|
|Diffusion|CNN|CIFAR-100|FID|14.34|**8.99**|
|BEiT|ViT-B|ImageNet-100|Linear Probing|71.8|**74.2**|

**Table RS: Comparison to the baseline with the full residual connection.** Across various models and setups,
our method consistently outperforms the baseline, demonstrating significant improvements in (1) linear probing accuracy, (2) KNN accuracy, and (3) FID score.

**2. Scalability and Selection of Optimal $\alpha_{\text min}$**

We have shown the scalability and robustness of our method through extensive experiments, over a broader range of model configurations (as outlined in the General reply).  The optimal value of our hyper-parameter $\alpha_{\text min}$ changes in a consistent manner with network depth, where the effective decay rate (cumulative decay over the entire network) is key and can be generalized to other model architectures in the same class.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes two architectural changes to ViT that can greatly improve MAE's linear probing performance. The motivation and intuition behind the modifications (e.g., how residual connections might harm abstract feature learning) are clear and align with hierarchical representation principles. The method is easy to implement (a decaying shortcut weight) and potentially adaptable for large-scale training. Weaknesses of the paper include its limited scope, as experiments focus narrowly on MAE and diffusion models, without demonstrating generalization to other self-supervised techniques or broader architectures. The method underperforms in fine-tuning, which is crucial for MAE’s strengths, and fails to achieve competitive linear probing results compared to SOTA methods like DINOv2. Additionally, the approach introduces sensitivity to hyperparameter choices with no systematic methodology for selection, making practical adoption challenging. The results also raise questions about robustness and scalability across different datasets and model depths.

The work remained borderline after rebuttal phase, but given the original claims and conclusions (e.g. covering generative representation learning) the AC feels that the main claims should have been fully substantiated for acceptance.