DeepCrossAttention: Supercharging Transformer Residual Connections
Authors: Mike Heddes
Adel Javanmard
Kyriakos Axiotis
Gang Fu
Mohammadhossein Bateni
Vahab Mirrokni
Venue: ICML 2025 poster (2025)
Keywords: residual network
cross attention
resnet
transformer
Primary area: deep_learning->attention_mechanisms
TL;DR: We speed up transformer training by introducing learnable, input-dependent residual connections combined with depth-wise cross attention.

Abstract: Transformer networks have achieved remarkable success across diverse domains, leveraging a variety of architectural innovations, including residual connections. However, traditional residual connections, which simply sum the outputs of previous layers, can dilute crucial information. This work introduces DeepCrossAttention (DCA), an approach that enhances residual learning in transformers. DCA employs learnable, input-dependent weights to dynamically combine layer outputs, enabling the model to selectively focus on the most relevant information in any of the previous layers. Furthermore, DCA incorporates depth-wise cross-attention, allowing for richer interactions between layers at different depths. Our language modeling experiments show that DCA achieves improved perplexity for a given training time. Moreover, DCA obtains the same model quality up to 3x faster while adding a negligible number of parameters (e.g., 0.2\%). Theoretical analysis confirms that DCA provides an improved trade-off between accuracy and model size when the ratio of collective layer ranks to the ambient dimension falls below a critical threshold.

DeepCrossAttention Supercharging Transformer Residual Connections
Mike Heddes Adel Javanmard
Kyriakos Axiotis 3 Gang Fu² MohammadHoseein Bateni Vahab Mirrokni
Abstract
processing (Vaswani, 2017), audio recognition (Gong et al., 2021), and computer vision (Dosovitskiy et al., 2021).
Transformer networks have achieved remarkable success across diverse domains, leveraging vari- ety of architectural innovations, including resid- ual connections. However, traditional residual connections, which simply sum the outputs of previous layers, can dilute crucial information. This work introduces DeepCrossAttenitn Attention (DCA), an approach that enhances residual learning in transformers. DCA employs learnable, input- dependent weights to dynamically combine layer outputs, enabling the model to selectively focus on the most relevant information in any of the previous layers. Furthermore, DCA incorporates depth-wise wise coossaattenton,, allowing for richer in- teractions between layers at different depths. Our language modeling experiments show that DCA achieves improved perplexity for given training time. Moreover, DCA obtains the same model quality up to 3x faster while adding a negligible number of parameters. Theoretical analysis con- firms that DCA provides an improved trade-off between accuracy and model size when the ratio of collective layer ranks to the ambient dimension falls below critical threshold.
A residual neural network (ResNet) is constructed by stack- ing layers known as residual blocks. Each residual block is characterized by the recursive equation #++11 f(x)) + 2t, which contains a residual function f along with an identity shortcut (also called an identity loop or skip connection). The residual functions typically used in these blocks include multi-layer perceptrons (MLPs), convolutional neural net- works (CNNs), and attention. By unrolling the recursion, equivalently see that each layer's input is the sum of all its previous layers' outputs (including the model's input). Figure 2 provides a schematic illustration of this concept.
Information dilution in residual networks. Residual con- nections increase the flow of information across the neural network. However, they also come with a potential lim- itation: Taking a straight sum of previous layer outputs implicitly treats all previous layers as equally important. This can dilute useful information present in select few layers (including the model's input) with potentially less useful information. We hypothesize that, because of this dilution, even though residual networks mitigate the prob- lem of neural network bottlenecks, they do not sufficiently resolve it. One way to resolve the issue of dilution would be to allow each layer to choose its inputs.
In order to confirm the existence and significance of the di- lution phenomenon we ask a simple question: Can residual networks easily learn to recover the input? This should be a basic task expected of any generative model - otherwise there would be information loss. However, if our dilution hypothesis is true, the answer would be negative. To test this, we create neural network consisting of number of low-rank layers, and add residual connections in order to mitigate the bottlenecks introduced by the low ranks. The resulting model full-rank. We compare this model with another model that employs learnable residual connec- tions, as in DenseFormer (Pagliardini et al., 2024), which we later also call GRN-v1, since it is the starting point of our generalizations. In Figure we see the results of the two models on two tasks: learning the identity transforma- tion and learning a random linear transformation. Perhaps surprisingly, we observe that the residual network is unable
1. Introduction
Residual connections play an important role in modern neu- ral network architectures because they stabilize the training of deep neural networks and improve model convergence and quality. Since their usage in the ResNet architecture (He et al., 20) residual connections have been widely adopted in both convolutional neural networks and transformer archi- tectures across various domains, including natural language
'Department of Computer Science, University of California, Irvine, USA Marshall School of Business, University of Southern California, Los Angeles, USA GGoogee Research, New York, USA. Correspondence to: Mike Heddes <mheddes udi.edu>, Gang Fu <thomasfu @ google.com>>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

DeepCrossAttentton:: Supercharging Transformer Residual Connections
10
ResNet
10:2
GRN-v1
MSE (log scale)
10 5
MSE (log scale)
10-4
6 10'
10-11
GRN-v1
10
4 x 10
10-17
20
x 10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Number steps
1e3
Number steps
1e3
(a) Learning the identity transformation by minimizing ff(x) - 212 where x is a 100-dimensional i.i.d. normal input and f is a low-rank lincar network.
(b) Minimizing the loss [= 2ll.. where æ a 100-d i.d. normal input, f is low-rank lincar network, and y = Ar + b, where A, b have i.i.d. standard normal entrics.
Figure 1. Training low-rank lincar models to learn the identity and a random transformation. Each model consists 10 lincar layers, each of rank and is trained using mini natchbach SGD,
fully reconstruct the input even after seeing 10s batches (10" examples), while the model with learnable residual weights is able to reach extremely small loss values, even with 100x fewer examples. This confirms that ResNet does not address neural network bottlenecks in a satisfctorry way, even though it learns a full-rank transformation, and under- scores the importance of using learnable residual weights to increase model capacity.
ers to the ambient dimension is below a threshold, which depends the complexity of the target task. In addition, the improvement in this trade-off can itself be characterized as a function of the collective ranks of the layers, ambient dimension and the complexity of the target task. We extend this insight to nonlinear models by working with the notion of bottleneck rank, proposed by Jacot (2023).
We additionally provide empirical results to support the the- oretical findings and demonstrate the effectiveness of DCA. Experiments on language modeling and image classifica- tion tasks demonstrate that DCA consistently outperforms the standard transformer architectures terms of perplex- ity, accuracy and training efficiency, DCA achieves lower perplexity for given parameter budget and training time. Furthermore, DCA exhibits improved training stability, mit- igating the occurrence of loss spikes frequently observed while training large models.
Our contribution. In this work, we propose DeepCrossAr- tention (DCA), new transformer architecture that gen- eralizes residual networks by employing learnable, input- dependent weights to dynamically combine layer outputs, enabling the model to selectively focus on the most relevant information in any of the previous layers and thereby pre- vent dilution of information in the hidden representations. Furthermore, DCA incorporates depth-wist wise cross-attention by enabling the queries, keys, and values in each transformer block to independently combine layer outputs, allowing for richer interactions between layers at different depths. This all achieved with a negligible number of additional pa- rameters, making DCA more effective than increasing the model size (for instance by increasing its width depth).
2. Related Work
Highway networks enable each layer to interpolate dynami- cally between its output fxx and its input æ using a gating mechanism (Srivastava et al., 2015). Residual connections (He al., 2016) popularized the direct flow of information from earlier to later layers using identity shortcuts. These in- novations proved crucial in stabilizing training and allowing for the construction of significantly deeper networks. Build- ing upon this concept, DenseNet (Huang et al., 2017) further enhanced information flow by concatenating the outputs of all preceding layers to each layer's input.
DCA can be viewed mechanism to adapt the model ar- chitecture dynamically for each input token. By optimizing the added parameters, DCA learns to effectively combine the outputs of earlier residual blocks. This allows the model to rearrange the residual blocks from purely sequential to fully parallel and any intermediate combination, without the need for explicit architectural design choices.
We analyse our generalization of the residual network theo- retically by focusing on linear low-rank model. We show that DCA achieves a better trade-off between accuracy and model size when the ratio of the collective ranks of the lay-
The following methods are the ones most similar to ours. They all build on the idea of DenseNet but apply an effi- cient aggregation of the previous layer outputs instead of
2

---

DeepCrossAttiettonn: Supercharging Transformer Residual Connections
concatenating them. DenseFormer (Pagliardini et al., 2024) performs the aggregation as a learned linear combination of the previous layer outputs. To reduce the computational load, they propose to apply their method only on subset of the possible layer connections. Building on Fonserormer, LAuReL (Menghani et al., 2024) presents three aggrega- tion functions, the best performing one applies a learned low-rank transformation to the previous layer outputs before the learned linear combination. Zhu et al. (2024) take different approach with Flyper-Connections. they consider fixed-size stack where layer outputs are added into with learned weight for every slot of the stack. Before each layer, the stack mixed by a matrix multiplication with a learned weight matrix. The input to a layer then obtained by a learned linear combination of the stack, instead of access- ing the previous layer outputs directly. They also present dynamic version of their method where the weights are derived from the inputs.
f+11
St
Et+1
St+2
f
f-1
Et
E
E
Figure 2. Two alternative schematic representations of standard ResNet. The top represents the recursive form, the bottom repre- sents the explicit sum.
G. (8.) ... 46(8600))
Narm V
Matmul
W/
Add
b
3. Method
Ret.U
We start with a detailed exposition of our proposed gener- alizations to the residual network architecture. We present three distinct proposals, each incrementally augmenting the complexity of the network structure. Building upon these proposals, we subsequently introduce DeepCrossAtteniie (DCA), a novel approach to enhance residual learning capa- bilities of the transformer architecture.
Weighted sum
GRN-v3
Rç(X)
Figure 3. Computation diagram of GRN-v3.
3.1. Generalized Residual Networks (GRN)
Notation. We denote residual function by ft R2d Rd where is the layer index and d the feature dimension. As example, in multi-layer perceptron residual network (MLP- ResNet), we have (xx) Vio (W,)) with W. E Rexd V E ddxk and a is a nonlinear function, such as sigmoid or ReLU,, that is applied component-wise. Then, the f-th residual block outputs I++1 (æ), are defined recursively as
We propose three generalizations of ResNets by consider- ing weighted linear combinations of previous layer outputs. The parameters of the modules and the generalizations are all optimized during training using the A optimizer (Loshchilov & Hutter, 2017).
Dimension-inneepedden weights (GRN-v1). We consider simple linear combinations as gt (x G,b,, y= y=G++1br+1
Sc+1 (a) = fx(g(())) + ge(ar).
Using this recursion, the output of the T-th residual block given by
with b, E gtxl which is initialized as all ones and opti- mized with the rest of the model parameters during training. This setting has been previously explored in the Dense- Former paper (Pagliardini et al., 2024).
T
(x+ Σ (x))
Dimension -depeddett weights (GRN-v2). this pro- posal, we allow b, B and consider
with the conventions that go(x) 0 and fo(go(z)) x.
We refer to Figure for schematic illustration.
ge (x = (G, 3 b,)1, y (Gr+1 by+1)1,
An alternative description, which we will use to introduce our generalizations, is the following. For every t, define the stack of layer outputs G2 E dxx as
where indicates the entry-wise (Hadamard) product. Note that in GRN-v1 the same weight vector b, is used for each of the d features. GRN-v2 generalizes this by using different weight vectors for different features, which are all stacked together in a matrix b, E dxd
G. : [
fo (x)) E dx
We then have gx(x) = G.1 and y = Gy 1 in the standard residual network, where 1 denotes the all ones vector.
input-depended weights (GRN-v3), In the next general- ization, we allow the weights to be input dependent. Specifi-
3

---

DeepCrossAttention:: Supercharging Transformer Residual Connections
cally, the weights are given by by + w with by, w. E IRdxt. The first component acts similar to the weights in GRN- v2, it puts different weights on different dimensions of the input. The second component W, is nonlinear mapping of the input features vector æ; but is the same for all the d dimensions. This combination gives us flexibility to have both dimension-dependent and input-dependent weights for slight increase in the number of parameters. GRN-v3 is expressed as
G2 = (x)) .. Fo(R)(()))
GRN
GRN
GRN
Norm
Norm
Norm
Masked Multi-
Cross Attention
Add & Norm
gx(x)) = (Gr 3 (b, + wi))1, w; w; = low y = GGT+ O (bT+1 + (r++))]
Feed Forward
Add
where wi Rd is initialized as all zeros and optimized with the rest of the model parameters during training and a: R-- + R is a non-linearity which is applied entry-wise. In this proposal we consider a to be the ReLU activation. The computation diagram of GRN-v3 is illustrated in Figure
Decoder Block
R2OX)
Figure 4. Computation diagram of a DCA decoder block.
Reducing memory and computation. Since the stack of layer outputs G, grows linearly with the depth of the model, this could lead to significant memory and computational overhead for deep models. Our experiments reveal that GRNs tend to weight inputs and the last few layer outputs the most. An example weight distribution is provided in Appendix Therefore, to increase efficiency, we propose to include only the first and last-k layers explicitly in G,. On the intermediate layers we apply standard ResNet, only involving simple addition. For example, if we set k 2. then G2 contains at most 4 vectors: the model inputs, the sum of the intermediate layers' outputs, and the last two layers' outputs ( and f1-2(9-2(z)). The GRNs then take this modified G; as their input.
4. Theoretical analysis
Motivated by language modeling tasks, we focus on the regime where the size of the training set (n) significantly exceeds the input dimension (n >> d). As we increase the number of model parameters, the representation capacity of the network improves, which helps with reducing the test error. We will be focusing on the the trade-off between the test error and the number of parameters, and argue that our proposed generalizations achieve a better trade-off than the standard ResNet.
We will first study a "stylized" low-rank linear model for which we characterize the test error-model complexity trade- off and demonstrate the benefits of our proposed generaliza- tions. Our analysis elucidates the role of various factors on this trade-off, such as collective widths of layers, complexity of the target task, and input dimension. We then discuss how some of these results can be extended to non-linear models and empirically demonstrate that the insights gained from our analysis are applicable to more complex models.
3.2. DeepCross Atentioonn
The generalizations introduced thus far are generally appli- cable to any ResNet. We now describe our main method which is specific to the transformer architecture. Deep- CrossAttention (DCA) generalizes self-attention by adding three independent instances of a GRN in each decoder block. In this proposal we consider the GRN to be GRN-v3. These three GRN instances are given the same stack of previous layer outputs as their input but return the queries, keys, and values for the attention module, respectively. This enables richer interactions between layers at different depths. Fig- ure shows the computation diagram of a DCA decoder block inside transformer, where the remaining skip con- nections ensure that the inputs are not added to the outputs of the decoder block, but are included in the inputs of both the attention and the feed forward module. Notably, DCA does not modify the underlying attention mechanism, but instead uses GRNs to dynamically compose attention inputs.
Due space constraint, proof of theorems are deferred to the supplementary material.
4.1. Low-rank linear model
Consider the setting where for each sample the response y E RR is given by
A
with E Rd representing the noise. Here A E RRdxd is a full rank matrix.
We consider a network with T layers where fi(z) V.
4

---

DeepCrossAttention:: Supercharging Transformer Residual Connections
(there is no activation). We let rt rank (V,) and define the collective rank т. := Σi= 7t. We assume T, < d, i.e. the collective rank of all layers still is lower than the ambient dimension d.
its test error (model risk) is given by
Risk(A) El(y -))
=E ||LA - A)
+ +o²
We next focus on four architectures: Baseline (where there is no residual connection), ResNet, GRN-v1 and GRN-v2 and characterize the class of models which can be expressed by each of these architectures. We assume each architecture to have layers.
RB(rrace( (A
A)xxx (A - A)
o²
- -I|A- A + to²
where we assumes that E[xx"] I (isotropic features). Since the term o² is constant (independent of model A) we will drop in sequel without effecting our discussion and focus on the excess risk. For a class of models C we use the notation ER" (C) to indicate the minimum excess risk achievable over the class C:
Baseline. In this architecture, there is no residual connec- tion and the model is given by y II V,a. We de- note by Caase the class of functions that can be represented by such architecture.
ResNets. In this case, we have y - -I-- (I V₁)a.. De- note by Cres as the class of functions that can be represented by such architecture.
2 A - |
ER" = min Āec
GRN-v1. In this case, we have y
with br+1 (T + dimmensioaal vector as deseribed in Sec- tion 3. Denote by CGRN-VI the class of functions that can be represented by such architecture.
Note that ER" (Chasse (T)) is obtained by the best r-rank ap- proximation to A and ER" (Cree) is obtained by the best T' -ran approximation to A I, both of which have sim- ple characterization in terms of the singular values of A and A - I, by using the celebrated Eckart- Eckart-Young-We Young-Mirsky theorem. Deriving ER" (CGRN-vI (T)) and ER" (CGenB(T)) are more complicated. the next theorem, we establish upper bounds on them.
GRN-v2. In this case, we have y = (GT+1 @ b1+1)1, where 1 d x (T+1) 1) matrix as described in Section We denote by CGRN-v2 the class of functions that can be represented by such architecture.
GRN-v3. In this case, we have ; (GT+1 O (bT+1 + 1r--))1,, where by+1 is d x 11) matrix and 0r+1 is dimensional vector as described in Section 3. We denote by CCRN-v3 the class of functions that can be represented by such architecture.
Theorem 4.2. Consider the singular value decomposition A I UEVT. For given m E [d], let Um. Em Vm be the top m singular vectors and singular values and define
A = A- 1 Uy Ey ViT where T == >t= TE. We then have
Err" ECro() || |||
Theorem 4.1. For the low rank linear model we have:
Cbase = {x Mæ: rank(M) S min(ro)(a.) Cres (x + (I + M) rank(M) STz).
Err" (CGRN-VI)
< A
1 FF
trace( A)²
d r.
Err (CGRN-22) < A
CGRN-V1 fx + (al- - M)#: Mank(M) S r.
Σ
A
trace( A )
-max
CGRN-V2 = fo + D + M) : (M) 1 r., is diagonal).
d r.
where (Au) are the diagonal entries of A. In addition,
CGRN-V3 fx + (D M) + (()x)a rank(M) S r., D is diagonal, w Rd
Err" Err" (CGRNCC)))a) S A
F max (T,, T2),
4.2. Trade-off between test error and model complexity
Ti trace A 2
1
22 max
r.
π Tld + 1)'
In the previous section, we characterized the class of models that can be expressed by each architecture. Next, study the trade-off between the optimal test error achievable by each model and the model complexity, defined as the number of its parameters.
02
Ta =
2 +
π (d 1)
Here Umax denotes the maximum eigenvalue of A+AT _ trace( A LTT P... and Vmax denotes the maximum eigenvalue of A+AT diag(A).
2
Note that all the classes of models characterized in Theo- rem are linear functions. For linear model x + Av,
2
5

---

DeepCrossAtientio:: Supercharging Transformer Residual Connections
8000
25000
Reduction in test error
4000
G1. 100
Reduction test
15000
G2. d 100
G1. d 200
0
10000
G2, d - 200
G1, 300
-2000
5000
4000
0
-5000
Amn
Figure 5. Gain in the model performance achieved by GRN-v1 and GRN-v2 over ResNet. The plots represents the lower bounds for G1 and G2 given in Theorem 4.3(i)). Observe that the gain at larger dimension d is higher. Left panel shows that the gain decreases as the 10). Right panel shows that the gain increases as the complexity of the target
collective rank of ResNet increases (Amin 5, Amax task (ie Amin/ Amax) increases (Amax 10 and r. 50
We proceed by discussing the model complexity for each of the architectures, in terms of model size. The number of parameters for ResNet given by 2dr., for GRN-v1 is given by 2dr. TT 1)/2, and for GRN-v2 is given by 2dr. (T 12.2. Note that by Theorem 4.2, if GRN-v1 and GRN-v2 achieve better Excess risk-model size trade-off compared to ResNet, then we can make this improvement arbitrarily strong by scaling A (and so A)
(it) Consider CGAN-v1 and CGRN-V2 the class of mod- els that can be expressed by the GRN-vI and GRN-v2 architectures with the same number of parameters as ResNet model with I layers and collective rank T+. Define G1 = ER= (Cres)-ER"" (CGRN-v1) and G2 ER" (Cres) (CGRNN- v2) the reduction in the optimal excess risk achievable by these classes compared to the optimal excess risk of ResNet. We have
In the next theorem, we focus on GRN-v1 and GRN-v2 and provide sufficient conditions under which they achieve better excess risk-model size trade-off. In the second part of the theorem, also lower bound the improvement that GRN-vI and GRN-v2 achieve in excess risk compared to ResNet, with using the same number of parameters.
G1 2 (d T+) (Vd+r+-vd)+ Vd)
>
G2 (d r.) -(V1 + 12 X
1
1
G3 2 d n(d + (o 2 + (T +1111
12
max
- -V1.6 - 1) -
Theorem 4.3. Assume that A O and let Amax and Amin > respectively denote the maximum and the min- imum eigenvalues of A I. Define = Amin/ Amax S Consider a ResNet model with collective rank
Our next result quantitatively shows the reduction in the collective rank one can achieve by GRNs, while maintaining the same test error as ResNet.
Dr Tt.
Proposition 4.4. Consider a ResNet with collective rank Ta- Likewise, GRN-v3 model achieve smaller test error with collective rank F., where F.. = -- 1- < with / (M(1+Ep) - and 1+
r. -EEnr < d. A GRN-vI or GRN-v2 model can achieve a smaller test error with collective rank r', where r'
If
d <1+
(4.1)
=-i r.-ddn"
then GRN-vI achieves a better excess risk-model size trade- off compared ResNet. In addition, if
Eo T(P
T+ (1+ +d x)
(4.2)
4.3. Insights from the analysis
then GRN-v2 achieves a better trade-off compared to ResNet.
Theorem 4.3 allows us elucidate the role of different factors on the gain achieved by GRNs.
Also, GRN-v3 achieves better trade-off compared to ResNet, if
Role of target task complexity. Note that Amin/Amas E [0, 1] is measure of complexity of the target task. Specifically, as x decreases, the matrix becomes closer to low rank matrix, and hence learning with low rank models becomes easier. Observe that the thresholds
S (1 17( 2 n)) 1.6,
(4.3)
where 7)
(x(1+6))-8 7+50 and Eo _( 1j
1+50
6

---

DeepCrossAttiettonn: Supercharging Transformer Residual Connections
given by the right hand side of (4. (4..3) are increasing in 6, i.e., for more complex tasks we see wider range of col- lective rank where GRNs outperforms the trade-off achieved by ResNet. Another way to interpret Theorem 4.3(1) is that for a fixed target task (and so fixed k), if the collective rank r, is above this threshold, the ResNet is already rich enough that it is hard to improve upon its trade-off.
5. Experiments
We conduct experiments on language modeling and image classification tasks to evaluate the effectiveness of DCA and to validate our theoretical insights. For the language model- ing tasks, the performance of DCA is compared against the standard transformer (Vaswani, 2017) on the LMIB (Chelba et al., 2013) and C4 (Raffel et al., 2020a) datasets. Unless stated otherwise, each model has an embedding dimension of 5 512 and an MLP dimension of four times the embedding dimension. By default, DCA uses a stack of all the previous layer outputs as input the GRNs. When DCA includes only the first and last-k layer outputs explicitly in the input stack (see Section 3.1), then this is denoted as k-DCA.
Role of collective rank. Observe that the lower bound on the gains G1, G2. G3 given by Theorem 4ii are decreas- ing in r.. In other words, when the collective rank r. of ResNet becomes smaller, the level of information dilution occurring in ResNet increases, giving GRNs a better lever- age to improve model perplexity with the same number of parameters.
Each model is trained with a sequence length of 128 and a batch size of 2048 over 64 TPUs for 500k steps, totaling 131B tokens. We use the A optimizer (Loshchilov & Hutter, 2017) with 31 = 0.9, B2 = 0.98, weight decay of 0.1, and a learning rate of 0.0016 with 1000 warmup steps and an inverse square root schedule (Raffel et al., 2020b).
Role of input dimension. Note that the upper bounds on To given by (4.1) to (4.3) increase with the input dimension d. Furthermore, the lower bounds on the gains G1, G2, G3 given in Theorem 4.3(i)) also increase with d. Therefore, for larger input dimensions, we have both a wider range for r. where GRNs outperforms the trade-off achieved by ResNet, and moreover, we obtain a larger gain in reducing model error.
Model depth scaling. For the first experiment, we pre-train a transformer and DCA on LMIB. We increase the model depth from to 42 layers and show the relation between perplexity (Jelinek et al., 1977) and model size in Figure 6. The figure shows that DCA obtains lower perplexity for a given parameter budget. Notably, the 30-layer DCA model obtains better perplexity than the 42-layer transformer, making DCA more parameter-efficient than adding layers.
We refer to Figure 5 for an illustration of these trends.
4.4. Extension to nonlinear models
We recall the definition of Bottleneck rank from (Jacot, 2023). For a function f Ω + Rd, its Bottleneck rank, denoted by rank (f, [2) is the smallest integer ht such that f can be factorized as f = h o g with inner di- mension & (i,e, 9 : Ω2 + R and h : R + Rd) It is also closely related to the Jacobian rank of a function de- fined as rf) = max rank[J f(x)). In general, (J" S ra (f), but for functions of the form f= oo A oo & (for linear map A and two bijections & and Φ), we have (f) = rankBx() = rank(A) (A). These two notions of rank satisfy the following properties (Jacot, 2023):
19
Transformer
DCA
18-
Perplexity
17
16
15
14
0.6
0.8
1.2 Number of parameters
1.0
1.4
1.6
rank(f o g) S min(rank(f))
1e8
rank(J + g) rank(s) + rank(g)
Figure 6. Perplexity on LMIB with 6, 12, 18, 24. 30, 36, and 42 layer transformer and DCA models.
Proposition 4.5. Consider an MLP with fi(z) Vip(U(z) with U, E IRiidd V, E gdxr, Denote by r. = Etn Ta the collective rank of the network. We have Chase C {f: rank (f) S min(rr)Fm.)-
First and last-k. DCA can be made more efficient by in- cluding only the first and last-k layer outputs explicitly in the input stack the GRNs (see Section 3.1). In this ex- periment, study the effect of is on 24-layer model's efficiency and quality. Table 1 shows that reducing k speeds up training while only slightly increasing the perplexity. Either small or large l obtain good training efficiency, as DCA then obtains the final perplexity of the transformer in a third of the time. Setting h = 2 results in model with 48%
Cres S {id + f: ranknw(J) S r..)
CARN-VI C {a- id + ff: rank (f) < r-j.
CGRN-V2 C {g: g(x)- = Dx+ f(x): rfn V
Tay D diagonal).
7

---

DeepCrossAttiention: Supercharging Transformer Residual Connections
lower inference latency compared to A: 24, thus setting k to be small results in efficient training and fast inference.
Table 2. Perplexity on LMIB for models of varying widths.
WIDTH
TRANSFORMER
DCA
DELTA
Table 1. Training speed in batches per second, normalized time for method to reach the perplexity of the transformer, and the final perplexity (PPL.) of the transformer and DCA with varying k.
64
45.75
42.94:10.07
-2.82
192
25.49
23.92:0.04
-1.57
384
18,8610.04
17.83=0.04
-1.03
768
14.70:10.04
14.11=0.07
-0.59
1024
13.61+0.01
13.2240.06
-0.39
METHOD
SPEED
TIME
PPL
TRANSFORMER
8.14+0.18
1.00
15. 15.1440.000 06
I-DCA
5.62:10.04
0.33
14.48:10.05
and 8-DCA models of increasing size on the C4 dataset. The results in Table 3 show that DCA consistently outperforms the standard transformer model. The absolute improvement in perplexity decreases for large models, which is consistent with the width scaling results. The perplexity throughout training is provided in Appendix G.
2-DCA
5.39-10.06
0.33
14.41100.04
4-DCA
5.01+00112
0.37
14.50++0.03
8-DCA
4.35:10.14
0.47
14.49:0.002
16-DCA
3.86:10.08
0.40
14.35::0.07
24-DCA
3.72:10.08
0.39
14.35:10.00
Training time. The effectiveness of a model architecture heavily depends on its training efficiency. Figure 7 shows the training time-perplexity trade-off for 24, 36, and 42 layer transformer and 2-DCA models. The figure shows that 2-DCA achieves better perplexity for given training time, highlighting the training efficiency of DCA. The training time versus perplexity results when DCA uses all previous layer outputs in the GRNs are provided in Appendix F.
Table 3. Perplexity on C4 for models of varying depths and widths.
D
W
PARAMS
TRANSF
8-DCA
DELTA
9
771
75M
27.876
26.443
-1.443
18
771
124M
23.013
21.810
-1,203
13
1111
179M
21.570
20.461
-1.109
18
1111
234M
19.756
18.824
-0.932
18
1600
449M
17.166
16.764
-0.402
Layers
Retrofitting pre-trained models. Since our method is iden- tical to standard residual network at initialization, adding DCA to a pre-trained model does not alter its function. In Table 4, we compare continuing training the pre-trained model with adding DCA to the pre-trained model. Incor- porating DCA results in perplexity improvement of 0.19 after 60k extra training steps, compared just 0.02 for the transformer. Thus, pre-trained models with residual architecture can also benefit from incorporating DCA.
24
36
42
pixiettyy
16
Model
Perplexit
Transformer
15
2-DCA
14
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
Training time (s)
le5
Table 4. Perplexity on LMIB for extended training of 6-layer mod- els. DCA added to 500k steps pre-trained transformer.
Figure 7. Perplexity LM1B versus the training time with trans- former and 2-DCA models of various depths.
STEPS
TRANSFORMER
DCA
DELTA
500K
18.98+0.001
18.98+0.01
0.00
Model width scaling. Our theoretical results indicate that the benefit of GRN is inversely related to the rank of the model. With this experiment, we validate whether the theo- retical results carry over to the transformer architecture by varying the model width. Table 2 shows the final perplexity of a 12-layer model with an embedding dimension ranging from 64 till 1024, pre-trained on LMIB. The delta col- umn, with the difference between the transformer and DCA, shows that the benefit of DCA is reduced the width of the model increases, which is consistent with our theoretical results. These results are in contrast with the depth scaling results, where the improvement of DCA is maintained for deeper models.
500K + 20K
18.96=0..02
18.81+0.01
-0.15
500K + 40K
18.9640.01
18.7940.03
-0.17
500K + 60K
18.96-00.01
18.7944.004
-0.17
a
Training stability. The occurrence of loss spikes is prob- lem when training large models as they can disrupt an ex- pensive training run (Chowdhery et al., 2023). Figures and 8, we indeed observe clear loss spikes with the trans- former model. Interestingly, training DCA is more stable, showing no significant loss spikes even for large models. This constitutes an important benefit of DCA.
Comparison with related work. We compare the perplex- ity of DCA with those obtained by the recent related works LAuReL. (Menghani al., 2024), DenseFormer (Pagliardini
Model scaling, For this experiment, we train transformer
8

---

DeepCrossAttiettinn: Supercharging Transformer Residual Connections
al., 2024), and Hyper-Connections (dynamic) (Zhu al., 2024) in Table 5. DCA improves upon the prior best method, hyper-connections, with difference in perplexity of 0.59, which is the biggest improvement among the methods.
classification accuracy, demonstrating that DCA effectively
generalizes the vision domain.
Table 8. Loss and Accuracy on ImageNet classification.
Table 5. Perplexity (PPL) and parameter count on LMIB using 6-layer model, comparing DCA with related work.
METHOD
Loss
ACCURACY
VIT
0,5698
76.4
VIT + DCA (OURS)
0.5284
77.1
METHOD
PARAMS
PPL
TRANSFORMER
49.65M
08.9
LAUREL-PA
49.75M
18.99+0.05
6. Conclusion
IXI-DENSEFORMER
49.65M
18.80:0.11
COPEE-CONNESTIOOS
49.68M
18.65:0.03
This paper introduces DeepCrossAttentioo (DCA), , a novel transformer architecture that enhances the flow of informa- tion across layers. achieves lower perplexity for a given parameter budget and training time for minimal increase in model parameters. DCA enables dynamic interactions between layer outputs by building on three generalizations of the standard residual network (GRN). We showed theoret- ically that GRN obtains a better test error-model complexity trade-off. In our DCA experiments we observe significant improvements in model stability, convergence, and quality.
DCA (OURS)
49.73M
18.06+00.01
Table 6. Perplexity (PPL.) C4 using 13-layer model and em- bedding dimension 1111, comparing DCA with related work. The baseline model has roughly 179M parameters.
METHOD
PPL
TRANSFORMER
21 11433
LAUREL-PA
20.951
IXI-DENSEFORMER (STACK SIZE=4)
21.168
21.077
Acknowledgment
AOPPER-OONNCTIONN (STACK SIZE=10)
20.718
Adel Javanmard is supported in part by the NSF Award DMS-2311024 the Sloan fellowship in Mathematics, an Adobe Faculty Research Award, an Amazon Faculty Re- search Award, and an iORB grant from USC Marshall School of Business. The authors are grateful to anonymous reviewers for their feedback on improving this paper.
8-DCA (OURS)
20.392
Ablation study. To determine the relative gain of each of the proposed generalizations, in Table 7 we show the per- plexity obtained by each method described in Section 3. The GRN versions use one GRN instance per decoder block. DCA, in contrast, uses three independent instances of GRN- v3 per decoder block. The biggest improvement in perplex- ity comes from GRN-v1, followed by DCA and GRN-v2.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences our work, none which we feel must be specifically highlighted here.
Table 7. Ablation study of DCA, showing the parameter count and the perplexity (PPL. on LMIB with 6-layer model.
ABLATION
PARAMS
PPL
References
TRANSFORMER
49.65M
18.9810.002
GRN-vI
49.65M
18.8010.11
Beyer. L., Zhai, X., and Kolesnikov, A. Better plain vit base- lines for imagenet- 11. arXiv preprint arXiv:2205.0/580. 2022.
GRN-V2
49.66M
18.4310.04
GRN-V3
49.68M
18.41+0.10
DCA
49.733 .9.73M
18.06=0.01
Chelba, C., Mikolov, T., Schuster, M., Ge, Q., Brants, T. Koehn, P., and Robinson, T. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv: 2013.
ImageNet classification. In addition to the language mod- elling experiments, we also experiment with image classifi- cation using the ImageNet dataset and the vision transformer (ViT) model (Dosovitskiy et al., 2021). Since the ViT model is transformer-based DCA can be incorporate in the same way as for the language models presented carlier. In Ta- ble 8, we present the results on the ViT-S/16 model (22M parameters) and follow the experimental setup by Beyer al. (2022). The results show a 0.7% improvement in
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.

---

DeepCrossAttinttinn: Supercharging Transformer Residual Connections
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., and Zhai, X. An image is worth 16x16 words: Trans- formers for image recognition scale. International Conference on Learning Representations, 2021.
Zhu, D., Huang, H., Huang, Z., Zeng, Y., Mao, Y., Wu, B.,
Min, Q., and Zhou, X. Hyper-connections, arXiv preprint arXiv:2409, 19606, 2024.
Gong, Y., Chung. Y.-A., and Glass, J. Ast: Audio spec- trogram transformer. In Interspeech 2021, pp. 571-575, 2021.
He, K., Zhang, X., Ren, S., and Sun, Deep residual leamn ing for image recognition. In Proceedings of the IEEE conference on computer vision and puttern recognition, pp. 770-778, 2016.
Huang. G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
Jacot, A. Implicit bias of large depth networks: notion of rank for nonlinear functions. The Eleventh International Conference on Learning Representations, 2023.
Jelinek, F., Mercer, R. L., Bahl, L. R., and Baker, J. K. Per- plexity measure of the difficulty of speech recognition tasks. The Journal of the Acoustical Society of America, 62(S1):S63-S66 1977.
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. arXiv preprint 1rXiv:1711.0511 2017.
Menghani, G., Kumar, R., and Kumar, S. Laurel: Learned augmented residual layer. In Workshop on Efficient Sys- tems for Foundation Models I1, 2024.
Pagliardini, M. Mohtashami, A., Fleuret, F., and Jaggi, M. Denseformer: Enhancing information flow in trans- formers via depth weighted averaging. arXiv preprint arXiv:2402.02622. 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21 (140):1-67, 2020a.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020b.
Srivastava, R. K., Greff, K., and Schmidhuber, Training very deep networks. Advances in neural information processing systems, 28, 2015.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
10

---

DeepCrossAttention: Supercharging Transformer Residual Connections
A. Proof of Theorem 4.1
We restate each of the claims in the theorem statement, followed by its proof.
Coase
{z Mæ: rank(M) I min(rr))=1).
Note that by the inequality rank(AB) Σ min(rank(A) (annk())r,m)))) rank( if M is of the form I V. then rank(M) -
min(r:) 1=1- For the other direction consider any matrix M with rank (M) ro S min(r:) and its SVD as M
PSQ' with P; P.QE ædxro with full column ranks. By setting, V1 PSQ and V2 =... VT QQ' have T
M
1==1 VI, because QTQ I and also rank(Vi) ro I min(r.) S rt.
Cre fx + (I + M) Ma)k(M) rr+}
We have
T
T VI
vI)+ = vi II Vi)
T
T
T
= vi
[T V/ +
IT+
V.
+ (I+
t=
1=
T-1
= I I+ VT +
EV (+V)
t=
T=t+1
Note that each of the summand is of rank most rt, it can be written I + M with rank(M) s Ein Pt. Hence
Ih (I + Vi)== x E Cres.
We next show that any I + M with rank(M)= :- S - 12 can be written as II (I +I+ V₁) with rank(V.) S for t e [T]. We show this claim by induction. For the basis (T = 1), we can take V1 = M. To complete the induction step, we need to find V Raxd such that rank V) rr and (I V)- (I M) I is of rank at most Σ[= rt. Then by the
induction hypothesis, we can write
T-1
(T + V) M)
II VV), 1=1
with rank(V; rs, which completes the proof. Without loss of generality, assume TT r; otherwise can take V7 = M and V, = 0 for t < T - 1.
To find such V we write M = =PQT with P,Q E Rdxr having full column rank. Define P,, Q E [Rdxrrr obtaining by considering the first FT columns of P and Q. Additionally, define
B PI (TP)) C = (=Q
(A.1)
We next construct V by setting V = BC Clearly, rank(V) PT. We also have
(I (I+ M) I=
II BC+ + + BCT)-1-1 = -I+ (I+BCC) (PD+
+ 1 B - -
-
Here we consider the notation P PA1] and Q [Q][Q=1) The second step above follows from the Woodbury
matrix identity. Rearranging the terms have
(I ++ V++ I (I BC - QT + (I+ -BU
(A.2)
The first term above is of rank at most rank Pank(P-.) (P-1) - rT 5 -t=1 Pr. We next show that the second and the third term cancel each other. Equivalently, we show that
T-1
PiQ = (I (C+ BB
(A.3)

---

DeepCrossAtientio:: Supercharging Transformer Residual Connections
To do this, next show that
P = (I BB Q -I=
(A.4)
Recalling (A.1) we have P1 B(I + P.). Also
C B = (I (1++FF))
=I+ Q -1------
= -I+
(A.5)
= -1+
P
(A.6)
Therefore, Pi = B) (I BC )BC Likewise, recalling (A.1) we have Q1 = C(I QTQ:)
)- Hence,
= = (1==--
using (A.6). This completes the proof of (A.4) and so (A.3).
Invoking (A.2) we get
(I + -C-
-
which is of rank at most r TT S ETe T-1
Tt, which completes the proof of the induction step.
CGRN-V1 = fx + (aI - M) rank(M) S r.
We prove this claim by induction. The induction basis (T = 0) follows readily since Gib; = byx. Assume the induction hypothesis for We have f((g(())) ViG,b, and so G [VGG,b, | Gt). Writing b bw1
bj
we obtain
Getibit1 V.G.b,bi + G.b.1
By induction hypothesis, G.b-1 is the set of functions of the form (al + M) with rank (M) S >E=r re.
Since rank(V) T; the set of functions that can be represented as G is subset of (al M)) with rank(M) < Ei=1 Te. Conversely, any given M of rank Ei=1 Te can be written M Mi V with rank(Mi) Σ rz and rank(V) rt. By induction hypothesis, (aI + M, a can be expressed by the tenm G, In addition, Væ can also be expressed by the term V,G,b,b,, by taking V, V. b, (0,0, 1)T, bi 1, which is possible since they are free from the choice of b
Hence, 1b41 the set of functions of the form (al + M) with rank (ank(M) s Ei Tz, completing the induction step.
CGRN-V2 (x (D + M) rank(M) rr,, D diagonal}.
The proof follows similar to that of GRN-v1. For the induction basis (7 0), have (G1 (Cr (#:bb) diag(b, )ar.
Assume the induction hypothesis for t. We have ((((())) = V.(G. O b,)1 and so G.+1 [V((G, @ 6,)1 Gr). Writing b
(22 B ((4)
we obtain
G O b.t1 V.
and hence
(1) V. O b.)1 (G. о (~1) 1. (G++1 O b++111 di
By induction hypothesis, (G, o (-11) ) is the set of functions of the form (D + M) with rank(M) < E re. By varying Vr, b and b("), the term diag( 6 V. @ b,)1 covers all functions of the form Væ with a d x matrices of rank r, (for given of rank rt, take V. V. b, = [ojo] ... [1], ( 1, which is possible since they are free from
12

---

DeepCrossAttention:: Supercharging Transformer Residual Connections
the choice of 1 Hence, (G++1 @ b is the set of functions of form (D + M) with rank(M) s Ei=1 " completing the induction step.
C (D M) (w rank(M) S Tr, D is diagonal, R
We prove that
T +
CGRN-V3
Σ Vi + aww x) aww))) + Dr, rank(V.) = rt. D is diagonal
we show the claim by induction on the number of layers. For the basis T = we have
(G1 @ (b1 + ww))1 = diag(bi + wi ) = diag bb) )z o(w]x)x
Assume the induction hypothesis for f. We have fz())())) = Vi(G. o c. with b. 1o(w Go) and G.+1 [ViGG, o 1 Gr). Writing c++1
[c4) [
(~1)
we obtain
G1+1 @ C++1 = Ming
and hence
(G++1 @ c++1 1 = diag( diag(c(4) ((1 Vi o 04)1 (G, o C+1 1.
(~1)
We take b, [0/0 [1] and w; = 0, and so [0 ... [1]. We also take C1+1 1. So,
(G++1 @ c+111 = Va + (Ge @ )1
By induction hypothesis, (G, @ (+11) covers all the functions of the form - Ee= Vaa oww + Da, which along with the above equation completes the induction step.
Finally, note that any matrix M Raxd of rank r, = Ei= re can written as Σ Vi, for some choices of matrices V, E Rgdxd of rank Tz.
B. Proof of Theorem 4.2
By Eckart-Young Boong-Mirskee theorem, U,. E,, VT is the best rank r* approximation to A-I, by which we obtain ER" (Cre))
| p.
We also have by definition,
ER" (CoRN-v1)= o.rank( m)
min
al À
(B.1)
F
Recall the SVD of UEVT and consider the following decompositions:
Σ,- 0
U U= U.. Ur...h V Vv. V...a):
Σ
0
Σ.
with Ur..L, ViL.d E (dx(ddr.), and E,.. a diagonal matrix of size r.. Since U is unitary matrix, we have
U,, U.. +
UT U II
We then that for any choice of a, Á, have
A- -= A I (1-a)I À
= A + U,, VoT + (1- - À
A U, E,, V +1- U.. + UU,, UT Á.
13

---

DeepCrossAtiention: Supercharging Transformer Residual Connections
Next, by taking Á = U, Vi +1 U, [,, V o.) the rank-r, matrix, we obtain
A aI A A (1 a) U,. U,
Invoking the characterization (B.1), we arrive at
ER*
ER (CGRN-VI) S min
+ (1 - a)) U r.dd 1F
2 ā² | -.. r..d 1| 2âtrace U
min
min
A
2 (d -r,)à² 2âtrace A)
1 trace A)²
A 2
where in the second equality, used the fact that U is unitary and so U r..l ||= =d rw. In addition, observed
that
A = A
U. Er V.T = UEV U. Vr. U.,, Era,l V.T ViLd
Therefore,
U T. Vr,, Er,, U r.,d U. U. - =V. Σ,, UT To,
AT.
and so Irace( A AU... Ur,!! U U... trace( A trace( A). This completes the proof of the upper bound on ER" ER" (CGRN-v1).
For ER" (CG we have
ER (CGRN-V22 2
min diagonal
A D U, E,, E,
V.TlI
min dingonal
| +
DI
- D
= min diagonal
(B.2)
In addition, since GRN-v2 optimizes over a larger class of models (using diagonals instead of scale of identity), we have
ER" (CGRN-v2) s ER" ER" (CGRN-vI) || r-r+ T
trace(A) (
(B.3)
Combining (B.2) and (B.3) we obtain the claimed upper bound on ER" (CGRN-v2).
We next proceed to bound ER* (CGRN-v3). By Theorem 4.1, this quantity is at most the optimal objective value of the following optimization problem:
min E ||AT Dx Mx aw x) x)x [2:]
(B.4)
M.M. m
subject to
rank(M) S r., D is diagonal, Rd
We caleulate the expectation in the objective over the gaussian vector x N(O, I). Define the shorthand B = A- M.
We write
E [|IBr Bx - ow x))]= = |B |B + Eoo w )= ||=l 20 T x)xx Bx]
B } Blo ( x))² || 2trace (B Blo x)xx [B)
(B.5)
These two expectations are characterized by our next lemma.
14

---

DeepCrossAtiention: Supercharging Transformer Residual Connections
Lemma B.1. Suppose that N(O, I) and let (2) =1 0) be the ReLu function. Then,
d+1
Blo(w'x)) x)²² txlleg] || 1
(B.6)
2
ww T 2 l1
Blo(w" (/tw)/))<]]1
||wlle2
I+
(B.7)
2T
Using the result of Lemma B.1 in (B.5) obtain
d ++1 |wlle, trace (B) + w Bw
E Bx ow ())
|B + 1
(B.8)
2
| ll.
With this characterization, we next proceed to calculate the optimal objective value of (B.4). We start by minimizing over w. To do this, we first fix ||wle = a, and optimize over the direction of w, and then optimize over a. Note that 2 is symmetric, the maximum is achieved when w in the direction of its
w Bw wh (B + B 3/w.. Since (B + // maximum eigenvalue. Define X as the maximum eigenvalue of (B + B We then have
+ a² d +d+1
2
min E Bx olw (tw))))=l(
= min |B| azu
V
(tracel B) + --max)
2
(trace(
B) + maxx )
B
(B.9)
xd + 1)
We next continue with minimization over M, D. Since we want to derive upper bound on the minimum objective value, we consider two choices of (M, D) motivated by the analysis GRN-v1 and GRN-v2.
Choice 1: Similar to the analysis of GRN-v1,
M U- E., VT. (1 a)U,,] and D al with
1 + d-r, --,, trace( A). With these choices have
B = A M - D
= A - I + (1--a))
A (1 - U, U,, (1 (1-a) I
A + (1 a) Ur.,1 U
A
trace A AUU U.T U
d d-
In addition,
Ur,, p.,d d and trace( A AU U.TT U...) = trace A). Hence,
B 2
A
1
trace A)
Furthermore, trace(B) B) = 0. Using these identities in (B.9) we obtain that the optimum objective value of (B.4)
satisfies the following:
OPT < A d d-
trace (A)
max
(B.10)
(dd +
with Vmax denoting the maximum eigenvalue of A+AT
trace A U UT r..d T-r.
Choice 2: Similar to the analysis of GRN-v2, we set Ur Er, VT and diag I + A). This way we have
B = A M D D
= A I M -
A
A diag( A)
15

---

DeepCrossAttietioon Supercharging Transformer Residual Connections
Hence, ||
| Ei=1 A and Irace (B) 0. Using these identities in (B.11) we obtain that the optimum
objective value of (B.4) satisfies the following:
2 (d +)
OPT S |
with i denoting the maximum eigenvalue of A+A!
diag( A)
2
Combining the bound from the two cases, we get
1/
D2 max
Err" (CoRN-v3) < OPT < ||| max
trace
max
A
+
d- d-
ald+1)
i=1
which completes the proof of theorem.
Proof. (Lemma B. Since the distribution of æ is rotation invariant, without loss of generality we assume w = =1w21ee, e1. We then have
(lo(w"))) Elo(w a)) || || Blo(())) || 11,2)
(B.11)
where æ=11 (22)..., We have
3 = =
= 0)] =
Also, since Φ.11 is independent of 01 have
d -
E (x)( |a 112 (z)
Combining the last three equations, get
d+1 1
Bio(w >|=12111,,11
To show the other claim, we note that
0
ifi
E
E[r
V2n
i
(B.12)
2
ifi=j=1.
Therefore matrix form we have ([(((1))) [=]
1 (I + ele)). This completes the proof of (B.7) by rotation
invariance of distribution of æ we can assume = |wlle el
C. Proof of Theorem 4.3
Let p := 2dr, where we recall that r, = EE= Ty. Note that p is the number of parameters for ResNet with layers and ranks r, for each layer We will compare the test error of GRN-v1 and ResNet with p number of parameters. This
corresponds to model in GRN-v1 with 7" layers such that 2d X f=1 + (T" -1)/2 = p. We set the shorthand T/ = Em= Te and let 01 2 ... 2 σa be the singular values of I. By Theorem 4.2 we have
d Σ
ER" (Cros)
of
ER" ER (R 5
Σ
1
o.
rd
i=.+1 +1
t=
16

---

DeepCrossAtiention: Supercharging Transformer Residual Connections
Therefore, ER" (CGRN-vV < ER" (Cres) if the following holds:
Σ o d r
Σ 0:
(C.1)
(Note that r4 < Ts since T < T). However note that the left hand side of this condition is upper bounded by
of s (r= - r.)
+1
Additionally, the right-hand side of the condition is lower bounded by
d Σ
2 (d - rr))
σi
d
1=r +1
So sufficient condition for (C.1) is that
(r, - - S (d r'))
Writing it in terms of K, we need
S dx² + +r, (1-11111 - x²2).
(C.2)
Our next lemma gives alower bound on r',.
Lemma C.1. Consider a standard Resnet model with collective rank r., and also a GRN-vI model with collective rank r', a GRN-v2 model with collective rank r", and GRN-v3 model with collective rank n which have the same number of parameters as the standard Resent model. We then have
To (Vd+r, va)² S ST, ST:: r. (v ST. ST..
(C.3)
(C.4)
-VI.VV++ - 1)² < rM ST.
(C.5)
Using Lemma C.I, condition C.2 is satisfied provided that
r. S dr. + (1 (1- 101r.-((vrrr. -vae]
Solving the above inequality for r and after some algebraic calculation, we simplify the above inequality as follows:
r. (1 (d++n(vn nVV + x))²
For GRN-v2, the argument goes along the same lines. Fixing number of parameters to this corresponds to model in GRN-v2 with 7 layers such that 2d ET dT" 1)/2 p. We use the shorthand r" = ΣE re. By Theorem 4.2,
1
ER" (CGRN-v2) = A|
trace A)
d rM
1
Σ
o
0
7
imr"+1
Following the same argument as the one for GRN-v1 (replacing r4, with r!)) we derive that GRN-v2 achieves a better trade-off than standard ResNet,
T. dk3 11 (1)-)))) x2).
(C.6)
17

---

DeepCrossAttention:: Supercharging Transformer Residual Connections
(Note that this is analogous (C.2) where r' is replaced by r!))
Using Lemma C.I, condition C.6 is satisfied provided that
S dn + (1 x²) (v By some algebraic calculation, this inequality can be simplified to
7+ S (1+
We next proceed with the case of GRN-v3. By Theorem 4.2 have
1
1 max Tld + 1)
ER* (CGRN--W S ||||
tracee A)
d-r.
with V/nax denoting the maximum eigenvalue of A+A eigenvalues we get
d-r. trace(
A) U P Rewriting this bound in terms of
2
1
1
ER" (CG(NNGN-
o
!!! Trm²11 d imp!!!
0
x 1)
Here we used the fact that the Up is the eigenspace of A and its eigenvalues are 0p1+1 2 ... od. To lighten the notation, we use the shorthand a = 011+11 and A: Ei oz Then in order to have ER" ER"RRCR""N---"" < ER" (Cres), it suffices to have
1 A d r
d- A)2.
t Tld +
Note that the left-hand side is upper bounded by (r, - )o. jo addition, this is quadratic inequality in A. Solving this inequality for A this corresponds to the following:
xdd+1)d-r")
d-r""
add+)(d r?)
(d-+")
Define the shorthand { = x(++)) Then the above can be written
E+ piu
(1+
d
(C.7)
(1 ++
We also have 2 (d-r""")Amm and S Amaxi so A/o 2 (d rM )r. Hence, the above condition holds
r1- d- (1 + E) -
E +
(C.8)
1+
is easy to verify that the right-hand side is decreasing S. In addition, by definition & and since r" < r. S d, we have
{ z Eo x(1=-1)- So sufficient condition for (C.8)
PM Eo
Eu d- p:r
(1 Eu)
k >
(C.9)
1 Eo
equivalently,
(k(1 + {o) - -oo) + Eo
1
1 Eo
d !!!
18

---

DeepCrossAttention:: Supercharging Transformer Residual Connections
Define 17
-6)+ 1+50
Rewriting the above inequality we need
(1 - ,/²) + drp² 2 r..
(C.10)
Next, by Lemma C.1, we have ru r"" S (VT.6 r+ -1)2, by which sufficient condition for (C.10) is as follows:
[r- - 1...-+r-111 (1 (f= (+)++ + dn²² 2 n..
By some algebraic calculation, this inequality can be simplified to
r. S (1+n n)) 1 1.6
This completes the proof of the first item in the theorem statement. To prove the second item in the theorem statement, we write
G1 == ER" (Cres) ER" (CGRN--VV)) 2 d rd
Σ
i=p +1
i=
2 (d - r
(r+ rer) (
d r. r. + r Xpmin))
2 dx2 min
- r. r.>,aax + (r. - -v -V - -2min)
(d Am) (vd r. VāT) Ammin)
where in the last inequality we used Lemma C.I.
A similar bound can be derived for GRN-v2, replacing r, with r in the argument. Specifically, we have
G2 ER" (Cres) ER" (CGRN-v2) dD2 r r. r rr"(22aa Xmnin)
2 d12 min
-re2 + (r. (VI - -1)3) 12 2min)
= (d r.)
-r, (2) - Xmiaa,
where in the last inequality we used the lower bound given for y" in Lemma C.I.
For GRN-v3, we have
G3 = ER" (Cres) ER" (CGRN-V33
2
1
Σ
Ti
2
+
2
max
r,
d - T!l
π
+ 1)
imp!!t +1
imp
1
1
(0
A)2
of
d
m!
+
d p!!!
imp!! +1
1
1
'02
2
(d - -!!!???? A2 - (r. - r )02
-(r-
d
T(d +1+1 1)
2
1
(1-
A² 1 2n (cl + 1)
1
02.
d
(ld + 1) k!/)
where we recall the shorthand A : El o2 and a == d, In the second inequality above, we used the fact that t. Sr. 5 = a, along with the inequality (a b)2 2 a / b2.
19

---

DeepCrossAttenttonn: Supercharging Transformer Residual Connections
Noting that A 2 (d r= Min and Amax continue from the above chain of inequalities, follows:
G3
d 1e
X
ra -r, !
2
T(d + 1)
2T (T + 1)
IV
Xinn
1x
1 + (>"mx A Xmmnn
π + +1)
2T 2T + 1)
1
Amin
A + (r+ 1.V16+++-1111) - mmin))
1)
2T(d+ +
क (d 1) [.
Amin +
r. - -1) - -Main)
2n + 1)
max
where we used Lemma C.I in step (a).
C.1. Proof of Lemma C.I
A standard Resnet model with collective rank r. has 2dr, number of parameters. A model in GRN-vI with collective rank r, has 2dr + (T" - 1)/2 parameters. Therefore, by assumption
2dr'_ + (T' - 12/2 2dr.
(C.11)
Since each layer has rank at least one, we also have r. 2 T' We define the shorthand & = (T' 1) (so r' 2 f Combining these two inequalities and writing them terms of &, we get
2d + / < 2dr..
Solving this inequality for we g 2v + dr, 2d. Using this bound in (C.11) get
2dr, s 2dr + 2( + dr. - dd2
Simplifying this inequality, we arrive at
r. -v-trarr-- sr.
The upper bound r', S rs also follows simply from (C.11).
For GRN-v2 model we follow the same argument. A model in GRN-v2 with collective rank r" has 2dr + dT -1)/2 parameters and so
2r" T" 1)/2 2r.
(C.12)
Since each layer has rank at least one, we also have r > T Define the shorthand E' = (T"("" 1). Combining the previous two equation, we get
2 + //2/2222222 < 2r.
Solving this inequality for g' we get E' 2 r. -
2. Using this bound back in (C.12) we obtain
2r. < 2r + +2(v1+r,-11 1
This simplifies r, [. T!
We follow the same argument for GRN-v3. A model in GRN-v3 with collective rank nd has 2dr + 4dT (T"" -1)/2+ "TT parameters and so
2r - (T" + 1)/2 2r..
(C.13)
Since each layer has rank at least one, we also have y! T Hence,
T12 + 5T 4r. S
20

---

DeepCrosAAtientio:: Supercharging Transformer Residual Connections
Solving for T get T < 11/2(v25 16r, - -5). Using this bound in (C.13), have
r" 2 r. 16
(v25 + 16r, 5)( 25 + 16r, 3)
Tx (40 + 16r, - 28 16r.)
16
(v25 16r, 4)2
2 r. 16
2r. (VI.6 + 2
The upper bound r" < r. follows easily from (C.13).
D. Proof of Proposition 4.4
The result follows from conditions (C.2), (C.6) and (C.10) which respectively provide sufficient conditions for GRN-v1, GRN-v2 and GRN-v3 to achieve smaller test error than a ResNet model, with the same number of parameters.
E. Proof of Proposition 4.5
The proof similar the linear case by induction on T. Note that for showing this direction (Crases Cres, CGRN-vI- CGRN-V2 being a subset of the rank constrained functions) only used the following two properties of the rank function which holds also for the Bottleneck rank: rank(f a g) I min(rank(f)) rgn)(g)) and rank(f + g) I raak(f)) + rank(g).
21

---

DeepCrossAtiention: Supercharging Transformer Residual Connections
F. Training time versus perplexity on the LM1B dataset
This appendix provides additional results on training time versus perplexity for DCA models. Figure shows the training perpleettyy trade-off for 12, 24, and 36 layer transformer and DCA models trained on the LMIB dataset. The figure shows that DCA achieves better perplexity for given training time (except for the first few training steps of the 36-layer model). Thus, highlighting the training efficiency of DCA.
Layers
17
Perplexity
16
Model
Transformer
DCA
14
1.5
2.5
Training time (s)
1e5
Figure 8, Perplexity LMIB pre-training versus the training time with transformer and DCA models of various depths.
G. Steps versus perplexity on the C4 dataset
This appendix provides additional results training steps versus perplexity for 8-DCA models. Figure 9 shows the training steps-perplexity trade-off for 75M, 179M, and 449M parameter transformer and 8-DCA models trained the C4 dataset. The results show the improved model convergence and quality of DCA.
Params
Model
32.5
75M
Transformer
30.0
179M
DCA
449M
Perplextty exity
27.5
25.0
22.5
20.0
15.0
2
Training step
Figure 9. Perplexity on C4 pre-training versus the number of steps with transformer and 8-DCA models of various sizes.
Table 9, Perplexity (PPL) LMIB with and without the model inputs for 6-layer GRN-v3.
METHOD
PPL
TRANSFORMER
20.878
GRN-v3 (LAST LAYERS)
20.301
GRN-V3 (MODEL INPUTS LAST LAYERS)
20.227
22

---

DeepCrossAtiention: Supercharging Transformer Residual Connections
H. Distribution of learned weights
Figure 10 shows the distribution of the learned bias values for each GRN-v3 instance 30-layer model. The layers tend to weight the imputs and the last few layers the most and frequently assign a negative bias for the intermediate layers, indicating that the layers are filtered as result of the ReLU activation. In Table 9, we show that indeed the GRN-v3 model perplexity improves as the model inputs included in addition to the last few layer outputs.
Layer
Layer 3
Layer
Layer
Layer 6
0.3
0.3
0.3
0.3
0.3
0.2
0.2
0.2
0.2
0.2
0.3
0.1
0.0
0.0
0.0
0.0
0.3
0.1
.0.
5.00
0:5
0:50
0.75
1.50
0.0
1.0
1.5
2.0
Layer
Layer B
Layer
Layer
Layer 11
0.3-
0.3
0.3
0.3-
0.3
0.2
0.2
0.2-
0.2
0.2
0.1
0.1
0.1
0.1
0.1
0.0
0.0
0.0
-0.1
-0.1
-0.1
-0.2
-0.2
-0.2
6
0.0
7.5
10.0
Layer 12
Layer 13
Layer 14
Layer 15
Layer 16
0.3-
0.3
0.3-
0.3-
0.3
0.2
0.2
0.2
0.2-
0.2
0.1
0.1
0.1-
0.1
0.1
0.0
0.0
0.0
0.0
0.0
-0.1
-0.1
-0.1
-0.1
0.2
0.2
0.0
5,0
7.5 10.0
10
5
10
10
15
Layer 17
Layer 18
Layer 19
Layer 20
Layer 21
0.3
0.3
0.3
0.3
0.3
0.2
0.2
0.2
0.2
0.2
0.1
0.1
0.1
0.1
0.1
0.0
0.0
0.0
0.0
-0.1
-2.1
-0.1
-0.2
-0.2
-0.2
-0.2
15
15
10
Layer 22
Layer 23
Layer 24
Layer 25
Layer 26
0.3
0.3
0.3
0.5
0.3
0.2
0.2
0.2
0.2
0.2
0.1-
0.1-
0.1-
0.1-
0.1
0.5
0.0
0.0
0.0
0.0
-0.1
-0.1
-0.1
-0.1
-0.1
-0.2
-0.2
10
15
15
10
20
10
20
20
Layer 27
Layer 28
Layer 29
Layer 30
Outputs
0.3
0.3
0.3
0.5
0.3
0.2
0.2
0.2
0.2
0.2
0.3
0.1
0.1
0.1
0.1
0.0
0.0
0.0
0.0
0.0
-0.1
-0.1
-0.1
-0.1
-0.1
-0.2
-0.2
-0.2
-0.2
-0.2
20
Layer
Layer
Layer
Lager
Layer
Figure 10. Distribution of learned bias values on LM1B pre-training with a 30 layer GRN-v3 transformer model. The solid line indicates the median value and the shaded area represents the 90th percentile.
23


Reviews:
[{"review_id": "tCCPOjrH7K", "reviewer": "Reviewer_nMe5", "summary": "The paper introduces DeepCrossAttention (DCA), a new mechanism that stores and uses intermediate features of transformers. DCA enables learnable, input-dependent weights to mix preceding intermediate features, enhancing the model's representation power. The authors also provides theoretical justifications regarding why DCA blocks have higher representation powers compared to standard ResNet architectures. Experiments on LM1B and C4 show that DCA achieves competitive performance compared to vanilla Transformer architectures.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our manuscript. We hope that our responses adequately address the reviewer's concerns. We believe these clarifications and additional results strengthen the manuscript.\n\n> Could the authors provide analysis regarding the memory usage during forward and backward pass?\n\nSince the difference in the number of model parameters is negligible, the main difference in memory usage comes from the number of activations that need to be stored. \n\nLet us first analyse the memory usage during inference. The vanilla transformer only stores one activation tensor which each transformer block adds to in order to compute the residual connection. DCA takes all the previous layer outputs as its input, thus the number of activations it needs to store scales linearly with the depth of the model, which significantly increases the memory usage for deep models. To mitigate the memory overhead, we propose k-DCA, which only uses the model’s input, last-k layer outputs, as well as the sum of the remaining intermediate layer outputs. This reduces the number of stored activations to k+2 times that of the vanilla transformer, which is independent of the model depth.\n\nDuring model training the memory usage of DCA is on the same order as the vanilla transformer. This is because the vanilla transformer also needs to keep all activations in memory to compute the gradients. As the DCA layer acts in a modular way on its input activations, the forward and backward passes can be performed by storing only one additional activation tensor. Based on profiling a 179M parameter model, we see that the peak memory footprint increases from 5.3GB to 6.6GB (this can likely be lowered further since our implementation is not tuned for efficiency). We also note that DCA is compatible with model sharding, which can greatly reduce the memory consumption per chip. Finally, DCA does increase the memory bandwidth because for each layer we need to read the activations of all other layers instead of just one. This additional cost is already factored into the runtime analysis presented in the manuscript which shows that DCA obtains lower perplexity for a given training time.\n\n>  Can the authors extend their comparison to include baselines with larger-scale models to assess the method's effectiveness at higher parameter counts?\n\nBased on the reviewer’s suggestion, we conducted additional experiments with larger-scale models to compare our method against related work. We have started additional runs comparing DCA against all baselines at 179M parameters (n_layers=13, emb_dim=1111, hid_dim=4444) on the C4 dataset. While our experiments are still running, we can provide the following perplexity results for the models at 400K (out of 500K) steps.\n\n| Method    | Perplexity, 400K steps |\n| --- | --- |\n| Transformer  | 21.772   |\n| 1x1-DenseFormer    |  21.313  | \n| Hyper-Connections   |  21.261  | \n| LAuReL-PA |  21.117 | \n| 8-DCA (ours)   | **20.625**   |\n\nThe preliminary results show that even at larger scale our method outperforms prior work. We will add the final results to the paper once all runs are complete.", "claims_and_evidence": "Most of the claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criterion make sense for the problem.", "theoretical_claims": "I did not check the proofs of theorems in the paper.", "experimental_designs_or_analyses": "I have checked the soundness/validity of the experimental designs in the paper.", "relation_to_broader_scientific_literature": "The main contribution of the paper is that it proposes a way to use a small number of additional parameters to improve the performance of the LLM.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength:\n1. The paper is well-written and easy to read and follow.\n2. The paper provides theoretical justifications alongside relatively large scale training experiments to validate the proposed method.\n\nWeaknesses:\n1. The majority of the paper is comparing DCA with vanilla transformers. And the comparison with baseline methods in Table 5 is limited to very small number of parameters, i.e., ~50 M. The authors should extend their analysis to larger-scale models.\n2. There is no discussion regarding memory requirements during the forward and backward passes. It would be beneficial for the authors to report the maximum memory usage in both phases to better check the method's scalability.", "other_comments_or_suggestions": "1. Could the authors provide analysis regarding the memory usage during forward and backward pass?\n2. Can the authors extend their comparison to include baselines with larger-scale models to assess the method's effectiveness at higher parameter counts?", "questions_for_authors": "No other questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Zgv7yswtIx", "reviewer": "Reviewer_zWQ6", "summary": "The authors introduce learnable residual connections to improve over standard residual connections used in ResNets and transformers. They highlight that simple residual connections struggle to recover the input (learn the identity function) on toy examples and their proposed learnable residual connections can overcome this problem. Theoretical analysis on a low rank linear model shows that their proposed method obtains lower risk given that the rank of the task is small enough. Further experiments on transformers are provided to highlight to empirically evaluate the method.\n\n######## update after rebuttal #########\n\nThe authors have clarified my questions, i will keep my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and insightful feedback on our manuscript. We address the three questions raised below.\n\n> Most efficiency gains seem to occur by using the first and last-k layer outputs in the GRN, for k=2. Moreover the perplexity gains from increasing k further are limited. This implies that not all layer representations contribute to the residual connection in GRN. However, the authors do not compare with only a learnt residual in each layer while ignoring all previous residuals (w1xt+w2f(xt)). Would this already be sufficient?\n\nThe method suggested by the reviewer has previously been proposed as LAuReL-LR. \nThe LAuReL-PA method that we used in our experiments is a stronger generalization of LAuReL-LR, as observed in the LAuReL paper. Thus, we opted to report LAuReL-PA in our manuscript instead. Based on our results, including previous outputs in addition to the last output leads to significant perplexity improvements.\n\nWe also conducted experiments without the model inputs as an explicit input to each GRN but we found that this did not perform as well as including the model inputs (num_layers=6, emb_dim=512 on lm1b):\n\n| Method   | Perplexity   |\n| -------- | ------- |\n| Transformer   | 20.878   |\n| GRN (last 4)  | 20.301 |\n| GRN (model inputs + last 3)  | 20.227 |\n\nWe would like to emphasize that even with k=2 all the layer representations do contribute to the residual connections in GRN. This is because all the intermediate layers are summed, as in a vanilla residual network, and passed as an additional input to the GRN. With k=2 each GRN thus has the following 4 inputs: the input to the model, the sum of all intermediate layer outputs, the second last layer output, and the last layer output.\n\n\n> Table 2 presents results where improvements with DCA diminish with increasing width. Can the authors estimate the rank in each layer to verify if similar trends like Fig 5 extend to an experimental setting.\n\nSince the models used in Table 2 incorporate nonlinear activations, the appropriate notion of rank is the bottleneck rank as described in Section 4.4. In our case the bottleneck rank is the same as the width of the model because the model strictly improves as the width increases, which indicates that the model is not able to represent the same function with smaller width. Our empirical results thus align with our theoretical results because in both cases the benefit of our method decreases as the rank (or bottleneck rank) of the model increases.\n\n> An analysis of the importance of layers given the learnt weights of DCA is missing in the experiments. I believe this would be crucial to highlight the differences between standard residual connections and DCA. Are there specific layers that obtain a larger weight in residual connections of DCA and if so which are they?\n\nIn Appendix H we provide insights into the importance of each layer by plotting the learnt weights for each GRN in Figure 10. The results show that the first and last few layers are most important as they obtain the largest weights. This insight led us to the design of the more efficient k-DCA which uses the model inputs and the last-k layer outputs together with the sum of all intermediate layer outputs as the input to the GRN. \n\nMoreover, we observed in the training dynamics that the model input is important for all the hidden layers, especially in the beginning of model training. We further verified this by removing the explicit connection to the model input which performed notably worse. We will enrich the discussion in Appendix H with these insights.\n\nWe hope that these responses adequately address the reviewer's concerns. We believe these clarifications and additional results strengthen the manuscript.", "claims_and_evidence": "Yes, the proposed model architecture has been validated empirically.", "methods_and_evaluation_criteria": "The method is evaluated across multiple datasets.", "theoretical_claims": "The authors show that the proposed GRN model can achieve lower risk.", "experimental_designs_or_analyses": "The experiments are designed to verify the performance of the newly proposed GRN architecture.", "relation_to_broader_scientific_literature": "The paper discusses residual connections and attention, two crucial mechanisms in modern deep learning.", "essential_references_not_discussed": "Most of the relevant literature is discussed.", "other_strengths_and_weaknesses": "Strengths\n\n1. The proposed method is simple and can provide significant improvements as well as faster training as shown in the experiments.\n2. The authors also provide some theoretical insights to validate their method, and show that as long as the rank of the target task is small enough, the GRN model can achieve lower risk.\n\n\nWeaknesses\n\n1. Most efficiency gains seem to occur by using the first and last-k layer outputs in the GRN, for k=2. Moreover the perplexity gains from increasing k further are limited. This implies that not all layer representations contribute to the residual connection in GRN. However, the authors do not compare with only a learnt residual in each layer while ignoring all previous residuals ($w_1 x_t + w_2 f(x_t)$). Would this already be sufficient?\n2. Table 2 presents results where improvements with DCA diminish with increasing width. Can the authors estimate the rank in each layer to verify if similar trends like Fig 5 extend to an experimental setting.\n3. An analysis of the importance of layers given the learnt weights of DCA is missing in the experiments. I believe this would be crucial to highlight the differences between standard residual connections and DCA. Are there specific layers that obtain a larger weight in residual connections of DCA and if so which are they?\n\nWhile the overall mechanism proposed with GRN is simple and can help achieve improved performance, it is unclear which of these layers enabled with residual connection are most important and how the optimal $k$ can be estimated for at task.", "other_comments_or_suggestions": "See above", "questions_for_authors": "See above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wf3JbSFD3b", "reviewer": "Reviewer_rA8T", "summary": "The authors propose DeepCrossAttention (DCA), a novel method that enhances residual connections in transformer architectures. In standard transformers, if we denote the input and output of the i'th block as $x_i$ and $y_i$ respectively, vanilla residual connections simply use $x_i = \\sum_{j < i} y_j$. DCA instead employs learnable, input-dependent weights to dynamically combine layer outputs: $x_i = \\sum_{j < i} y_j \\cdot (\\text{ReLU}(\\langle y_j, w_{ij} \\rangle) + b_{ij})$ where $w_{ij}$ and $b_{ij}$ are parameter vectors.\n\nFor attention layers specifically, DCA creates queries, keys, and values by independently combining previous layer outputs using separate parameter vectors, allowing for richer interactions between layers at different depths.\n\n## Theoretical Contributions\nThe authors analyze stacked low-rank linear projections, demonstrating that DCA achieves a better accuracy-model size trade-off when the ratio of collective layer ranks to ambient dimension falls below a critical threshold. They extend this analysis to nonlinear models using the concept of bottleneck rank.\n\n## Empirical Contributions\nExtensive language modeling experiments on LM1B and C4 datasets demonstrate that DCA:\n- Achieves better perplexity for a given parameter budget\n- Reaches equivalent model quality up to 3x faster\n- Exhibits improved training stability with fewer loss spikes\n- Adds only a negligible number of parameters (approximately 0.2%)\n\n## Strengths\n1. The method is elegant in its simplicity while delivering substantial convergence improvements, particularly for models with smaller hidden dimensions\n2. Comprehensive empirical validation across multiple scales (75M to 449M parameters) using various metrics\n3. Thorough ablation studies isolate the contribution of each proposed component\n4. Clear improvements over related methods (LAuReL, DenseFormer, Hyper-Connections)\n\n## Weaknesses\n1. Experimental validation is limited to language modeling tasks; testing on other modalities (vision, audio) would strengthen the paper's claims about the method's generalizability\n2. The theoretical analysis, while sound, is less compelling than the experimental results; the authors could have expanded the empirical evaluations section instead", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our manuscript. We address the two concerns raised by the reviewer below.\n\n> Experimental validation is limited to language modeling tasks; testing on other modalities (vision, audio) would strengthen the paper's claims about the method's generalizability.\n\nBased on the reviewer’s suggestion we have performed additional experiments on ImageNet classification using vision transformers. Since the ViT model is also transformer-based, we were able to incorporate DCA the same way as for the language models presented in the manuscript. We present our results on the ViT-S/16 model from https://github.com/google-research/big_vision (22M params):\n\n| Method    | Training loss | Accuracy |\n| --- | --- | --- |\n| ViT  | 0.5698   | 76.4 |\n| ViT+DCA |  **0.5284** | **77.1** |\n\nThis indicates that our results generalize to the vision domain.\n\n> The theoretical analysis, while sound, is less compelling than the experimental results; the authors could have expanded the empirical evaluations section instead.\n\nIn response to the reviewer’s feedback, we have significantly expanded the empirical results section of our manuscript. The updated version will include the ImageNet results presented above, as well as additional comparisons with prior work using larger-scale models (see also the response to reviewer nMe5) to provide a more comprehensive empirical evaluation of our method.\n\nWe are confident that these additional experiments and clarifications address the reviewer's concerns and further strengthen the paper.", "claims_and_evidence": "Yes, the claims are supoorted theoretically and empirically.", "methods_and_evaluation_criteria": "Yes LM on C4 and LM1B makes sense but the evaluation could be further expanded.", "theoretical_claims": "I checked the correctness of proofs and did not find any issues.", "experimental_designs_or_analyses": "Yes, the experiments are sound.", "relation_to_broader_scientific_literature": "The proposed method builds on related work and the novelity over past work is not significant - however the authors do show improved reuslts over the methods that they build on.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "-", "other_comments_or_suggestions": "-", "questions_for_authors": "See weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes an alternative to residual connections in Transformers, which dynamically reweights previous layers' outputs based on input dependent weights. On language modeling experiments, it was shown that the proposed approach improves upon model convergence compared to both standard Transformer and related baselines. All reviewers agree that this is a sound idea, and appreciate its simplicity. Questions are raised mostly around the empirical side, eg lack of larger scale evaluations, tasks other than language, analysis on memory consumption etc. The authors did a pretty good in answering these, especially in providing new experimental results. Although it is still highly desirable to further flesh out the empirical side of this work, due to the generality of the claim, this work does already demonstrate promising results and I'd lean towards accepting if there is space.

One additional point, I think the authors should cite and discuss the Highway Networks paper, which is a well established prior work and shares similar design spirit.
Decision: Accept (poster)