Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds

We thank the reviewer for their constructive comments. Their remarks will allow us to improve the clarity of our manuscript greatly.

Claims And Evidence

Clarifying the limit: By "deep," we mean constant ($O(1)$) depth. Our results hold in the large $D$, proportionally large $N$ limit.

Stationary points vs. minimizers: We agree and will refine the mapping between SE fixed points and the free energy around line 229 (see also response to Reviewer xpYy).

Methods And Evaluation Criteria

Order of learning: The reversed learning order is not specific to the algorithm. On our synthetic task, all algorithms learned the last layer first. In TREC, the reversal likely stems from task, model, and algorithm jointly. Our theory predicts that Bayes-optimal learning proceeds layer-wise from the last layer. While this matches some empirical observations, the learning order depends intricately on task and architecture and deserves further study. We will note this.

Experimental Designs Or Analyses

The diagonal structure of $Q$ is task-specific. Iterating SE from generic $Q$ shows it remains diagonal. Intuitively, the target weights are random and approximately orthogonal in high dimensions, implying the same for the GAMP estimates.

Fig. 1 (right): define weak recovery threshold. Defined in Theorem 2.3: the smallest sample complexity at which GAMP’s estimates have non-zero correlation with the target. We will add a pointer in the caption.

Fig. 2: nonzero start The first point shows overlap after one iteration—not initialization. We will clarify.

Supplementary Material

Thank you for the note, we will improve that appendix to be better self-contained.

Other Comments Or Suggestions

line 21 We indeed mean "proportionally large", we will reword.

Suggestions on eqns (2)-(3): Thank you, we will follow your suggestion for clarification. As we discuss in l-157, we consider frozen value matrices fixed to identity. We will stress this point.

In (3) $1_M$ denotes: In (3) $1_M$ refers to the identity matrix.

barer input and other typos: We thank the reviewer for their careful reading. These typos will be corrected.

Questions For Authors

1. c is the skip connection strength, introduced in (3).

2. They are indeed PxP matrices. We will stress this point.

3. They are Kroneker deltas: $\delta_{ab} = 1$ if $a=b$ and zero otherwise.

4. We mean that the prior on the weights $W^*$ is nonseparable. The denoiser $g_{out}$ is defined in (12) and its functional form can be arbitrarily complicated, both column-wise and row-wise non-separable, as in the case of deep attention models, as detailed in Appendix E.

5. We mean that it would be natural to expect we need less order parameters: in general one would need $PL \times PL$, but if the first $L$ layers are identical we just need $4P \times 4P$ by imposing that the overlap $Q$ has a certain structure. More in detail, we would be imposing that the overlap of a layer $l$ and the $L-$th layer is identical for any $l\le L-1$, that the overlap between any pair of layers $l_1,l_2\le L-1$ is identical, and finally that the self-overlap of the weights of any layer $l\le L-1$ is identical.

6. We refer to the main answer above to this same question.

7. The SE equations do involve the dimension $M$ through the denoiser $g_{\rm out}: \mathbb{R}^{M\times M} \to \mathbb{R}^{P\times M}$, involved in the SE equation (14), and defined in (12). A technical challenge was, in fact, precisely to ascertain the denoiser $g_{\rm out}$, which we unravel in Appendix E. Another challenge was to actually show the equivalence between SMIMs and MIMs (see appendix A).

8. $K$ denotes the output dimension—for sentiment classification, this corresponds to the predicted class probabilities. In our setting, the output is an $M \times M$ matrix, so $K = M^2$, interpretable as a target attention matrix capturing token interactions. Sentiment analysis, with $K$ being the number of sentiments, is a good example that we will include.

9. The setup in GB23 and BMN20 (Berthier, Montanari, Nguyen, 2020) for non-symmetric matrices allows for non-separable non-linearities along both the axis i.e across different $N$ samples and across different coordinates of the parameters. See for instance equation $(5)$ in GB23 and equations $(1),(2)$ in BMN20. In Equations (61)-(62), $B^{t}$ contains the flattened parameters $\hat{W}^t$ and therefore has dimensions $DM \times PM$. The number of samples $N$ appears in the dimensions of $\Omega^t \in \mathbb{R}^{N \times PM}$ which contains the corresponding flattened inputs.

10. Adding fully trainable non-linear layers is a challenging pursuit which would warrant completely new technical ideas, and we defer the study thereof to future work. Non-learnable non-linear layers could be added without a big additional technical challenge.

We thank the reviewer for their constructive comments. Numerically, we observe that in the setting of Fig. 2, GAMP and Gradient Descent display similar layer-wise learning dynamics. To put these observation on completely rigorous theoretical grounds, one would need to adapt the dynamical-mean-field-theory analysis of (Gerbelot, Troiani, Mignacco, Krzakala, Zdeborová 2024) to SMIMs. While we believe such an extension to be possible, it would warrant significant additional investigation, and we leave this for future work.

We thank the reviewer for their constructive comments. The focus of this work is indeed theoretical, and lies in making several theoretical findings on learning in attention models, extending previous related works. We discussed how the uncovered phenomena also hold with Gradient Descent (Figure 2 Bottom) and on some simple real world case (Figure 3 Right), but totally agree that it would be interesting to have a more thorough mechanistic study of layer-wise learning in more involved transformer models. On the other hand, we believe such an analysis is outside the scope of the paper at hand, and hope our work will inspire a more applied set of researchers to conduct such an investigation.

We thank the reviewer for their constructive comments.

Claims And Evidence:

Fig. 3 numerically supports the conjecture, while Sec. 3.2 discusses its infinite-depth limit where inner layers ($\ell \leq L - 1$) become interchangeable.

Relation To ... Literature:

The reviewer is right that our work builds on prior techniques. However, as noted, the new structure introduces substantial technical challenges.

Weakness

1. On optimality:

(1) GAMP is a polynomial-time algorithm that converges with few iterations, as predicted by state evolution. All examples in our plots used at most 50 iterations—we will state this explicitly.

(2) First-order optimality: GAMP minimizes prediction error among all first-order algorithms (e.g., Celentano et al. 2020), making our thresholds optimal in that class.

(3) Bayes-optimality: Our Th. 2.1 shows that when state evolution has a unique fixed point, it corresponds to the MMSE estimator. In such cases, GAMP achieves statistical optimality and our thresholds match the fundamental limits. This holds in all our examples.

2. Thank you—this helps clarify terminology. The weak recovery threshold for initial recovery is the sample complexity beyond which $\hat{W}$ escapes the zero fixed point, i.e., gains non-zero row-space. In contrast, the subspace recovery threshold for a subspace $U$ of the target row-space is the sample complexity above which $U$ is included in the row-space of $\hat{W}$. Def. (17), equivalent to Def. 2 in [1], formalizes this: $v^\top Q_t v > 0$ ensures $v \in U$ is in $\hat{W}W^\star$'s row space, where $Q_t = \hat{W}W^{\star \top}/d = \hat{W} \hat{W}^\top/d$.

For SMI targets, GAMP dynamics are more complex than for single-index: Learning is coupled across subspaces—progress on one can aid others. This manifests as GAMP passing near a series of saddle points. The thresholds in (17) reflect the minimal sample complexities required to escape these saddles, matching those in [1] (Def. 5), found via the stability of (18).

The threshold in (16) is the minimal complexity to recover any 1D subspace, thus a special case of the more general subspace recovery threshold

3. The weak recovery threshold for initial recovery marks the point where there exists a subspace $U$ that satisfies the condition (17). For example, in sec. 3.1, span$(w^\star_2)$ is recovered first, so its subspace recovery threshold (17) coincides with the weak recovery threshold for initial recovery (16). In contrast, the threshold for span$(w^\star_1)$ is higher, and depends on how much of $w^\star_2$ has been learned. It is not equal to (16): One can define a more precise threshold $\alpha_{w^\star_1}(q_2)$—the sample complexity for recovering span$(w^\star_1)$ given overlap $q_2$ with $w^\star_2$, assuming minimal $q_2$ in (17). Thus, (17) gives a fine-grained notion of threshold for a specific subspace, while (16) captures the threshold for recovering any subspace.

4. In (18), $Q_t$ denotes the overlap between $\hat{W}^t$ and $W^\star$ when the estimator uses side-information of the form $\lambda W^\star + \sqrt{1 - \lambda} \times \text{noise}$ (see l. 254, right column). Eq. (14) corresponds to the case $\lambda = 0$, with no side-information beyond the dataset ${x^\mu, y(x^\mu)}$. Thus, (14) and (18) describe distinct AMP algorithms, depending on the presence of side-information. Introducing side-information is technically required to define the thresholds (17) via the limit $\lambda \to 0$. For clarity, we will denote the overlap as $Q^t_\lambda$ in Theorem 2.3.

5. The thresholds in Fig. 1 indeed instantiate the theory of Sec. 2. We will revise $\alpha_1$ in (17) to $\alpha_U$. In subsection 3.1, $\alpha_1=\alpha_{{\rm span} (w^\star_2)}$, i.e. the sample complexity necessary to retrieve the subspace spanned by the second layer weights, is the weak recovery threshold from (16). The sample complexity $\alpha_2$ corresponds to the sample complexity required to further retrieve the full row-space ${\rm span} (w^\star_2, w^\star_1)$. This corresponds to (17) for $U= {\rm span} (w^\star_2, w^\star_1)$.

6. Writing a MIM as an SMIM naturally leads to studying targets where each input is a set of vectors—a common setup in sequence models. SMIMs can be seen as MIMs with a block-structured weight matrix (see Appendix B). While our proofs build on MIM results, the specific block structure we consider is new and well-suited to attention models. SMIMs are more natural for gradient-based training, where gradients are computed for all weights. The equivalent MIMs involve tied/zeroed weights, incompatible with standard gradient descent. This is similar to how CNNs are special cases of FCNs, but still distinct in practice due to gradient computation.

Questions For Authors

We expect GAMP to saturate the SQ bounds, as shown for MIMs in [1] (Lemmas 3.3 and 4.3). Weak recovery thresholds relate to generative exponents: 0 if threshold is zero, 1 if finite, $\geq 2$ if infinite.