Residual Connections Harm Generative Representation Learning

## Abstract

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.

## 1 Introduction

Residual networks (ResNets) (He et al., ) define a connection structure that has achieved near-universal adoption
into modern architectures for deep learning. At the time of their development, supervised learning
(*e.g.,* ImageNet (Deng et al., ) classification) was the driving force behind the evolution of convolutional
neural network (CNN) architectures. Residual networks solved a key issue: CNNs constructed of more than approximately
20 convolutional layers in sequence became difficult to train, leading to shallower networks outperforming deeper
ones, unless additional techniques, such as auxiliary outputs (Szegedy et al., ) or batch
normalization (Ioffe & Szegedy, ), were employed. Both ResNets, and their predecessor, highway
networks (Srivastava et al., ) provide elegant solutions to this trainability problem by endowing the network
architecture with alternative shortcut pathways along which to propagate gradients. Highway networks present a more
general formulation that modulates these shortcut connections with learned gating functions. However, given their
sufficient empirical effectiveness, the simplicity of ResNet’s identity shortcuts (residual connections) makes them a
preferred technique.

Figure: Figure 2: Our *decayed identity shortcuts* introduce a depth-dependent scaling factor to shortcuts in a residual network, thereby modulating the contribution of preceding layers and fostering greater abstraction in deeper layers. A simple schema for controlling decay factor $\alpha$ suffices to improve feature learning in both MAEs and diffusion models, as well as diffusion model generation quality.
Refer to caption: x2.png

While solving the gradient propagation issue, residual connections impose a specific functional form on the network;
between residual connections, each layer (or block of layers) learns to produce an update slated to be added to its
own input. This incremental functional form may influence the computational procedures learned by the
network (Greff et al., ). Alternatives to residual and highway networks exist that do not share this
functional form, but implement other kinds of skip-connection scaffolding in order to assist gradient
propagation (Larsson et al., ; Huang et al., ; Zhu et al., ). Thus, shortcut pathways, rather than a
specific form of skip connection, are the essential ingredient to enable the training of very deep networks.
Nevertheless, nearly all modern large-scale models, including those based on the transformer
architecture (Vaswani et al., ) incorporate the standard identity shortcut residual connection.

This design choice holds, even as deep learning has shifted into an era driven by self-supervised training. The
shift to self-supervision brings to the forefront new learning paradigms, including those based on contrastive (Wu et al., ; He et al., ; Chen et al., ; Caron et al., ; Grill et al., ), generative (Goodfellow et al., ; Karras et al., ; Ho et al., ; Song et al., ; ; Rombach et al., ), and autoencoding (Kingma & Welling, ; He et al., ; Li et al., ) objectives.
Many systems in the generative and autoencoding paradigms rely on “encoder-decoder” architectures, often styled
after the original U-Net (Ronneberger et al., ), which contains additional long-range shortcuts between corresponding
layers in mirrored symmetry about a central bottleneck. With representation learning as a goal, one typically desires
that the middle bottleneck layer produce a feature embedding reflecting abstract semantic properties. The interaction
of skip-connection scaffolding for gradient propagation with encoder-decoder architectures, self-supervised training
objectives, and bottleneck representations has not been carefully reconsidered. This is a worrisome oversight,
especially since, even in the supervised setting with standard classification architectures, prior work suggests that
unweighted identity shortcuts may be a suboptimal design decision (Savarese & Figueiredo, ; Fischer et al., ).

Intuitively, identity shortcut connections may not be entirely appropriate for capturing high-level, semantic
representations as they directly inject low-level, high-frequency details of inputs into outputs, potentially
compromising feature abstraction. We explore this issue within generative learning frameworks, including
masked autoencoders (MAEs) (He et al., ) and diffusion models (Ho et al., ), leading paradigms
for self-supervised image representation learning and generation. Our experiments demonstrate that identity shortcuts
significantly harm semantic feature learning in comparison to an alternative we propose: gradually decay the weight of
the identity shortcut over the depth of the network, thereby reducing information flow through it
(Figure 2). With increasing layer depth, our approach facilitates a smooth transition from a residual
to a feed-forward architecture, while maintaining sufficient connectivity to train the network effectively. Unlike
prior work on learned gating (Srivastava et al., ) or reweighting (Savarese & Figueiredo, ) mechanisms for
residual connections, our method is a forced decay scheme governed by a single hyperparameter.

A parallel motivation for our design stems from Huh et al. (), who show that features from residual
blocks have higher rank than those produced by comparative feed-forward blocks. The smooth transition between
residual and feed-forward behavior induced by our decay scheme regularizes deeper features toward exhibiting
low-rank characteristics. Section 6 experimentally explores the correlation between our decayed
identity shortcuts and low-rank feature representations. Figure 1 previews the corresponding
improvements to representation learning. Our contributions are:

- $\bullet$
We introduce decayed identity shortcuts, a simple architectural mechanism which enhances semantic feature
abstraction in masked autoencoders and diffusion models.
- $\bullet$
We identify a key correlation between our decayed identity shortcuts and low-rank inductive bias, empirically
validating that our method improves classification accuracy and yields low-rank features with distinct cluster
structures.
- $\bullet$
Our novel design within an MAE yields a substantial performance boost in linear probing on
ImageNet-1K (Deng et al., ) (72.7% from a baseline 67.8%).
- $\bullet$
In diffusion models, our design improves both feature learning and generation quality.
- $\bullet$
Ablation studies on ImageNet-100 show that smaller models equipped with decayed identity shortcuts outperform
larger ones using standard residual connections. A VIT-S/16 model (Dosovitskiy et al., ) with our shortcut
design outperforms a baseline VIT-B/16 (78.5% *vs.* 76.5%).

## 2 Related Work

Self-supervised representation learning.
Recent advancements (Achiam et al., ; Kirillov et al., ; Rombach et al., ; Team et al., ; Shi et al., ; Ramesh et al., ) in deep learning follow a common scaling law, in which a model’s performance consistently improves
with its capacity and the size of the training data. This effect can be observed in large language models (LLMs),
which are trained on vast amounts of internet text, enabling them to perform some tasks at human
level (Laskar et al., ) and exhibit remarkable zero-shot capabilities (Kojima et al., ). These models
are trained using next-token-prediction, allowing them to be trained without labeled data. In contrast, the progress
of this scaling law in computer vision has largely depended on annotated data. For instance, the Segment Anything
model (Kirillov et al., ) leverages 1 billion human-annotated masks, and state-of-the-art image
generators (Ramesh et al., ) require training on huge datasets of text-image pairs (Schuhmann et al., ).
However, the vast volume of unlabeled visual data and desire for continued scaling motivates a transition to
self-supervised learning.

At present, two families of approaches to self-supervised visual representation learning appear particularly
promising: contrastive learning (Wu et al., ; He et al., ; Chen et al., ; Caron et al., ; Grill et al., ), which trains a discriminative model to maximize mutual information across image augmentations,
and generative learning, via masked image modeling (Bao et al., ; He et al., ; Chen et al., ), which trains to
reconstruct occluded pixels, or via diffusion denoising (Song et al., ; Ho et al., ; Song et al., ), which trains to reverse a process that mixes images with Gausssian noise. Some hybrid
approaches (Zhou et al., ; Huang et al., ; Li et al., ) combine both families. Despite advancements, neither
has demonstrated the same scalability (Singh et al., ) as seen in LLMs. This challenge is additional
motivation for reconsidering the foundations of self-supervised network architectures.

Residual and skip-connection architectures.
Highway networks (Greff et al., ) first propose an additive skip connection structure to provide a scaffolding
for gradient propagation when training very deep (*e.g.,* 100 layer) networks. Motivated by the gating mechanisms within
LSTMs (Hochreiter & Schmidhuber, ), this solution uses learned gating functions to weight each combination of identity
and layer output branches. Residual networks (He et al., ) are a simplification that removes these learned
coefficients. DenseNet (Huang et al., ) and FractalNet (Larsson et al., ) demonstrate that access
to gradient paths of multiple lengths are the core requirement of training scaffolding, by introducing skip-connection
structures with other functional forms. DenseNet utilizes feature concatenation instead of addition, while FractalNet
imposes a recursive tree-like architecture combining subnetworks of multiple depths.

Zhu et al. () explore variants of ResNets and DenseNets with fewer points of combination between different
internal paths, demonstrating that a sparser scaffolding structure may be more robust as network depth increases to
thousands of layers. Savarese & Figueiredo () add a scalar gating functional to the layer output in residual
networks, yielding a hybrid design between residual and highway networks; learning this scalar gating provides a
consistent benefit to classification accuracy. Fischer et al. () develop a weighting scheme for residual
connections based upon a sensitivity analysis of signal propagation within a ResNet. To date, none of these potential
improvements have seen broad adoption.

Low rank bias in neural networks.
Over-parameterized neural networks exhibit surprising generalization capabilities, a finding seemingly in contradiction
with classical machine learning theory (Neyshabur et al., ). This phenomenon implies the existence of some form
of implicit regularization that prevents the model from overfitting. From the perspective of neural network
parameterizations, Arora et al. () suggest that linear models with more layers tend to converge to minimal
norm solutions. In the context of CNNs, Huh et al. () demonstrates that stacking more feed-forward layers compels
the model to seek solutions of a lower rank, and Jing et al. () reinforce this finding by adding more layers
to an autoencoder’s bottleneck, thereby creating a representation bottleneck. In vision transformers,
Geshkovski et al. () examine the connection between attention blocks and mean-shift clustering
(Cheng, ), showing that repeated attention operations result in low-rank outputs. Moreover,
Dong et al. () reveal that eliminating the shortcut connection from residual attention blocks causes
features to degenerate to rank 1 structures doubly exponentially. From a different perspective, recent
work (Radhakrishnan et al., ; Beaglehole et al., ; Radhakrishnan et al., ) shows training algorithms
implicitly induce low-rank behavior in neural networks. Radhakrishnan et al. () study the dimensionality
reduction behavior of a recursive feature machine (Radhakrishnan et al., ) and effectively verify performance
on low-rank matrix recovery.

## 3 Method

Prior works show that deeper feed-forward architectures have an inductive bias towards producing low-rank feature maps, while ResNets do not display the same behavior (Huh et al., ). However, despite this bias, deeper feed-forward architectures are typically less effective and generalize worse than ResNets (He et al., ). We aim to combine the properties of both feed-forward networks and ResNets, using the low-rank prior to enhance the abstraction capability of the network while maintaining the core benefits of the residual block, including stable training and the capacity to construct deeper models.

### 3.1 Decayed Identity Shortcuts

Feed-forward layers. Consider a neural network of $L$ layers.
For each layer $l$ parameterized with ${\bm{\theta}}_{l}$, the operation of a feed-forward neural network can be described as:

$$ ${\bm{x}}_{l+1}=f_{{\bm{\theta}}_{l}}({\bm{x}}_{l}),$ (1) $$

where ${\bm{x}}_{l}\in\mathbb{R}^{d}$ represents the output from the preceding layer, and $f_{{\bm{\theta}}_{l}}$ denotes the transformation applied at the current layer. Although it is widely known that pure feed-forward architectures are susceptible to vanishing gradients when building deeper models, Huh et al. () demonstrates that feed-forward modules offer implicit structural regularization, enabling deep models to generate abstract representations at bottlenecks.

Residual connections.
To address the optimization problem of vanishing gradients in deeper neural networks,
ResNets (He et al., ) construct each layer as a residual function, resulting in a modification to Eq. 1:

$$ ${\bm{x}}_{l+1}={\bm{x}}_{l}+f_{{\bm{\theta}}_{l}}({\bm{x}}_{l}).$ (2) $$

This design builds shortcuts from input to output, allowing gradient magnitude to be preserved regardless of the depth of the model. However, a consequence of this design is that the output stays close to the input in practice (Greff et al., ), defeating the need to construct complex transformations over depth. The same phenomenon is also observed in highway networks (Srivastava et al., ), which adopt learnable gates $H_{\phi}({\bm{x}})\in[0,1]^{d}$ in both the residual and skip branches:
${\bm{x}}_{l+1}=H_{\phi}({\bm{x}}_{l})\cdot{\bm{x}}_{l}+(1-H_{\phi}({\bm{x}}_{l
}))\cdot f_{{\bm{\theta}}_{l}}({\bm{x}}_{l})$. Although this flexible design allows the model to build the abstraction level over depths, similar to feedforward networks, Srivastava et al. () finds $H_{\phi}\approx 1$ for most units, suggesting the model prefers copying the input.

Decayed identity shortcuts for unsupervised representation learning.
Setting aside the optimization benefits brought by residual connections,
we rethink the role of the residual connections from the viewpoint of representation learning. Abstraction can be viewed as invariance to local changes of input and is crucial to the disentanglement of the feature space (Bengio et al., ). Prior work suggests that a shortcut path of residual connections tends to preserve high-frequency fine-grained input information (Greff et al., ), resulting in decreased feature abstraction. We hypothesize that this lack of abstraction harms the capability of the model to learn meaningful low-level features and that ensuring an abstract structure in the deeper layers of the neural network will help improve representation learning, especially for unsupervised tasks that often use indirect proxy objectives, such as pixel-wise reconstruction loss. Motivated by this hypothesis, we propose to downweight the contribution from the shortcut path:

$$ ${\bm{x}}_{l+1}=\alpha_{l}{\bm{x}}_{l}+f_{{\bm{\theta}}_{l}}({\bm{x}}_{l}),$ (3) $$

where $\alpha_{l}\in[0,1]$ is a rescaling factor to the residual path, controlling the information flow through the skip connection. Fully expanding this relation for a network with $L$ layers indexed from 0 to $L-1$, we have that:

$$ ${\bm{x}}_{L}=\left(\prod_{l=0}^{L-1}\alpha_{l}\right){\bm{x}}_{0}+\sum_{l=0}^{ L-2}\left(\prod_{i=l+1}^{L-1}\alpha_{i}\right)f_{{\bm{\theta}}_{l}}({\bm{x}}_{ l})+f_{{\bm{\theta}}_{L-1}}({\bm{x}}_{L-1}).$ (4) $$

We see that the contribution of the input ${\bm{x}}_{0}$ is scaled by each $\alpha_{l}\leq 1$ while each subsequent network block output $f_{{\bm{\theta}}_{l}}({\bm{x}}_{l})$ omits scaling factors up to $\alpha_{l}$. Hence, the contribution of early features of the network is especially down-weighted, preventing the network from passing fine-grained detailed information to the bottleneck $X_{L}$.

Decay schema. Rather than adopting a naive choice of $\alpha_{l}$ as a constant across all layers, we choose $\alpha_{l}$ to be a function parameterized by the layer index $l$, where the contribution from the shortcut path is monotonically decreasing when $l$ increases:

$$ $\alpha_{l}=1-\delta_{\alpha}l,$ (5) $$

where $\delta_{\alpha}:=\frac{(1-\alpha_{\rm min})}{L}$, $\alpha_{L-1}\equiv\alpha_{\rm min}$ is a minimum scaling factor applied at the final layer $L-1$. Our formulation brings two primary benefits. First, $\alpha_{l}$, as a linear interpolation between 0 and 1, acts as a smooth transition between residual connections and feedforward layers,
bringing us the optimization benefits seen in the residual connections, while simultaneously encouraging learning the deeper layers to learn more abstract representations.
Second, similar to the naive formulation, our method only introduces one extra hyperparameter $\alpha_{\rm min}$, which is not data-dependent and does not need to be learned.

### 3.2 Implementation Strategy

Skip connections for autoencoders.
Since our method progressively decays the residual connections over network depth, it encourages the most abstract features to be learned by later layer.
However, learning a highly abstract bottleneck is detrimental to the training objectives that aim for pixel-wise reconstruction,
as they necessitate the preservation of information across all feature levels.
To address this, we incorporate standard skip connections between the encoder and decoder, enabling the encoder to directly pass information
from shallow layers to the decoder while learning increasingly abstract representations in the deeper encoder layers.

Stabilizing training with residual zero initialization.
The model exhibits rapid feature norm growth at the beginning of training for $\alpha_{\rm min}\leq 0.7$.
We suspect that the model learns to amplify the output feature norm of $f_{{\bm{\theta}}_{l}}({\bm{x}})$ to counteract the significant decay applied to the residual connection.
This growth leads to training instability and negatively impacts training convergence. To address this issue, we follow the implementation of previous works (Ho et al., )
and initialize the weights of the final output layer in each $f_{{\bm{\theta}}_{l}}$ to zero instead of using the original Xavier uniform initialization (Glorot & Bengio, ).
This approach significantly enhances training stability by limiting the rate of feature norm growth and enables us to explore training with even lower values of $\alpha_{\rm min}$.

## 4 Experiments on Masked Autoencoder (MAE)

For masked autoencoders (MAEs) (He et al., ), we replace the residual connections in the encoder’s MLP and attention blocks with decayed identity shortcuts.
The MAE operates by accepting images with a random subset of pixels masked out and learning to recover the discarded pixels.
As Section 3.2 describes, we add skip connections between the encoder and decoder to facilitate learning abstract
representations at the bottleneck. Since the original MAE has twice the number of encoder layers as decoder layers, we inject output
from every other encoder layer into the corresponding decoder layer. To match spatial dimensions, injected encoder features are combined
with learnable masked tokens before channel-wise concatenation. The implementation details for the training and evaluation are shown in
Section A. He et al. () show the desired representations appear at the end of encoder; we therefore
apply our decaying schema only to the encoder.

### 4.1 Representation Learning on ImageNet-1k

We pre-train MAE on the ImageNet-1K train split (Deng et al., ) and follow recent works to evaluate the learned representations using linear probing and end-to-end finetuning, where we added a single linear head over learned representations to predict the image categories

For hyperparameters of pretraining MAE, including the learning rate schedule, total training epochs, and mask ratio, we follow the best settings found in the original paper.
Please see the appendix for detailed experimental setups.

**Table 1: Accuracy of linear classifier based on pre-trained learned representations on the ImageNet-1K dataset. We evaluate our learned representation using the standard evaluation protocol: linear probing (LP) and fine-tuning (FT). With our simple modification, we substantially improve the MAE ViT-B/16 baseline for linear probing by encouraging the model to learn an increased abstract features over depth. We achieve competitive performance compared to I-JEPA which uses an explicit feature loss.**
| Method | Objective | Augmentation | FT | LP |
| --- | --- | --- | --- | --- |
| MoCo-v3(Chen et al., ) | Contrastive Loss | Full | 83.2 | 76.7 |
| DINO(Caron et al., ) | Contrastive Loss | Full | 83.3 | 78.2 |
| Con MIM(Yi et al., ) | Contrastive Loss | Affine + Mask | 83.7 | 39.3 |
| ADDP(VIT-L) (Tian et al., ) | Feature Loss (VQ) | Affine + Mask | 85.9 | 23.8 |
| Latent MIM(Wei et al., ) | Feature Loss | Affine+ Mask | 83.0 | 72.0 |
| Data2Vec(Baevski et al., ) | Feature Loss | Full+Mask | 84.2 | 68.0 |
| CAE(Chen et al., ) | Recon. + Feature Loss | Affine+Mask | 83.8 | 70.4 |
| I-JEPA(Assran et al., ) | Feature Loss | Affine+Mask | - | 72.9 |
| MAE(He et al., ) | Recon. Loss | Affine+Mask | 83.6 | 67.8 |
| Ours ($\alpha_{\rm min}$ = 0.6) | Recon. Loss | Affine+Mask | 82.9 | 72.7 |

We report the results in Table 1, where
we show the linear probing performance of various self-supervised methods,
which we categorize by their objectives, and data augmentation processes.
In the top half of the table, we present methods that employ a contrastive loss.
Although these methods produce the best probing accuracies, their success depends on a carefully designed data augmentation process,
which may need to be tuned for each different data distribution.
In the bottom half, we show several methods based on generative architecture, including ours, which do not rely on contrastive objectives.
With the exception of Data2Vec, these methods only use a standard random affine data transformation (with random masking),
which need not be distribution-reliant. Among these methods, MAE only uses a pixel-wise loss, I-JEPA , Latent MIM, and CAE use a latent feature alignment loss, and CAE uses both.
Our method simply extends MAE by constructing an implicit feature bottleneck and shows significant improvements over the MAE baselines (72.7% *vs.* 67.3%),
outperforming Data2Vec, Latent MIM and CAE and giving a probing accuracy competitive with I-JEPA, without needing to use explicit feature alignment.

End-to-end fine-tuning, unlike linear probing which only trains a single linear layer, updates the entire network for image classification. Since the features can shift significantly from their pre-training state during end-to-end updating, we argue that this may not accurately reflect the quality of the learned representations. For example, DINO demonstrates superior performance in various downstream vision tasks compared to MAE, but its fine-tuning performance is worse than MAE. Similarly, ConMIM and ADDP exhibit poor linear probing performance, suggesting lower-quality representations, yet their fine-tuning performance surpasses that of contrastive learning methods. Nevertheless, we still provide the fine-tuning results for reference.

Figure: (a) MAE
Refer to caption: x3.png

### 4.2 Embdding analysis.

We qualitatively evaluate the feature learning in Figure 3(a) by visualizing the t-SNE embeddings of the learned features for a random subset of test images. The embedding produced from our features displays much clearer separation than baseline, which has struggles to differentiate the categories.

To provide another qualitatively evlauation of our learned representations, we adopt the pixel-wise embedding approaches proposed by Zhang et al. () to group the
representations from the last layer of the encoder into a lower dimensional space. We use their default hyperparameters to cluster representations across images of COCO validation set. We visualize in Figure 4 the top 3 eigenvectors
for both our approaches with $\alpha_{\min}=0.7$ and baseline. From the visualization, ours learns abstract representation and the object from the same categories have similar color, indicating a global consistent semantic grouping. The baseline MAE, on the other hand, doesn’t show clearly global semantic patterns overall and the low-level representations, e.g., edges, are also high-lightened, indicating a representations from multiple semantic hierarchies are entangled.

We further evaluate the clustering quantitatively, following the postprocessing protocol  (Zhang et al., ) to produced unsupervised semantic segmentation and report the results as the mean intersection of union (mIoU). Ours (10.41 mIoU) achieve 6.31 mIoU improvement over baseline (4.10 mIoU), which support the qualitatively comparison.

### 4.3 Ablation Studies on ImageNet-100

We conduct ablations on several properties of our framework on ImageNet-100. A summary of results can be found in Tables 2 and 3.

**Table 2: Linear Probing Accuracy on ImageNet-100 for our method varying $\alpha_{\rm min}$ and architecture size. We conduct ablation studies and demonstrate that linear probing performance for both architectures increases as $\alpha_{\rm min}$ decreases until around 0.5-0.6. While the larger ViT-B/16 architecture achieves the highest accuracy of 83.6, it is noteworthy that the smaller ViT-S/16, when utilizing our method, outperforms the baseline setting (standard Residual Connection at $\alpha_{\rm min}=1$) of ViT-B/16.**
| Backbone<br>$\alpha_{\rm min}$ | 0.5 | 0.6 | 0.7 | 0.8 | 0.9 | 1.0 |
| --- | --- | --- | --- | --- | --- | --- |
| ViT-B/16 | 82.3 | 83.6 | 81.8 | 79.8 | 79.2 | 76.5 |
| ViT-S/16 | 78.6 | 78.5 | 78.1 | 75.2 | 73.5 | 69.2 |

Decay rate $\alpha_{\rm min}$.
The only parameter of our framework is $\alpha_{\rm min}$, the minimum scaling factor applied to the identity shortcut at the final layer. In Table 2,we show linear probing scores for varying values of $\alpha_{\rm min}$. We observe that $\alpha_{\rm min}$ must be sufficiently small to regularize the flow of information through the residual connection effectively. A reasonably small $\alpha_{\rm min}$ prevents the deeper layers of the encoder from relying heavily on the residual connections, allowing for more abstract representations in the bottleneck.
This yields up to a 7.1% improvement over the ViT-B/16 baseline (${\alpha_{\rm min}}=1$) in linear probing accuracy on ImageNet-100.
On the other hand, if $\alpha_{\rm min}$ is too small, for example, $\alpha_{\rm min}\leq 0.4$ for ViT-B/16, we observe that the training becomes unstable.

Architecture size. In Table 2, we also train both the ViT-B/16 and the smaller ViT-S/16 backbone using varying $\alpha_{\rm min}$. Our framework is especially effective on the smaller architecture, increasing linear probing performance by 9.4%, compared to the baseline setting (standard Residual Connection at $\alpha_{\rm min}=1$). This is consistent with the observation that larger models generalize better (Huh et al., ) than smaller models. Our method is able to significantly improve the smaller ViT-S/16 and slightly close the gap between the differently-sized architectures.

Skip connections. Another critical design choice in our network is to include skip connections that are not in the original MAE. As discussed in Section 3.2, if the MAE does not use skip connections, the bottleneck layer must preserve all information to reconstruct the input image accurately. This is opposed to learn abstract representations at bottleneck. These contrary effects significantly degrade the representation learned by the model, leading to a 22.1% drop in the linear probing score, as we report in Table LABEL:tab:skip_connections.

Different decay schema.
We also explore decay schema, with results summarized in Table LABEL:tab:decay_schema: (1) Scaling both branches of the residual blocks simultaneously by applying a constant factor, $\alpha=\sqrt{0.5}$, to both ${\bm{x}}$ and $f_{{\bm{\theta}}_{l}}({\bm{x}})$. (2) Scaling only $f_{{\bm{\theta}}_{l}}$ using the same constant factor, $\alpha=\sqrt{0.5}$. (3) Applying our proposed schema exclusively to either the attention or MLP branch.

Among these, (2) shows no significant improvement over the baseline, while (1) yields some improvement but still underperforms compared to our approach. By analyzing (1) and (2), we demonstrate that the representation gains are due to down-weighting the skip connection branch. Notably, recent diffusion models (Karras et al., ; Song et al., ; Karras et al., ; ) have employed (1) in their designs. However, applying decay only to the MLP or attention branch reduces the overall decaying effect across the network, resulting in lower performance compared to our schema, which achieves the best performance among the tested designs.

Table: (a)

## 5 Experiments on Diffusion Models

Diffusion models.
We use U-ViT (Bao et al., ), a ViT-based diffusion model with skip connections between the encoder and decoder, as the baseline for our diffusion model experiments. Recent studies (Yang & Wang, ; Baranchuk et al., ) suggest that diffusion models learn the best semantic representations near the decoder’s latter stages. Therefore, we apply our proposed decay mechanism up to the end of the decoder. While this design might be suboptimal, as the smallest decay factor may not align with the layers holding the best semantic representations, we demonstrate in practice that this simple approach effectively enhances both the learned representations and the quality of generated outputs.

Experimental details.
We utilize the default scheduler and sampler from U-ViT (Bao et al., ), replacing only the residual connections with our proposed decayed shortcut connections. We train unconditional diffusion models on CIFAR-100 and ImageNet-100 without using image class labels. Additionally, we train a class-conditional diffusion model on ImageNet-100 to validate our design across different tasks. For ImageNet-100, instead of training directly on pixels, we adopt a latent diffusion (Rombach et al., ) approach by running the model in the latent space of a pretrained VAE, which reduces input resolutions from 256x256x3 to 32x32x4. We use U-ViT-Mid for ImageNet-100 and U-ViT-small for CIFAR-100. For model and training details, please refer to  Bao et al. ().

We evaluate the learned representations with linear probing and we train a linear classifier over the frozen representations. We report the results as the best configurations, including the choices of layer index and noise level, that yields the best performance

Results. Our results are presented in Table 4, where we demonstrate that replacing residual connections with our proposed decayed identity shortcuts consistently enhances representation quality and image generation across both datasets and tasks (conditional and unconditional generation). Notably, this improvement is achieved without introducing any additional learnable parameters. We provide the visualization of generated images in the appendix for qualitative comparisons (Figure 10).

**Table 4: Performance using diffusion models. In diffusion models, we demonstrate that our proposed decayed identity shortcut (with $\alpha_{\rm min}<1.0$) enhances probing accuracy and improves generation quality across various datasets and configurations in both classes conditional and unconditional setups.**
|  | Linear Probing (Acc.)$\uparrow$ | Generation quality (FID)$\downarrow$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset<br>$\alpha_{\rm min}$ | 1.0 | 0.8 | 0.7 | 0.6 | 1.0 | 0.8 | 0.7 | 0.6 |
| CIFAR-100 (Uncon.) | 62.47 | 63.58 | 66.86 | 64.63 | 14.34 | 11.65 | 8.99 | 11.71 |
| ImageNet-100 (Uncond.) | 72.8 | 74.5 | 76.1 | 75.8 | 44.40 | 40.96 | 41.17 | 43.51 |
| ImageNet-100 (Class Cond.) | - | - | - | - | 6.93 | 5.75 | 5.11 | 4.98 |

## 6 Discussion on feature rank

In this section, we try to answer a key question:
How and why do residual connections impact the abstraction level of the deeper layers in a neural network? We delve deeper into how our design reinforces the low-rank bias of neural networks and try to connect our method to ideas in existing works (Huh et al., ). To this end, we visualize the training dynamics of our method and analysis the feature rank of our approach to provide a holistic analysis.

Low-rank simplicity bias. Huh et al. () investigate the low-rank simplicity bias in deeper feed-forward neural networks, which drives neural networks to find low-rank solutions.
At the same time, they make an empirical observation that deeper residual networks do not show a similar rank contracting behavior.

Effective rank.
For analysis purpose, Huh et al. () quantify the rank of the learned representation using the effective rank, which is a continuous measure.
For a matrix ${\bm{A}}\in\mathbb{R}^{m\times n}$, the effective rank $\rho({\bm{A}})$ is defined as the Shannon entropy of the normalized singular values (Roy & Vetterli, ):

$$ $\rho({\bm{A}})=-\sum_{i}^{\min(n,m)}\bar{\sigma}_{i}\log\bar{\sigma}_{i},$ (6) $$

where $\bar{\sigma}_{i}=\sigma_{i}/\sum_{j}\sigma_{j}$ denotes the $i^{\text{th}}$ normalized singular value. Intuitively, this measure is small when a few singular values dominate and large when singular values are evenly spread, hence giving a good continuous approximation for matrix rank.
In the following subsections, to compute the effective rank,
we use the singular values from the covariance matrix ${\bm{A}}_{\bm{\theta}}$ of the last-layer features, where ${\bm{A}}_{\bm{\theta}}(i,j)$ denotes the covariance of the learned class tokens for the $i^{\text{th}}$ and $j^{\text{th}}$ samples.

Figure: (a) Effective Rank of MAE over time.
Refer to caption: extracted/5901827/images/legend.png

Inspired by Huh et al. (), we conjecture that the improvement to feature learning capability of our method can mainly be attributed to the decayed identity shortcuts promoting low-rank features at the bottleneck.
We measure the training dynamics of the models presented in Table 2 (MAEs trained on ImageNet-100) in terms of accuracy and the effective rank. In Figure 6(c), we quantify the correlation between effective feature rank and probing accuracy using Pearson correlation coefficient to validate the hypothesis.

In Figures 6(a) and 6(b), we present the training dynamics of our model, highlighting effective rank and linear probing accuracy for different values of $\alpha_{\text{min}}$. During the early epochs, models with lower $\alpha_{\text{min}}$ tend to exhibit both lower effective rank and higher probing accuracy, supporting our hypothesis. As training progresses, the correlations between $\alpha_{\text{min}}$ and effective rank become less precise. We suspect this is due to the training dynamics, where models with lower $\alpha_{\text{min}}$ experience faster growth in feature rank, reflecting the model’s effort to compensate for the decay factors. Despite this, we can still conclude that lower $\alpha_{\min}=[0.5-0.6]$ results in lower feature rank and better probing accuracy compared to higher $\alpha_{\min}=[0.7-1.0]$.

## 7 Conclusion

Huh et al. () raise a key insight in their work – that how a neural network is parameterized matters for fitting the data – and investigate the inductive low-rank bias of stacking more linear layers in a network.
In this work, we observe that the ubiquitous residual network (He et al., ) may not be the ideal network parametrization for representation learning and propose a modification of the shortcut path in residual blocks that significantly improves unsupervised representation learning. We explore the connection between our reparameterization of the residual connection and the effective rank of the learned features, finding a correlation between good representations and low-rank representations.

Our work calls into question a fundamental design choice of neural networks that has been used in many modern architectures. By rethinking this choice, the door is open for further reparametrizations and improvements to unsupervised representation learning.
The results we show provide a prompt for more extensive investigations into the connection between low effective rank and high-quality abstract representations, as well as the exploration of underlying theoretical mechanisms for this relationship.

## Appendix A Training and Evaluation Details

### A.1 Model training.

Our training configurations primarily followed the guidelines established by He et al. ().
In the ImageNet-1K experiment, our model was trained for 800 epochs, utilizing the AdamW Loshchilov & Hutter () optimizer with a constant weight decay of 5e-2 for a batch size of 1024.
We set the maximum learning rate to 6e-4. Initially, the learning rate started at 0 and linearly increased to its maximum over the first 40 epochs, after which it followed a cosine schedule to gradually decrease to zero by the end of the training period.
It is worth noting that the learning rate per sample, or effective learning rate, in our setup matched that of He et al. (), although our maximum learning rate was set lower due to our batch size being a quarter of theirs.
We applied random resizing, cropping, and horizontal flipping during training as part of our augmentation scheme.
To enhance the quality of the learned representations in most experiments, we employed the normalized pixel loss, as proposed by He et al. ().
In the ImageNet-100 experiment, we employed the identical training configuration used in the ImageNet-1K experiments. We train our model with 4 NVIDIA A40 GPUs and a completed trianing usually takes 20 hours on ImageNet-100 and 200 hours on ImageNet-1k.

### A.2 Evaluation with Linear Probing.

For the ImageNet-1k dataset, we use the exact same evaluation protocols employed in He et al. (), which includes random data augmentation.

For the ImageNet-100 dataset, we employed a simpler evaluation protocol: We train the linear classifier with a batch size of 1024 for 200 epochs, where the learning rate starts at 1e-2 and then decays towards 0 using a cosine scheduler. During this evaluation, we do not apply any data augmentation.

Figure: Figure 7: We present our enhanced UNet Transformer architecture for Masked Auto-encoder. (1) Left: Our customized encoder blocks, equipped with our proposed decay identity shortcuts. (2) Middle: Standard transformer blocks as the decoder blocks. (3) Right: We incorporate the decay identity shortcuts exclusively within the encoder blocks of our UNet transformer and employ standard transformer blocks for the decoder. To support abstract representation learning at the bottleneck, *i.e.,* the last layer of the Encoder 12, we adopt the UNet Ronneberger et al. () architecture and create skip connections that transmit every other encoder feature directly to the decoder.
Refer to caption: x6.png

### A.3 Modified Architecture

We present a visualization of our UNet transformer design, as outlined in Section  3.2, in Fig. 7. It’s important to note that decayed identity shortcuts are exclusively implemented within the encoder block. Additionally, we establish skip connections from alternating blocks in the encoder to the decoder, following the UNet Ronneberger et al. () architecture’s design principles.

## Appendix B Further experiments

### B.1 Further Ablation of Maximum Decay Rate on ImageNet-1K

**Table 5: Linear probing accuracy of our method by varying $\alpha_{\rm min}$, the architecture size, and the dataset. We extend the ablation studies detailed in Table. 2 by including linear probing results for the ImageNet-1K dataset. The results indicate that the most favorable $\alpha_{\rm min}$ consistently falls within the range of $[0.5,0.6]$ for ImageNet-1K and ImageNet-100 experiments.**
| Dataset | Backbone<br>$\alpha_{\rm min}$ | 0.5 | 0.6 | 0.7 | 0.8 | 0.9 | 1.0 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| ImageNet-1K | ViT-B/16 | 69.8 | 72.7 | 68.9 | - | - | - |
| ImageNet-100 | ViT-B/16 | 82.3 | 83.6 | 81.8 | 79.8 | 79.2 | 76.5 |
| ImageNet-100 | ViT-S/16 | 78.6 | 78.5 | 78.1 | 75.2 | 73.5 | 69.2 |

We present the results of ablating the choices of $\alpha_{\rm min}$ on ImageNet-1K dataset in Table.5.
From the table, we show that the optimal $\alpha_{\rm min}$ for ViT-B/16 on ImageNet-100 matches the optimal one for ImageNet-1K, while a lower $\alpha_{\rm min}$ is preferred for a smaller architecture ViT-S/16.

### B.2 Reconstruction quality.

Figure: Figure 8: Qualitative comparison of images reconstructed by MAE with and without our method. We observe our method learns features with higher linear probing accuracy without compromising reconstruction quality. Row 1: ground truth test image. Row 2: images masked at $75$%. Row 3: reconstructions with our method. Row 4: reconstructions with baseline MAE.
Refer to caption: extracted/5901827/images/mae_recon_comparison.png

We qualitatively evaluate test images reconstructed by an MAE using our framework and images reconstructed by the original MAE. We show the reconstructed images in Figure 8. While the focus of our work is entirely to improve the representations learned by an encoder, we observe that our framework does not harm the reconstructions. Hence, there is no qualitative tradeoff for our increase in linear probing accuracy.

### B.3 Abstraction and Low-rank in the Supervised Setting

Figure: (a) Effective rank of ResNet for different depths at the convergence of the training.
Refer to caption: extracted/5901827/images/legend.png

In this experiment, we modify the standard ResNet-18 model to experiment with different depth models.
By default, the ResNet-18 has a total of 8 residual blocks that are equally distributed into 4 layers.
To increase model depth, we repeat residual blocks in the 3rd layer to obtain models varying between 8 and 16 total layers. At convergence, we observe that the models of different depths achieve a similar test accuracy.
However, despite similar accuracies, in Figure 9(a), which visualizes the effective rank over depth for different values of $\alpha_{\rm min}$, we see that the effective rank decreases over depth.
Furthermore, smaller values of $\alpha_{min}$ consistently lead to features with lower effective rank.

Next, in Figure 9(b), we try to verify our conjecture by visualizing the evolution of effective rank during training when choosing different $\alpha_{\rm min}$ in our method.
For this experiment, we choose to train the standard ResNet-18 using our decayed identity shortcuts. In this setup, we observe that the optimal choice of $\alpha_{\rm min}$ slightly improves the test accuracy of the classification network: 94.4% with $\alpha_{\rm{min}}=0.7$ *vs.* 93.6% with $\alpha_{\rm{min}}=1.0$.
We observe that the effective rank of the final features decreases with decreasing $\alpha_{\rm min}$. This supports our
hypothesis that (1) decayed identity shortcuts
substantially decrease the rank of bottleneck features and (2) decreasing feature rank may help improve learned features.

Figure: Figure 10: Qualitative comparison of images generated by diffusion models. Our method, decayed identity shortcuts with $\alpha_{\rm min}=0.6$, shows improved representation learning and produces higher-quality generated images compared to the baseline, which employs full residual connections ($\alpha_{\rm min}=1.0$).
Refer to caption: x9.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
strengths: + Simple modification with effective performance gains.
+ Clear writing of the method & motivation.
+ Diverse experiments for MAE, diffusion models, and effective rank analysis.
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.
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).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
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?
soundness: 3
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.
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.
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”

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
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
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
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.
soundness: 2
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.
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.
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).

Meta-review: 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.

Decision: Reject