Despite the widespread success of Transformers across various domains, their optimization guarantees in large-scale model settings are not well-understood. This paper rigorously analyzes the convergence properties of gradient flow in training Transformers with weight decay regularization. First, we construct the mean-field limit of large-scale Transformers, showing that as the model width and depth go to infinity, gradient flow converges to the Wasserstein gradient flow, which is represented by a partial differential equation. Then, we demonstrate that the gradient flow reaches a global minimum consistent with the PDE solution when the weight decay regularization parameter is sufficiently small. Our analysis is based on a series of novel mean-field techniques that adapt to Transformers. Compared with existing tools for deep networks (Lu et al., 2020) that demand homogeneity and global Lipschitz smoothness, we utilize a refined analysis assuming only $\textit{partial homogeneity}$ and $\textit{local Lipschitz smoothness}$. These new techniques may be of independent interest.
Paper
Similar papers
Peer review
Summary
The paper considers the theoretical mean-field limit of Transformers where width and depth go to infinity and studies the approximation error and convergence properties of gradient flow with weight decay. Both a residual self-attention layer and a residual feedforward layer are approximated by an ODE which averages the two encoders, whose solution models distribution of the parameters of both blocks throughout the depth of the Transformer. It is shown that the dynamics of the discretized model approximates the Wasserstein gradient flow under some regularity assumptions. Moreover, it is shown under partial homogeneity and universal kernel assumptions that if the Wasserstein flow weakly converges to a stationary distribution, then the discrete model must also converge with arbitrarily small risk.
Strengths
* The paper is a nontrivial extension of existing mean-field analyses of two-layer neural networks and ResNets to Transformers, which is challenging due to the existence of both feedforward and attention layers, the latter of which is new in the literature. * The assumptions for the networks and input data are quite general and can encompass e.g. the softmax mechanism and in-context learning settings. * The analyses and techniques are sufficiently novel and detailed and serve as an initial characterization of the dynamics of large scale Transformers, in particular the global convergence analysis of Section 4.
Weaknesses
* The obtained error bounds (Theorem 3.1) are so large that they likely have little practical relevance beyond the initial moments of training. In particular, the hidden constants are in the worst case super-super-exponential with respect to the time horizon. Specifically, the bound is exponential in $\phi_T(N,D,C\times B_\tau)$ where $B_\tau$ is exponential in $R_\tau$, which in turn is exponential in $\tau$. While bounds exponentially diverging in time are common in the literature (which is a priori expected from an ODE discretization argument) as mentioned in the paper, the dependency in Theorem 3.1 is much worse and some effort needs to be made to at least justify this. For example, one exponentation seems to be removable if e.g. $\phi_T$ is bounded. Currently the dependency in the hidden constants are not made obvious without going through the proofs in the Appendix, which I feel should be addressed more transparently. * The convergence analysis (Theorem 4.1) also suffers from the same problem. The analysis requires a time horizon $\tau_0$ large enough so that $W_2(\rho^{(\tau)}, \rho_\infty)$ is exponentially small (w.r.t. $R_\infty$) compared to the desired error $\epsilon$. This horizon is then fed into the approximation bound of Theorem 3.1, again yielding very large constants unless convergence is achieved in the very early stages of training. (This issue seems less critical since it only affects the $C_1$ term, although the $C_2\lambda$ term is still super-exponential in $R_\infty$.) * The recurring rate $L^{-1}+\sqrt{\log L/M}$ is not shown to be tight, leading me to further question the utility of the provided bounds. Why is the scaling for $L,M$ different? How does the rate compare with the analysis for ResNets or other ODE systems with 'width' & 'depth' dimensions? * A nitpick: the terms 'universal constant' or 'universal bound' are used quite frequently, however they should be reserved for constants that do not depend on *any* problem parameters.
Questions
* Can the dependency on time horizon be alleviated with stronger model assumptions? * Is the $L^{-1}+\sqrt{\log L/M}$ rate tight/expected? (see Weakness) * The widths of both the feedforward layer and the attention layer are both set equal to $M$, however the number of heads cannot typically be expected to be very large compared to the width of feedforward layers. Does the analysis generalize to when they are different? * Some recent papers [1,2] have also studied mean-field limits of Transformers from different perspectives, it would be nice to add a comparison in the related works section (although the latter is a contemporary work). [1] https://openreview.net/forum?id=xm2lU7tteQ [2] https://arxiv.org/abs/2405.15712
Rating
5
Confidence
3
Soundness
4
Presentation
2
Contribution
2
Limitations
See Weaknesses.
Summary
The authors theoretically investigate the mean field limit of residual transformer models. As depth and width go to infinity, thanks to the residual structure, the forward pass can be modeled as an ordinary differential equation, and the training gradient flow converges to a Wasserstein gradient flow. The author show well-posedness of this PDE, provide non-asymptotic time uniform bounds on quantities of interest and show global convergence to the global minimum of the empirical risk as $L,m/log(L) \to+\infty$ and $\lambda \to 0$.
Strengths
The tackled problem is extremely relevant and the derived results are sound. Given the size of the work, the authors did a good job overall in organizing the contents.
Weaknesses
1. The content is intense and the presentation doesn't help a novel reader. I believe polishing of the manuscript would improve the quality. More concretely, I would add a small paragraph in the main manuscript summarizing the main ideas behind the novelties. In the appendix instead it is really easy to get lost in a long proof, loosing track of the main goal. I suggest to include a small paragraph at the beginning of each section with the main steps required for the final proof, to guide the reader and make even the technical part more accessible (something on the line of what the authors did in D.3 would be perfect). 2. While this did not and will not impact my score, for manuscripts of this size I believe at least a really small experimental section (even in the appendix) to test the claims would make the results more sound.
Questions
On the technical side: 1. While I understand the need of assumption i) in theorem 4.1, it is not obvious to me how stringent it is. Already for an extremely simple problem like $Q(\rho) = \int \rho \log(\rho) + \lambda \int x^2 d\rho$ the minimizer is a Gaussian, that has unbounded support. Given the last example, the heuristic justification given by the authors seems to not be sufficient. Do the authors see a way to transfer this condition on the growth rate of $R(\Theta)$ at infinity? 2. The immediate extension of the deterministic case would be to consider the full Fokker-Plank, as it may emerge as a model of stochastic optimization. Already in the simplest case, the PDE under consideration would be $\partial_\tau \rho^{(\tau)} = div(\rho^{(\tau)} \nabla U) + \sigma \Delta \rho^{(\tau)}$. The stationary distribution in this case would be $\rho \propto e^{-U/\sigma}$, that has unbounded support for $U$ defined on all parameter space. This suggest that already in the simplest case (just pure isotropic diffusion added), the extension of the results may not be straightforward. I would appreciate the authors' comment on this, and I would like to know if they see an easy way to weaken that hypothesis as it's something one has no control on. 3. While mentioned in the proofs, it is not clear in the main manuscript what $\delta$ is. For example, in theorem 4.1 the statement is "with high probability", but it is not mentioned on what. I would add two words in the theorems/lemmas that are in high probability to specify with respect to what. 4. It sounds strange to me that theorem 3.1 gives a bound on a compact set, while theorem 4.1 on an unbonded time horizon, with what appears to be exactly the same constant. Is the constant in theorem 4.1 really the same of the one in 3.1 or it includes other terms?
Rating
7
Confidence
4
Soundness
4
Presentation
2
Contribution
4
Limitations
The authors correctly assessed the limitations in the manuscript, I have no further suggestions.
Summary
This paper analyzes gradient flow on Transformer networks. It is shown that for wide and deep Transformers, gradient flow converges to the Wasserstein gradient flow and reaches a global minimum.
Strengths
This is a well-written paper and it seems the results are strong and clean. However, I have very little background in this area and cannot verify the correctness of the statements. I hope the AC can find another qualified reviewer and delete this review if possible.
Weaknesses
N/A
Questions
N/A
Rating
6
Confidence
1
Soundness
3
Presentation
3
Contribution
3
Limitations
N/A
Summary
This paper studies transformers in their mean field limit. They use a ResNet architecture with infinite depth. The residual blocks are made of two steps: one transformer step and one standard MLP. They also use infinite width for both steps. Then, they show that the gradient flow in this limit is well posed, by using the notion of Wasserstein gradient flows since the mean field model is defined on space of measures. They also show consistency with the original discrete transformer model. Then, they study the global convergence properties of the gradient flow: they prove a sort of global convergence result, which requires some hypotheses.
Strengths
The writing of the paper is quite clear. Proving mean field convergence results has been studied a lot in the literature on theoretical deep learning. The result of the authors is interesting because they deal with the transformer architecture which ends up being a sum of two maps in the mean field representation. The authors can prove similar results in this context, which are new to the best of my knowledge of the literature.
Weaknesses
One of the main weakness for a submission to NeurIPS is the format of the paper. The 9 pages of the main paper are devoted to a (well done) explanation of the main results. There is no concrete ideas of proof in the main paper. More than 40 pages of technicalities are necessary to validate their result. The novelty of the results is the fact that transformer architecture is taken into account in the mean-field limit. However, the stated results do not show to what extent this architecture helps in obtaining the results. In fact, the transformer architecture appears more as something that hinders the standard proof on global convergence. The paper does not shed light on the interest/peculiarity of the transformer architecture. I would have appreciated if some simple experiments could be done to assess the range of validity of the results and in particular testing the importance of the transformer architecture. The hypothesis in section 4 to obtain a global convergence result seems not checkable in practice. In my opinion, it is not a global convergence result but rather an "if" theorem. Although I know that this assumption has been put forward in other previous papers, it is still an very demanding hypothesis.
Questions
What are specific features of the transformer architecture that are necessary to make the results valid? In other words, can you extend your results to more general architecture than transformers? The justification of the hypothesis on the separability property in C5 is difficult to understand. Can you elaborate more on this assumption and try to report precisely on the progress of the literature on this assumption? I do think it is not fair to claim global convergence for such kind of "if" theorem with a hypothesis that is uncheckable. I suggest the authors to rephrase their contributions.
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
2
Limitations
Not applicable.
Summary
The paper presents a rigorous theoretical analysis of the global convergence properties of gradient flow in training Transformers in the setting of in-context learning. By constructing the mean-field limit, the authors show that as the model's width and depth increase to infinity, the gradient flow converges to a Wasserstein gradient flow. This convergence ensures that the training process reaches a global minimum when the weight decay regularization parameter is sufficiently small. The paper introduces techniques adapted to Transformers, leveraging partial homogeneity and local Lipschitz smoothness to demonstrate the close approximation between discrete and continuous models and to establish the global convergence of gradient flow.
Strengths
Strength: 1. The paper studies the Transformer model in the context of in-context learning. Although the mean-field approach has been applied to infinite width and depth ResNet models before, the Transformer model is currently the most widely used and thus worthy of detailed study. This paper extends the analysis by incorporating two distinct encoders, filling the gap in understanding the Transformer model in infinite limits. 2. The paper introduces new techniques by extending previous approaches used for ResNet models. The authors refine the analysis by assuming only partial homogeneity and local Lipschitz smoothness, which is a key extension. The idea of considering the continuous dynamics as a function of the average behavior of f and h is new and interesting. 3. The final result has a mild requirement for the dependence of depth and width, where one would only need depth $L = \Theta(\epsilon^{-1})$ and the number of Transformer block $M = \Theta(\epsilon^{-2} \log (\epsilon^{-1}))$ to achieve a loss as small as $\epsilon$. (However, the reviewer has some questions about the result, so this strength can still be questionable.) 4. The paper is well-organized and clearly written. The authors effectively communicate the ideas, contributions, approaches, and limitations of their work, making it accessible and comprehensible to the reader.
Weaknesses
## Potential fallacy in the claim of Corollary 4.1: 1. The reviewer is convinced by most results presented in the paper, but has concerns about the claim in Corollary 4.1. The main issue arises from the fact that the authors treat many objects as constants, even though they clearly depend on some key quantities. In short, the reviewer thinks that $C_1$ in Theorem 4.1 depends significantly on $\epsilon$, which makes the claim in Corollary 4.1 questionable. The concern stems from the observation that $C$ in Theorem 3.1 may depend exponentially on $\tau$. Following the proof of Theorem 4.1, to achieve a loss as small as $\epsilon$, one would need $Q(\rho^{(\tau_0)}) \le \epsilon$. This introduces a non-trivial dependence of $\tau_0$ on $\epsilon$. Given that $C$ may depend exponentially on $\tau_0$, $C_1$ in Theorem 4.1 would thus depend exponentially on $\tau_0$, leading to a complex dependence on $\epsilon$. This dependence cannot be ignored, and treating $C_1$ as a simple constant to achieve Corollary 4.1 appears incorrect. (The reviewer believes this is how Corollary 4.1 is currently derived.) Based on the above, the reviewer thinks that 1. The current approach in proving Corollary 4.1 is flawed. 2. A naive application of Theorem 4.1 might only lead to asymptotic bounds for $L$ and $M$, instead of the linear and quadratic terms presented nicely in the current version of Corollary 4.1. ## Limitations of settings and results Though the reviewer agrees that studying and achieving a good understanding of the case where both width and depth go to infinity is valuable, there are some limitations in the results presented in the paper. 1. The main result, Theorem 4.1, depends on Assumptions 1-4. Although the reviewer considers Assumptions 1-4 to be reasonable and likely to hold under certain settings, the authors do not make it clear when these assumptions would hold. Consequently, the setting of the Transformer model might differ from classical Transformer models (e.g., it might only hold under the reparameterized case (2.3) instead of keeping all value, key, query, and output matrices). The reviewer will address this point more in the Questions section. 2. The main result, Theorem 4.1, relies heavily on a few additional assumptions that are not justified in the paper: 1. the Wasserstein gradient flow weakly converges to some distribution $\rho_{\infty}$ 2. the uniform boundedness of $\rho^{(\tau)}$ for large time $\tau$ 3. the separation property for $\alpha_1$ with the support expansion of $\alpha_2$ to $K$. The reviewer acknowledges that some of these assumptions are also adopted in other mean-field approaches and appreciates the authors' efforts to justify them intuitively in Section C.5. However, they are not examined carefully and may harm the validity of the statement. For example, for the second assumption, the reviewer agrees that a large regularization $\lambda$ can implicitly bound the norm of the solution. However, this might require a large $\lambda$, which could make the results less meaningful. Notice that the current bound shown in Proposition 3.2 implies that $R_{\tau}$ may depend exponentially on $\tau$, thus it is not clear that a small $\lambda$ can result in a bounded solution. And based on Theorem 4.1, a small $\lambda$ would be required to achieve a small loss. 3. The result does not provide a non-asymptotic bound for the required training time. In other words, there is no guarantee that $\hat{Q}(\Theta^{(\tau)})$ will be small for any specific $\tau$. This is because there is no control over the convergence rate of $\rho^{(\tau)}$ to $\rho^{\infty}$, which would require a stronger assumption beyond the current weak convergence assumption of $\rho^{(\tau)}$. 4. The loss and training dynamics are only considered for the population risk under $\mu$. There is no consideration of the effect of an empirical dataset. This is a clear limitation, but the reviewer can accept this as an initial step toward studying the Transformer model. ## Minor: 1. Some parts of the statements are not well organized. For example, Assumption 2 is an assumption about norm bounds of function $f$ and its gradient. The line "define ReLU'(x) = ..." has no direct connection with the assumption. The reviewer thinks that the authors are trying to argue that by using ReLU'(x) as the gradient of the ReLU function, $f$ defined in (2.3) would satisfy Assumption 2. If this is the case, the reviewer suggests the authors separate the statement into 1) what the assumption of $f$ is. and 2) when the assumption would hold. Another minor point is that the Fréchet derivative still acts differently from the gradient, so some careful treatment might be needed if the authors claim that the result in Theorem 4.1 applies to Transformers with ReLU activation. Typos: 1. Line 111, it should be "h(Z, w) = W_2 \sigma_M(W_1 Z)" instead of $W_1 H$. 2. Line 156, in Assumption 2, $\Psi_j$ in the description should be $\Phi_P$. 3. Line 86: "~~the~~ our deep Transformer model" 4. Line 901, For the line that bounds $R(\Theta)$, there is no $C_{\lambda}$ in the first line, and the transition from the first line to the second line is currently shown as $1/2 \epsilon + ... \le 1/4 \epsilon + ...$
Questions
1. Please address the above concern for Corollary 4.1 carefully. The answer may greatly affect the final rating from the reviewer. 2. Questions related to assumption 1-4: 1. For the purpose of confirmation, have you shown in which cases Assumption 1-4 are met for the Transformer model? My current understanding is that they would be met when you reparameterize $f$ as in (2.3) but not when you keep the W_V, W_K, W_Q and W_O matrices (at least Assumption 2 won't be satisfied in this case). Is that correct? The review is not against the reparameterization and considers this is still a Transformer model, but want to understand the difference and the claim of the paper more carefully. 2. As the focus of the paper is on the Transformer model, the main results are based on the structure defined in (2.1) (2.2), (2.4) and Assumptions 1-4 over f and h. The review wonders if it is possible to provide a separate corollary for a commonly used form that satisfy Assumption 1-4, such as (2.3), making it clearer on an example where minimum assumptions are used. Or in other words, are there technical difficulties in showing Assumption 1-4 for (2.3)? Is the universal approximation capabilities the key obstacle?
Rating
5
Confidence
4
Soundness
2
Presentation
3
Contribution
3
Limitations
The authors have adequately addressed the limitation of the paper.
Given the above points, the reviewer considers the paper to be an interesting attempt at tackling important questions. The paper introduces valuable techniques for addressing problems related to Transformer networks and has the potential to become a strong contribution. However, in its current form, the paper is incomplete and would benefit greatly from revision and resubmission to a future venue. The reviewer suggests the following improvements: 1. Please rewrite Corollary 4.1 and other theorems, ensuring that the dependencies for all constants are clearly specified. 2. Prove that assumptions 1-4 are satisfied for Transformer networks in the form of (2.3). It would be helpful if some parameters in the assumptions, including $K, K_T, K_P, \Phi_T, \Phi_P, \Phi_{PP}, \Phi_{TP}, \Phi_{TT}$ could be represented in more explicit forms with clear bounds. 3. If the authors intend to include ReLU activation beyond smooth activations, they should demonstrate that ReLU activation also satisfies Assumptions 1-4 and Proposition 3.1. 4. For the assumptions used in Theorem 4.1, even though it may be difficult to prove when these assumptions hold, it would be beneficial to provide heuristic examples to help readers understand when and why the assumptions are valid. For instance, the example provided by Reviewer accV in question 1 (Q3 in the rebuttal) offers a good illustration that this bounded assumption may not generally hold. Therefore, the authors need to make a stronger effort to justify the assumptions. In the current response to that question, phrases like "remains possible" are not convincing enough. Again, the reviewer understands that theoretical work in this area is challenging and that compromises are often necessary. For example, the weak convergence assumption is difficult to avoid in the mean-field regime, and previous work on ResNet models [Ding et al. 2021] has also shared the same asymptotic nature as in this paper. This is why the reviewer initially leaned toward accepting the paper, hoping that some of these concerns could be addressed. However, as outlined above, the reviewer does consider the work to be incomplete and believes that the paper does not meet the standard for a conference paper due to its lack of direct application to Transformer models, its weak (asymptotic) results, and its reliance on assumptions that are difficult to verify. The reviewer apologizes for the change in rating, noting that this adjustment reflects a reassessment of the initial rating rather than a direct response to the rebuttal. [Ding, Zhiyan, et al. "Overparameterization of deep ResNet: zero loss and mean-field analysis." Journal of Machine Learning Research 23.48 (2022): 1-65.]
We appreciate your careful feedback. We understand the concerns you have raised regarding Corollary 4.1, the asymptotic nature of our results, and the applicability to Transformer models. However, we would like to argue that these weaknesses you have mentioned are essentially the weaknesses of almost all existing mean-field studies of deep neural networks. In fact, our paper contributes to pushing the common assumptions in the mean-field literature towards more practical scenarios. Therefore, we respectfully disagree with your comment that our work is incomplete. Corollary 4.1 should not need major revision, as its mathematical formulation and technical proof are accurate. To ensure rigor, it suffices to delete the sentence in line 302 "Here, $L_0$ scales as $\Omega(\epsilon^{-1})$, and $K_0$ as $\Omega((1+\delta)\epsilon^{-2})$", and remove the paragraph around lines 304-306. Regarding your concern that the dependence on training time $\tau$, the depth $L$ and the width $M$ is asymptotic, we would like to emphasize again that our result remains, to our best knowledge, the state-of-the-art global convergence guarantees for large-scale Transformers. Moreover, we would also like to point out that our work has two main results: (i) connections between the discrete Transformers and their continuous limit, and (ii) global convergence guarantees of Transformers. We feel that it is unfair to deny all of our contributions based solely on that our global convergence guarantees for discrete Transformers require asymptotically large $\tau$, $L$ and $M$. Regarding your concern about the applicability to Transformers, we honestly believe that the Transformer architecture we consider is already among the ones considered in theoretical works that are closest to practical Transformer architectures. We will add concrete examples of Transformer architectures and rigorously prove that Assumptions 1-4 hold for these architectures in the camera-ready version. As clarified, these assumptions easily hold for activation functions that are sufficiently smooth. We will also clarify the case of the ReLU activation function, noting that it may not generally satisfy these assumptions, but it is possible for Assumptions 1-4 to hold when the data follow certain favorable distributions. Once again, we appreciate your detailed and comprehensive feedback. We sincerely hope that you can consider re-evaluating our contributions. Thank you for your time and attention. Best regards, Authors
The reviewer appreciates the authors' further response and acknowledges the effort made in addressing the concerns raised. Here are the reviewer's further comments: 1. The reviewer agrees that the term "major revision" used in the previous feedback regarding Corollary 4.1 was not accurate. For the sake of correctness, removing the lines related to the dependencies of $L$ and $M$ should be sufficient. When the reviewer suggested a revision, it was more about improving the clarity of the statement. For instance, as also noted by reviewer mtEE, the authors have treated many quantities as constants even though they still have dependencies on the problem parameters. The reviewer believes that this may have contributed to the inability to identify the incorrect conclusion in the initial paper. Another minor revision could involve removing the $C_{\lambda}$ constraint and instead simply choosing $\lambda \le (2C_2)^{-1}\epsilon$, and then adjusting the proofs on line 901 with the correct constants. 2. The reviewer also acknowledges that many current works in the mean-field regime exhibit asymptotic behavior. The reviewer recognizes that this work represents state-of-the-art results in this challenging area. However, the asymptotic nature of the results, particularly regarding the width and depth, still limits their practical applicability. 3. As noted in the previous response, the concern regarding the concrete application to the Transformer architecture remains the key issue that leads to a reject score. The reviewer appreciates the authors' willingness to include concrete examples of Transformer architectures and to address the case of ReLU activation. The reviewer agrees that these additions seem feasible, and if incorporated, the reviewer would be willing to raise the score and accept the paper. However, while these additions make sense intuitively and the authors seem confident in their ability to implement them, the reviewer remains skeptical that these theorems can be adequately developed and integrated into the camera-ready version, given the current form of the paper and the volume of work required. Nonetheless, including these components would make the work more convincing regarding the convergence story for Transformer networks.
Thank you very much for acknowledging that the asymptotic assumptions on $L$ and $M$ are common in the mean-field analysis literature. We also appreciate your suggestion on removing $C\_{\lambda}.$ We will follow your suggestion, choose $\lambda \le (2C\_2)^{-1}\epsilon,$ and adjust the proofs accordingly. Regarding your concern about the concrete application to the Transformer architecture, we would like to clarify that verifying Assumptions 2-4 (please note that Assumption 1 is irrelevant to the Transformer model, and is only a fairly mild assumption on the data) for concrete examples of Transformer architectures with smooth activations is fairly intuitive and the proof is mainly based on a series of tedious calculations. Below, we give a concrete proposition and its brief proof as follows. --- Consider $f(Z,\theta)=VZ \mathrm{softmax} (Z^T W Z ) $ $\qquad $ (Eq 1) with the collection of parameters $\theta = \mathrm{vec}[V,W],$ where $\mathrm{softmax}$ denotes the column-wise softmax function. Moreover, consider $h(Z,w)=W\_2\mathrm{HuberizedReLU}(W\_1H)$ $\qquad $ (Eq 2) with the collection of parameters $w= \mathrm{vec}[W\_1,W\_2].$ Here, $\mathrm{HuberizedReLU}$ denotes the entry-wise HuberizedReLU activation function defined as $\mathrm{HuberizedReLU}(x) = \left\\{ \begin{aligned} &0, &&\mathrm{if} \ z\leq 0;\\\\ &z^2/2, &&\mathrm{if} \ z\in [0,1];\\\\ &z - 1 / 2, &&\mathrm{if} \ z\geq 1. \end{aligned} \right. $ Then, we can consider a Transformer model defined by equations (2.1), (2.2), and (2.4) in the paper, where the functions $f$ and $h$ are specified above. We suppose that this Transformer model is applied to a learning task with data that satisfies Assumption 1. We have the following proposition. **Proposition.** Consider the Transformer model defined by equations (2.1), (2.2) and (2.4), with $f(Z,\theta)$ and $h(Z,w)$ defined in (Eq 1) and (Eq 2) respectively. Then Assumptions 2-4 all hold. In our following comments, we will present a proof sketch of the proposition above. Due to the large character count of equations, we split the proof into several separate replies. We apologize for the long comments. To simplify our response, we omit the detailed derivations for the function $h(Z,w)$, which corresponds to the MLP part, in our verification of Assumptions 2 and 3. We feel that the fact that $h(Z,w)$ satisfies Assumptions 2 and 3 is relatively more intuitive, especially given the proofs for $f(Z,\theta)$. We will also omit some calculation details to keep our response reasonably simple. We will make sure that all the omitted details will be added to the revised version of our paper.
Proof. We first introduce some notations. We remind the readers that for a matrix $\mathbf{A}$, we denote by $\\| \mathbf{A} \\|\_2$ the spectral norm of $\mathbf{A}.$ We also denote by $\\| \cdot \\|\_{1-\mathrm{col}}$ the maximum $\ell\_1$-norm across all columns of a matrix. Denote $Z=(z\_1,\dots,z\_{N+1})\in\mathbb{R}^{D\times N+1}.$ Then the function $f$ can be rewritten as $f(Z,\theta)=VZ\mathrm{softmax}(Z^TWZ)=(f(Z,\theta)\_{:,i})\_{1\leq i\leq N+1}$, where $f(Z,\theta)\_{:,i}=\sum\_{j=1}^{N+1}P\_{ij}z\_j$ and $P\_{i,:}=\mathrm{softmax}(Z^TWz\_i).$ Next, we calculate the derivatives of $f(Z,\theta)\_{:,i}$ with respect to $Z$ and $\theta$ as follows: For $Z$: the Jacobian $J\in\mathbb{R}^{(N+1)D\times(N+1)D}$ is $J=(J\_{ij})\_{1\leq i,j\leq N + 1},$ where $J\_{ij}=\frac{\partial f\_{:,i}}{x\_j}\in\mathbb{R}^{D\times D}.$ After calculation, we obtain $J\_{ij}=ZQ\_i\big[E\_{ji}Z^TW+Z^TW^T\delta\_{ij}\big]+P\_{ij}I,$ where $Q\_i:=\mathrm{diag}(P\_{i:})-P^T\_{i:}P\_{i:},$ $E\_{ij}$ is the matrix with zeros everywhere except one the $(i,j)$-th entry, and $\delta\_{ij}$ is the Kronecker delta. For $\theta$: Define $A\_i=Z^TWz\_i.$ After calculation, we have $\nabla\_{\mathrm{vec}[V]} f(Z,\theta)\_{:,i}=\sum\_{j=1}^{N+1}P\_{ij}\mathrm{diag}\Big([B\_{kj}(z\_j)]\_{1\leq k\leq D}\Big),$ $\nabla\_{\mathrm{vec}[W]} f(Z,\theta)\_{:,i}=\sum\_{j=1}^{N+1}P\_{ij} \Big(\mathrm{diag}\Big([z_l^T B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big)[\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}]\_l\Big)_{1\leq l \leq N+1} Vz\_j,$ where $B\_{kj}(z)$ is the $D\times D$ matrix with zeros everywhere except $z\_j$ for the $k$-th row. --- We then verify the assumptions one by one. For Assumption 2 (i), we have $ \\| f(T,\theta) \\|\_{2-\mathrm{col}} = \\| VT \mathrm{softmax} (T^T W T ) \\|\_{2-\mathrm{col}} \leq \\| V\\|\_2 \cdot \\| T \\|\_2 \cdot \\| \mathrm{softmax} (T^T W T ) \\|\_{2-\mathrm{col}} \leq \\| \theta \\|\_2 \cdot \\| T \\|\_{2-\mathrm{col}} \cdot \\| \mathrm{softmax} (T^T W T ) \\|\_{1-\mathrm{col}} \leq \\| \theta \\|\_2 \cdot \\| T \\|\_{2-\mathrm{col}},$ where the second-to-the-last inequality follows by the fact that $\ell\_2$-norm can be upper bounded by the $\ell\_1$-norm, and the last inequality follows by the fact that each column of the softmax output has an $\ell\_1$-norm equaling one. Therefore, the first condition in Assumption 2 with $K = 1$ is verified for the function $f$ in (Eq 1). For $h$ in (Eq 2), we have $ \\| h(T,w) \\|\_{2-\mathrm{col}} = W\_2\mathrm{HuberizedReLU}(W\_1T) \\|\_{2-\mathrm{col}} \leq \\| W\_2 \\|\_2 \cdot \\|\mathrm{HuberizedReLU}(W\_1T) \\|\_{2-\mathrm{col}} \leq \\| W\_2 \\|\_2 \cdot \\|W\_1T \\|\_{2-\mathrm{col}} \leq 2 \cdot \\| w \\|\_2^2 \cdot \\|T \\|\_{2-\mathrm{col}}$, where the second inequality follows by the property of HuberizedReLU that $|\mathrm{HuberizedReLU}(x)| \leq |x|.$ This demonstrates that Assumption 2 (i) with $K=1$ holds for $h$ in (Eq 2) as well. For Assumption2 (ii), we have $\\| \nabla\_{\mathrm{vec}[V]} f(T,\theta)\_{:,i} \\|\_{2} \leq \sum\_{j=1}\^{N+1} P\_{ij} \\| \mathrm{diag}\Big([B\_{kj}(T\_j)]\_{1\leq k\leq D}\Big) \\|\_{2} \leq \sum\_{j=1}^{N+1} P\_{ij} \\| T\_j \\|\_{2} \leq \\|T\\|\_{2-\mathrm{col}}. $ Similarly, we have $\\| \nabla\_{\mathrm{vec}[W]} f(T,\theta)\_{:,i} \\|\_{2} \leq \sum\_{j=1}^{N+1} P\_{ij} \Big\\| \Big(\mathrm{diag}\Big([z_l^T B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big)[\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}]\_l\Big)_{1\leq l \leq N+1} VT\_j \Big\\|\_{2} \leq \sum\_{j=1}^{N+1} P\_{ij} \sum\_{l=1}^{N+1}\\|T\_l^T T\_j \\|\_{2}\cdot (\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i})\_l \cdot \\|V\\|\_{2} \cdot \\| T\_j \\|\_{2} \leq \\|T\\|\_{2-\mathrm{col}}^3 \\|\theta\\|\_{2}.$ Combining the two equations above gives Assumption2 (ii) with $\phi\_P( \\|T\\|\_{2-\mathrm{col}} ) = \\|T\\|\_{2-\mathrm{col}} + \\|T\\|\_{2-\mathrm{col}}^3$. For Assumption2 (iii), we have $\\| J \\|\_{2}\leq \sqrt{\sum\_{1\leq i,j\leq N+1}\\|J_{ij}\\|^2\_2}\leq (N +1) \max\_{1\leq i,j\leq N+1} \\|J_{ij}\\|\_2.$ For any $1\leq i,j\leq N + 1,$ we have $\\|J_{ij}\\|\_2\leq P_{ij} + \\|T\\|\_2 \\|Q\_i\\|\_2 \\|E\_{ji}T^TW+T^TW^T\delta\_{ij}\\|\_2 \leq 1 + 2 \\|T\\|\_2^2 \\|W\\|\_2 \leq 1 + 2 \\|T\\|\_F^2 \\|\theta\\|\_2. $ Hence, we have $\\| J \\|\_{2}\leq 2N \\|T\\|\_F^2\cdot (1 + \\|\theta\\|\_2).$ The above equation demonstrates that for $f$, Assumption2 (iii) holds with $\phi_T ( N, D, \\| T \\|\_F ) = 2N \\|T\\|\_F^2$. We have verified Assumption 2 for the attention layer encoder $f$. The verification for $h$ is similar and easier. We omit the derivation for $h$ here to shorten our response, but we will make sure to include the full verification in the revised version of the paper.
Next, we verify Assumption 3. Given that we are currently considering the example where the encoder employs a smooth univariate activation function, we can prove stronger results by removing the expectation $\mathbb{E}_{\mu}$. (i) and (iii): Given the calculation of derivatives we have presented above, it suffices to show that $\mathrm{diag}\Big([B\_{kj}(z\_j)]\_{1\leq k\leq D}\Big)$ and $\mathrm{diag}\Big([Z^T B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big)\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}Vz\_j$ are both locally Lipschitz continuous with respect to $Z$ and $\theta.$ Since each of $\mathrm{diag}\Big([B\_{kj}(z\_j)]\_{1\leq k\leq D}\Big),$ $\mathrm{diag}\Big([Z^T B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big),$ $\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}$ and $Vz\_j$ is obviously bounded by an increasing function of $N,D,\\|\theta\\|, K_T,L_T,$ it suffices to show that each of them is locally Lipschitz continuity with respect to both $Z$ and $\theta.$ This is straightforward as they are all sufficiently smooth. (ii) and (iv): Because the norm of the difference of two Jacobian matrices $\\|J^1-J^2\\|\_2$ is bounded by $\sqrt{\sum\_{1\leq i,j\leq N+1} \\|J^1\_{ij}-J^2\_{ij}\\|\_2^2},$ it suffices to show that $J\_{ij}$ is locally Lipschitz continuous with respect to both $\theta$ and $Z.$ Again each component of $J\_{ij}$ that depends on $Z$ or $\theta,$ i.e. $Z,Q_i,W,P\_{ij},$ is bounded by an increasing function of $N,D, K_P,L_T,K_T,\\|\theta\\|,$ and is locally Lipschitz continuous given sufficient smoothness. Hence, (ii) and (iv) also hold. The proof for the HuberizedReLU MLP encoder is similar to the above and is not conceptually complicated, if not easier. We omit the details here, and save the space for a more detailed discussion about Assumption 4, as we expect that you may be more skeptical about the proof regarding the “universal kernel” assumption in Assumption 4. --- For Assumption 4, we consider the pair $(g,\alpha) = (h,w),$ and the partition $\alpha = (\alpha\_1,\alpha\_2)$ with $\alpha\_1 = W\_2,$ $\alpha\_2 = W\_1$. We also let a compact set $\mathcal{K} = \\{ W_1: \\| W_1 \\| \leq 1 \\}$. Then Assumption 4(i) on the partial $1$-homogeneity property straightforwardly holds: $h(T,W\_1,c\cdot W\_2) = c\cdot W\_2\mathrm{HuberizedReLU}(W\_1H)= c\cdot h(T,W\_1,W\_2).$ Regarding Assumption 4(ii) on the universal kernel property, we first note that according to the choice $(g,\alpha) = (h,w)$, this assumption is purely an assumption on the MLP part of the Transformer. Here we give the detailed proof as follows. First of all, according to the classic universal approximation theory (see the wiki page of “universal approximation theorem” and [1,2,3] for more details), we know that two-layer fully-connected networks with non-polynomial activation functions and without any constraints on its parameters are universal approximators. Therefore, we know that the function class $ \mathrm{span} \\{ W\_2\mathrm{ReLU}^2(W\_1T): W_1 \in \mathbb{R}^{\mathrm{dim}(W_1)}, W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)} \\} $ is dense in $\mathcal{C}(\\|T\\|\_{2-\mathrm{col}}\leq B,\mathbb{R}^{D\times(N+1)})$. Moreover, by the definition of HyberizedReLU, for any $B>0$ and any $\hat{W}_1$, $\hat{W}_2$, there exist small constant $c$ such that $c\cdot \hat{W}_1 \in \mathcal{K}$, $c\cdot \\|\hat{W}_1\\| \leq B^{-1}$, and $ c^{-2} \cdot \hat{W}\_2\mathrm{HyberizedReLU}(c\cdot \hat{W}\_1T) = c^{-2} \cdot \hat{W}\_2\mathrm{ReLU}^2(c\cdot \hat{W}\_1T) = c^2\cdot c^{-2} \cdot \hat{W}\_2\mathrm{ReLU}^2(\hat{W}\_1T) = \hat{W}\_2\mathrm{ReLU}^2(\hat{W}\_1T)$, where the second equation follows by the positive $2$-homogeneity of $\mathrm{ReLU}^2$ activation. This implies that $ \\{ W\_2\mathrm{ReLU}^2(W\_1T): W_1 \in \mathbb{R}^{\mathrm{dim}(W_1)}, W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)} \\} \subseteq \\{W\_2\mathrm{HyberizedReLU}(W\_1T): W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)}\times \mathcal{K}\\} $. Therefore, we conclude that $ \mathrm{span} \\{W\_2\mathrm{HyberizedReLU}(W\_1T): W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)}\times \mathcal{K}\\}$ is dense in $\mathcal{C}(\\|T\\|\_{2-\mathrm{col}}\leq B,\mathbb{R}^{D\times(N+1)})$. This finishes the validation of Assumption 4. --- [1] Funahashi, Ken-Ichi (January 1989). "On the approximate realization of continuous mappings by neural networks". Neural Networks. [2] Cybenko, G. (1989). "Approximation by superpositions of a sigmoidal function". Mathematics of Control, Signals, and Systems. [3] Pinkus, Allan (January 1999). "Approximation theory of the MLP model in neural networks". Acta Numerica.
We hope that by presenting the details that you have requested, we can convince you that our theory can be applied to fairly practical Transformer models, and that the volume of work required should not be an issue. We are dedicated to add additional details on the explanations and verifications in our revised paper. We sincerely hope that you could take our discussion above into consideration and reevaluate our result. We appreciate your time and effort in reviewing our paper.
The reviewer thanks the authors for their further efforts. While the reviewer is not able to examine the full proof within the rebuttal period, the reviewer will take a closer look into it in the coming days. Based on the provided proof, the reviewer believes that it can be incorporated into the revised paper and is therefore willing to increase the score, provided no major concerns arise upon closer inspection. As a side note, the reviewer finds it interesting to explore how the parameters specific to the Transformer architecture ($\phi_P, \phi_T$ and similar quantities in Assumptions 1-4) influence the convergence properties in the final theorem. Once again, the reviewer appreciates the authors' efforts during the rebuttal period to address the concerns raised.
Thank you for the detailed reply. I will maintain my score as I still feel the technical issue of Q1 ends up weakening the relevance of the paper, however I personally do think the approach and ideas are very interesting.
First of all, I would like to thank the authors for their thorough rebuttal. I am satisfied with the answers and I would like to keep my score, conditionally on the authors addressing Q1 in the revised version.
The reviewer first thanks the authors for their detailed and thoughtful response. The reviewer also appreciates the authors' honesty and transparency in the rebuttal. Unfortunately, the reviewer does not feel that the concerns have been adequately addressed. As a result, the reviewer has decided to change the rating. The reviewer understands how frustrating this can be, especially given the thorough rebuttal provided by the authors in an effort to address all the questions raised. The reviewer struggled with the rating during the initial review phase and provided the rating with the hope that some of these concerns would be resolved. However, after careful and extensive consideration, the reviewer believes that the paper remains incomplete in its current form. Before delving into specific concerns, the reviewer would like to acknowledge the strengths of the paper, as noted in the initial review. The paper is well-organized and clearly written, and it tackles a challenging and important problem: the theoretical understanding of training Transformer models. The reviewer recognizes that theoretical work, particularly in complex areas like Transformer networks, is inherently difficult. The reviewer is open to compromises, such as accepting reasonable assumptions that may be difficult to verify or acknowledging results with suboptimal dependence on problem parameters, including asymptotic results. However, based on the authors' response, the reviewer believes the following deficiencies are significant: 1. Corollary 4.1 is incorrect and requires major revision. 2. The dependence on training time $\tau$, the depth $L$ and the width $M$ is purely asymptotic. 3. The paper lacks clear results for the Transformer network (e.g., in the form of (2.3)) or for functions with ReLU activation. First of all, having a misleading or incorrect conclusion for one of the main claims of the paper is a serious issue, as it undermines the credibility of the rest of the work. This problem also leads to the asymptotic nature of the results, which significantly diminishes the importance of the paper’s contributions from the reviewer’s perspective. To clarify for those who may be interested, the reviewer believes that the logic behind Theorem 4.1 and Corollary 4.1 is as follows: A specific training time is required to ensure that the limiting dynamics reach a small error $Q(\rho^{(\tau_0)})$, and $L$ and $M$ must be chosen large enough so that the discrete training dynamics remain close to the limiting dynamics within the time $\tau_0$. The lack of control over $\tau_0$ to achieve a small loss leads to the asymptotic nature of the time (with additional weak convergence assumption), depth, and width. These first two deficiencies alone prevent the reviewer from giving the paper a high score. However, what is most concerning upon reevaluating the paper is the third point. The paper claims to provide convergence results for large-scale Transformers, yet no specific Transformer model is present in any of the main theorems. A Transformer should be defined in the form of (2.4) or at least in the form of (2.3), rather than as a concatenation of general functions $f$ and $h$. In the rebuttal, the authors used many "believe"s including the "belief" that (2.3) would satisfy Assumptions 1-4 and the "belief" that the results apply to the ReLU function with its Fréchet derivative. While the reviewer agrees with these beliefs on an intuitive level, the absence of a theorem that includes at least the form in (2.3) or the ReLU function makes these "beliefs" shaky and the work incomplete. It is somewhat surprising that the authors claim to achieve results for Transformers, yet provide no specific results related to the attention mechanism (2.4) (or its reparameterized form (2.3)) and the ReLU function. This presentation is therefore also misleading, making it difficult to connect the main results in the paper with a Transformer network. Given the rebuttal, it seems that Assumptions 1-4 for (2.3) and ReLU activation for Proposition 3.1 are checkable, so including these components would form a more complete story for the convergence of Transformers in the asymptotic regime.
I thank the authors for their answers. About Q2/A2: I do not think that an experiment on the MNIST dataset and just claiming global convergence would be a valuable addition to illustrate the claims in the paper. Studying the size of parameters to obtain this global convergence would be more interesting. In any case, having read the answers to Q3 and Q4, I reckon this work appears as an improvement over the current literature and could motivate some further progress. This motivates me to increase my score.
Thank you for your suggestions on the experiments and for increasing your score! We will make sure to add experiments to study the size of parameters that can guarantee to global convergence.
Decision
Accept (poster)