LAuReL: Learned Augmented Residual Layer

Thank you for your comments and suggestions.

Ablation study with practical footprint metrics

Thank you for suggesting this experiment. This section demonstrates improvements of LAuReL variants individually, as well as when they are combined.

For the purpose of comparing different LAuReL variants on the LLM pre-training task, we set up the following baseline. We pre-trained on the C4 corpus with $\approx$ 10B tokens. We used a $4 \times 4$ Google Cloud TPU v6e topology, but we expect similar results with a comparable GPU setup.

In order to simplify the comparison across many ablations (and also to avoid the noise in downstream evals at 10B token scale), we report model performance using the test loss, which is a good proxy for downstream model quality.

We train our main baseline with 24 layers and 157.2M params, along with a larger baseline with 28 layers for comparison. We run all the LAuReL variants (RW, LR, PA) on top of the regular baseline with 24 layers. We also run two combinations of variants (RW+LR, RW+LR+PA). We report the number of params, test loss, peak memory as reported by profiling tools, and average step time. Lower is better for all metrics. The table below shows the results; note that all LAuReL experiments use L=24.

All LAuReL variants perform better than the large baseline in terms of the test loss while using much fewer parameters, lower peak memory, and lower average step time. Given the above tradeoffs in loss, memory, step time, etc. we recommend trying the LAuReL variants in this order: RW $\rightarrow$ LR $\rightarrow$ PA / RW+LR $\rightarrow$ RW+LR+PA. We will add these numbers in the revision.

Can α and β be learned by the network

That’s a good question. Since $f(.)$ is non-linear, absorbing $\alpha$ within the function is not equivalent to having an explicit $\alpha$ outside $f(x)$. Additionally, the normalization would apply on top of the weighted combination of $f(x)$ and $x$, so the individual scalars would still be useful in learning the relative weights of the two components. We also demonstrate that the -RW variant is useful on its own in the above ablation study.

Relevant Work

Thank you for the references. DenseNet (Huang et al., 2018) connects every pair of layers in the network and hence in the vanilla version, all the activations need to be in memory. This is prohibitively expensive for deep LLMs and other modern transformers. When introducing dense-blocks, all previous activations within the block need to be visible to any given layer within the block; this requires refactoring the model architecture into dense blocks.

On the other hand, LAuReL requires minimal changes. In fact, in LAuReL-PA, which is the most similar to DenseNet, we make three design choices to achieve memory efficiency and performance. Firstly, each layer only looks at the past $k$ activations. For the above experiments, $k=3$ was sufficient. Secondly, we also propose using low-rank linear functions to further reduce memory usage due to activations. Thirdly, the LAuReL-PA variant uses learned scalars ($\gamma_{i}$, $\gamma_{i-1}$, …) to learn the weights of the previous activations (which we found to be crucial), whereas DenseNet assumes a simple sum of the previous activations.

Additionally, as seen above, LAuReL-RW and -LR variants provide significant improvements over naive-scaling, and can be combined with the -PA method.

Identity Mappings in Deep Residual Networks (He et al., 2016): This paper introduces variants of residual connections with different types of 'gating', which look similar the -RW variant, except that they use a much larger number of parameters for either the exclusive gating or the 1x1 conv gating ($D_in \times D_out$ params per layer), which is much more expensive than the 1/2 params per layer of the -RW variant. For the rest of the paper, the authors introduce and focus on the pre-activation residual connection which places the activation functions such that it helps with optimization.

We will include these discussions in the revision.

We hope these address your concerns - thanks!

Thank you for your kind comments and valuable feedback!

Justification into why specific variants perform better under certain conditions

Deep networks generally do better with reasoning, math, coding etc. tasks. However, residual connections are crucial for such networks, and LAuReL helps augment these residual connections with learned components. So at a high-level we expect LAuReL to help with reasoning/math/coding tasks.

Specifically we have the following intuition for why the different variants work:

• LAuReL-RW: Helps with learning the importance of the residual input ($x_i$). This might not be as useful in the earlier layers, but would be important in other later layers to tackle the vanishing gradient problem using a higher learned weight for the residual input.

• LAuReL-LR: Helps with allocating learning capacity for the linear part of the network (the residual input, i.e., $x_i$ in Figure 2), such that the main network can use it’s capacity towards learning better nonlinear functions ($f(x)$), while LAuReL contributes the linear components ($x_i + ABx_i$) to the residual stream.

• LAuReL-PA: A hybrid of -RW and -LR variants, where multiple previous activations are used in a learned weighted manner, along with learning linear functions on top of them. This allows layers accelerated access to previous activations, along with learning their relative importance.

Overall, LAuReL provides a general formulation for the residual connection, along with practical variants that can operate on the residual stream in different ways as highlighted above. LAuReL variants can be combined with each other, as seen in the ResNet experiments, as well as the small-scale ablations on an LLM reported in response to Reviewer gnJr.

Interestingly, LAuReL variants consistently show better performance than naively adding a layer on top of the baseline (ResNet experiments in Table 1, small-scale LLM ablations in response to Reviewer gnJr, and LLM-2 naive scaling experiments in response to Reviewer jdcS). This demonstrates that operating on the residual stream using LAuReL has a non-trivial impact on model convergence, which cannot be matched by adding a full additional layer whereas LAuReL variants add $\leq$ 0.1% params in these experiments.

Practical considerations for very deep networks (e.g., scaling to 100B+ parameter models)

While it is hard to experiment with 100B+ params, we can extrapolate from the ablations done in the paper as well as in response to Reviewer gnJr and Reviewer jdcS.

• The -RW variant is straightforward to try and does not have a tangible impact on latency and memory. We expect initialization of these scalars to also be important. For earlier layers, the initial weight given to the non-linear component can be higher.
• The -LR variant is cheap enough to also be included. We expect a slight increase in memory and latency footprint. $r$ should be scaled up as $D$ grows. In our experiments up to 4B params, the ideal ratio of $D / r$ was between 24--32.
• The -PA variant is also helpful if the network is very deep. Using $k={3, 4, …}$ was helpful in smaller-scale LLM experiments. When keeping a larger value of $k$, the accelerator memory usage should be monitored and appropriate rematerialization strategies should be employed. Although we do not see a large increase in memory usage with the -PA variant in the smaller-scale experiments. If the above variants work well with some headroom available in latency and memory, a combination of RW+LR or RW+LR+PA can be tried.

We hope these address your concerns - thanks!

Thank you for your comments and suggestions.

Tuning $r$

Since is $r$ the rank of $A$, $B^T$, we expect $r \ll D$. Indeed, if $D = 512, 768, 1024, …$, this leaves a small range of discrete values for $r$ (unlike hyperparameters such as learning rate, weight decay, that can take continuous values). In our experience $r = {32, 48, 64}$ work well for LLMs.

LLM Depth / Parameter Scaling with Error Bars:

Similar to the comparison with naive depth-scaling with ResNet (Sec 3.1), we ran scaling experiments with LLM-2. Originally, both baseline/LAuReL had 40 layers. Adding the 41st layer required turning off a minor architectural change (which enforced the number of layers to be divisible by 2). To be fair, we re-ran the baseline and LAuReL.

We present the number of params and average training step times.

Note LAuReL adds only 0.1% parameters and incurs a step time penalty of 2.42%. On latency, it is slightly above Baseline$^{+1}$ . This is because LAuReL is invoked for each layer, but the per-layer invocation cost is small.

The table below shows the downstream quality of the baselines and LAuReL.

LAuReL wins on all tasks except WMT23; it is more parameter-efficient than naive-scaling.

When to use Laurel-PA?

LAuReL-PA works well for deep networks where there is a risk of vanishing gradients. A suitable rematerialization and sharding, given the extra activations in memory, would help; the new ablation results (see Reviewer gnJr) are positive. We expect LAuReL-PA to do better than works such as DenseNet (Huang et al. '16), where the number of activations in-memory is quadratic.

Comparison with Parameter Sharing

Techniques like param-sharing are complementary, since they can be applied in conjunction with LAuReL.

Comparison with Highway Nets/Residual Gates

Highway Nets (Srivastava et al., 2015) is similar to LAuReL-RW, except that the Transform Gate requires ($D^2$ (weight matrix) + $D$ (bias)) params, in addition to the latency incurred by a full-rank matmult. But we use 1-2 scalars per Laurel-RW layer, with no significant latency impact. Similarly, the Residual Gates (Savarese et al., 2017) also is similar to LAuReL-RW, except they use ReLU as the gating function. However, LAuReL is more general including variants like -LR, -PA, which can be combined.

Green Percentages in Table 2 / 3

Apologies. Green and bold font indicated statistically significant improvements for ResNet and LLM-2. The percentages for a couple of tasks (TyDiQA, MGSM, etc.) had a typo, but the absolute values were correct.

Intuition for asymmetrical improvement in tasks

Deep networks generally do better with reasoning/math/coding. However, residual connections are crucial and LAuReL helps augment them with learned components. So we expect LAuReL to improve such tasks.

However, for certain other tasks the network might be bottlenecked on the number of parameters in the MLP layers. Thus it is hard to pinpoint why a particular task improves more than others with LAuReL.

We hope these address your concerns and you can reconsider the assessment/score - thanks!

Thank you for your comments and suggestions.

Ablation study with practical footprint metrics

Thank you for suggesting this experiment. For the purpose of comparing different LAuReL variants on the LLM pre-training task, we set up the following baseline. We pre-trained on the C4 corpus with $\approx$ 10B tokens. We used a $4 \times 4$ Google Cloud TPU v6e topology, but we expect similar results with a comparable GPU setup.

In order to simplify the comparison across many ablations (and also to avoid the noise in downstream evals at 10B token scale), we report model performance using the test loss, which is a good proxy for downstream model quality.

We train our main baseline with 24 layers and 157.2M params, along with a larger baseline with 28 layers for comparison. We run all the LAuReL variants (RW, LR, PA) on top of the regular baseline with 24 layers. We also run two combinations of variants (RW+LR, RW+LR+PA). We report the number of params, test loss, peak memory as reported by profiling tools, and average step time. Lower is better for all metrics. The table below shows the results; note that all LAuReL experiments use L=24.

All LAuReL variants perform better than the large baseline in terms of the test loss while using much fewer parameters, lower peak memory, and lower average step time. Given the above tradeoffs in loss, memory, step time, etc. we recommend trying the LAuReL variants in this order: RW $\rightarrow$ LR $\rightarrow$ PA / RW+LR $\rightarrow$ RW+LR+PA. We will add these numbers in the revision.

We found that r={32, 48} work well for LAuReL-LR, and k={3, 4} works well for LAuReL-PA. We went with r=32, and k=3 respectively.

Relationship to DenseNet

Thank you for the reference. DenseNet connects every pair of layers in the network and hence in the vanilla version, all the activations need to be in memory. This is prohibitively expensive for deep LLMs and other modern transformers. When introducing dense-blocks, all previous activations within the block need to be visible to any given layer within the block; this requires refactoring the model architecture into dense blocks.

On the other hand, LAuReL requires minimal changes. In fact, in LAuReL-PA, which is the most similar to DenseNet, we make three design choices to achieve memory efficiency and performance. Firstly, each layer only looks at the past $k$ activations. For the above experiments, $k=3$ was sufficient. Secondly, we also propose using low-rank linear functions to further reduce memory usage due to activations. Thirdly, the LAuReL-PA variant uses learned scalars ($\gamma_{i}$, $\gamma_{i-1}$, …) to learn the weights of the previous activations (which we found to be crucial), whereas DenseNet assumes a simple sum of the previous activations.

Additionally, as seen above, LAuReL-RW and -LR variants provide significant improvements over naive-scaling, and can be combined with the -PA method.

We will include these discussions in the revision.

ViT

Unfortunately these experiments are still in-progress at the time of the rebuttal deadline; we will add them as soon as the runs finish. However, we expect LAuReL to provide improvements on top of ViT baselines, given LAuReL has shown improvements on ResNet as well as three LLM baselines (LLM-1, LLM-2, and the new small LLM baseline mentioned above). Note that the latter three are transformer-networks, very much like ViT.

We hope we have addressed your concerns and questions and we hope you can reconsider your assessment/score. Thank you.