The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

Summary

The paper analyzes the covariance matrix of transformer networks at initialization, which commonly serves to analyze the training stability of deep networks. Specifically, the limiting distribution at network initialization is characterized by an SDE in the infinite depth-and-width limit regime. To enable the analysis, the authors make minor modifications to the attention and MLP formulation, which do not seem to impact their functionality much. Finally, promising preliminary numerical results are presented, which indicate that the transformer network with modified attention and MLPs can train stably and produce non-trivial performance without normalization layers.

Strengths

The paper is well written and relatively easy to follow. The analysis looks solid as far as I can tell. In particular, both the self-attention and FFN blocks, both with skip connections, are analyzed, and only relatively minor modifications are made to make the analysis tractable. The paper also presents promising experimental evidence that the proposed shaped attention can be trained stably without normalization layers.

Weaknesses

I appreciate the analysis on BERT-style language modeling in the appendix. It would be great to train the model with a comparable setup to BERT-Base and evaluate it on GLUE. That would then be a first piece of evidence that the proposed modified could be deployed in practice. A curious thing is that the shaped attention in Eq. (9) can have negative values and values larger than 1, and no longer consists of proper PMFs. While this might not be an issue in practice, it is still a deviation from the principled formulation of the attention being normalized to 1. Can the authors think of an alternative formulation for which the normalization proprerty would still hold (or at least for which negative attention values could be avoided)? Apart from that, I only see minor weaknesses. For example that the analysis is only for the network at initialization, but I agree with the authors that this can be addressed in follow-up work. Typos: Missing space in L25 (“train. One”), L122 (“important properties”)

Questions

* Could the proposed analysis be applied to decoder models with causally masked attention? Would masking the shaped attention lead to issues? * Could the dependence on the average token discussed in L264 be an artifact of the bias towards diagonal values in Eq. (9)? Are there potential issues with long-range dependencies when introducing this bias? * Assuming the authors used the embedding dim for the MLP hidden dim (both equal to 200), how would the results change when the depth, width and MLP hidden dim would be chosen with ratios like e.g. in BERT base? * Besides trainability without normalization layers, did the authors observe other signs of improved stability? For example fewer spikes in the loss curve or lower training loss variance?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There are no major limitations, please see “Weaknesses”. The authors are upfront that they only present an analysis for the network at initialization, and that more sophisticated analyses should be done in follow-up work. I’m not an expert in SDE-based analysis of neural networks and did not check the proofs in depth.

Summary

This paper studies a transformer architecture in the infinite width and depth limit where the ratio between the two remains constant. It is shown that, at initialization, the dynamics of the sample covariance matrix as it propagates through the layers converges to an SDE in the limit, where the SDE is indexed by the depth to width ratio. Moreover this limit is non-commutative and randomness preserving. The authors then propose the shaped attention - by constructing the attention matrix to resemble the identity function, and by properly scaling the residual connections and the softmax temperature, a more stable SDE is reached at the limit, theoretically allowing for more stable training.

Strengths

1. The paper presents a rigorous analysis of the infinite depth-width limit, showing convergence of the dynamics (propagation of the data covariance at initialization) to a stochastic differential equation. 2. The proposed shaped attention architecture is intriguing, meriting further investigation as to its potential benefits.

Weaknesses

My main issue with the paper is the lack of any convincing experimental result on the potential benefit of the proposed architecture. While it seems reasonable that a more stable forward process would lead to a more stable training and better performance, this is not always the case in practice, and indeed the paper barely makes any effort to confirm this. The experiment that is presented in section D of the appendix is in my opinion not convincing for the following reasons: 1) only 2 random seeds are considered for generating the results which is not close to enough, and 2) the relative improvement is inside the standard deviation anyway (for d=24). Without further experimental confirmation, i'm afraid the entire motivation and analysis is shaken to the point of being pointless from a practical perspective (I am willing to increase my score if such evidence is presented). Having said that, i do not think that an immediate practical application is mandatory for any paper, however given the abundance of literature on the subject of infinite width/depth limits i see the contribution of the paper at its current form as incremental, more fitting a workshop than the a main conference contribution.

Questions

Is there any reason to prefer an SDE limit over an ODE one (assuming a feature learning limit rather than the NTK)?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Limitations are adequately addressed

Summary

This paper investigates the convariance matrix of a modified softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. They develop a theory to describe the limiting distribution of the convariance matrix using a SDE indexed by the depth-to-width ratio. Based on this result, they propose a shaped attention to achieve a well-defined limit and ensure the stablity of the network.

Strengths

1 The authors derive neural convarince SDEs for ReNets and softmax-based attention, respectively. Then, they combine the attention and residual layers, and give the SDE for full transformer architeture. To the best of my knowledge, the theoretical contributions are novel. The analysis technique is standard and the proof is convincing. 2 The theoretical contribution is also useful to guide the hyper-parameter tuning. The idea of shaped attention is very interesting, and the numerical simulations demonstrate the effectiveness of the proposed method.

Weaknesses

1 All simulations are conducted at random initialization. It is better to conduct several experiments to show how the proposed initialization affect the training dynamics. 2 There are several typos, e.g., lines 25, ''.One'' -> ''. One'', line 122 , ''importantproperties"-> ''important properties", line 242 "remains -> remain".

Questions

Please see weaknesses.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations, and there is no potential negative societal impact of their work.

Summary

**Please note that this is an emergency review, completed in a limited amount of time whilst I am traveling.** This paper studies the properties of a class of attention-based networks, modified from the standard Transformer, in the limit in which depth and width tend proportionally to infinity. Technically, it extends recent work by Li, Nica, and Roy (NeurIPS 2022) on SDE-based descriptions of MLPs with random weights. The central result is a modification of the attention mechanism, referred to by the authors as shaped attention, that gives a non-trivial limit.

Strengths

In my abbreviated reading, I found the paper clear and enjoyable to read. I think it is a worthy contribution to the rapidly-growing literature on double-scaling limits of deep and wide networks, as it takes a substantial step closer to the architectures used in practice relative to prior works. The non-commutative scaling regime of ResNets studied by the authors nicely compliments recent work by Hayou and Yang on commutative infinite-width and -depth limits of residual networks. I found the observations regarding finite-time numerical instabilities to be particularly interesting. Due to time constraints, I have not had the chance to go through the proofs in detail.

Weaknesses

To my mind, the main weakness of this work is that it does not address inference. However, this reflects the larger issue that it is challenging to perform inference in conditionally-Gaussian processes (perhaps outside the nearly-Gaussian limit, where things can be studied perturbatively). Thus, this limitation does not significantly dampen my enthusiasm for the authors' work.

Questions

- Lines 46-47: In discussing the perturbative regime, work by Dyer and Gur Ari (2019) and by Zavatone-Veth, Canatar, Ruben, and Pehlevan (2021) should be cited alongside the work of Roberts, Yaida, and Hanin (2021). - Lines 48-53: For completeness, it might be useful to cite work on the proportional limit from random matrix theory, e.g. by G. Akemann and colleagues. - Line 243: "Covariane" -> "Covariance" - Lines 322-333: In relation to my comment about inference above, I think this paragraph would be substantially improved if the authors could offer a more concrete suggestion of how their approach could be leveraged to study training and generalization. - Lines 381-382: Why is the arXiv version of Li, Nica, and Roy cited in place of the NeurIPS 2022 publication?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

I think the authors adequately discuss the limitations of their work.

Replying to Rebuttal by Authors

The authors have addressed most of my concerns. Solid contribution.