DeepCrossAttention: Supercharging Transformer Residual Connections

We thank the reviewer for taking the time to review our manuscript. We address the two concerns raised by the reviewer below.

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.

Based 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):

This indicates that our results generalize to the vision domain.

The theoretical analysis, while sound, is less compelling than the experimental results; the authors could have expanded the empirical evaluations section instead.

In 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.

We are confident that these additional experiments and clarifications address the reviewer's concerns and further strengthen the paper.

We thank the reviewer for their time and insightful feedback on our manuscript. We address the three questions raised below.

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?

The method suggested by the reviewer has previously been proposed as LAuReL-LR. The 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.

We 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):

We 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.

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.

Since 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.

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?

In 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.

Moreover, 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.

We hope that these responses adequately address the reviewer's concerns. We believe these clarifications and additional results strengthen the manuscript.

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.

Could the authors provide analysis regarding the memory usage during forward and backward pass?

Since 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.

Let 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.

During 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.

Can the authors extend their comparison to include baselines with larger-scale models to assess the method's effectiveness at higher parameter counts?

Based 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.

The 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.