Transformer Block Coupling and its Correlation with Generalization in LLMs
large language models;transformers;hidden representations
Yes
No ethics review needed.
See weaknesses section above.
- The concept of using Jacobians to study transformer block coupling provides a fresh perspective on understanding the internal mechanics of LLMs.

- The authors conduct an extensive empirical analysis across a variety of LLMs, lending credence to their hypothesis regarding the significance of transformer block coupling.

- The paper is well-organized, with clear definitions and logical flow, making it easy to follow the complex concepts and their implications.
This paper presents an insightful exploration into the internal dynamics of Large Language Models (LLMs). The authors propose a novel framework to evaluate the coupling of transformer blocks through their Jacobian matrices, offering a structured perspective on their inter-token and cross-layer relationships. The study concludes that this coupling is positively correlated with model generalization performance and appears to be a more significant factor than other hyperparameters such as parameter budget and model depth.
- The most concern lies in the strength of the generalization claims made by the authors. The paper presents empirical evidence suggesting a strong correlation between transformer block coupling and model performance on Open LLM LeaderBoard. However, it seems that some new models (e.g., LLaMA 3, Phi-2) which involve more pertaining steps (tokens) may have a more significant increase in performance than coupling. Revisiting the Figure 2, the coupling seems to emerge at certain training steps and then remain plain. So if the correlation only holds on a limited scope of pretraining steps?

- The model performance is evaluated on Open LLM LeaderBoard. However, for some reasoning tasks, there may be some emergent capabilities [1] that only larger models have. In such settings, is the correlation between coupling and model performance still holding?

[1] Chain-of-Thought Prompting Elicits Reasoning in Large Language Models
Yes
No ethics review needed.
1) what is the value of fig 6/9?
2) what are your computational costs?
3) anything surprisingly found in the jacobians?
- the paper offers a new perspective on by coupling training dynamics with model performance. 
- the method demonstrates somewhat robust as shown with the regression fit
- potentially high impact on model architecture/ training and diagnosis (if better understood)
the paper analyzes and correlates jacobian singular value decompositions cross tokens and cross layer-depths. as a result, they define a 'coupling metric', and show that high coupling apparently correlates to high benchmark values.
as such some intel on the inner processings of llm's is generated but the interpretation is widely left to the user.
the interpretation of this metric is vastly unclear. it is kinda clear that model training dynamics converge over training steps, in particular with common decaying lr-schedulers applied. so i wonder rather if one can interpret the score as a 'model training convergence rate' rather then performance-metric. falcon-7b (the entire model cluster > -.7 in fig1) demonstrates that high-coupling does not necessarily correlate with openllm-benchmarks afterall, tracing the activations/ gradients already shows such a convergence as well. afterall, on a converged model, activations only nuanced shift between layers. 

similarly i'm unsure how to interpret fig 6 etc.
computational costs of the methods are also not discussed (?).

overall i appreciate the solid work and do think there is merit in this methodology per se. however more focus should be made to practical implications of this rather abstract metric, and rigorously discussed.
Yes
No ethics review needed.
Figure 3: The LSS is defined as $\frac{L}{\lVert \tilde{x}^l_{i}-\tilde{x}^0_{i} \rVert}$ -- clearly a ratio of two positive numbers. How can it be that LSS values reported in Figure 3 are negative?. The authors themselves mention that "Note that $\text{LSS}\geq1$" (in Line 216) which is inconsistent with Fig. 3

Figures 4 and 5: Readability will be greatly improved if some indications of the relevant "axes" is added to the figure itself (Layer/Time in Fig 4/5; embedding dimension on one of the panels; ). Since the different models basically exhibit "identical" results (at least qualitatively), I will advise the authors to keep just one column (one model) for each Fig in the main text, and use the freed space to better visual explanation (other models can be moved to Supplementary).
Even more importantly: a scale/colorbar should be included. Currently it's not even clear if different sub-panels have the same scale or no. The choice to present absolute values is also a bit odd and not well-motivated.
The paper offers an intriguing observation, expanding on previous works that have reported similar phenomena in resnets. It makes a strong use of the large number of openly available LLMs to study hidden representations. The observation that the "coupling" performance correlates with model performance is of potential interest to the community and could encourage further research.
This paper studies the inference dynamics in LLMs, by looking at several properties of the "trajectories" that tokens representations form as they evolve through the network. It is found that the linearized approximations of different blocks in a trained transformer tend to align with each other (in the sense of their top singular vectors being aligned), and that the magnitude of this alignment correlates with overall model performance. Several other properties of the trajectories are also discussed.
The paper seems to be somewhat missing a more coherent organization and message. In particular:
- What is the motivation to consider the additional metrics of "Linearity" (Section 3.2) and "Layer-wise exponential growth" (Sec 3.3)? Can these be related, theoretically or empirically, to the "coupling" phenomena that is the main focus of the paper? (There's a small attempt to answer that in Section 6.2, I don't think it is very satisfying -- see below).
- The coupling phenomena is presented as an "interesting observation", but the paper doesn't offer much more than that, and so the reader is left somewhat unsatisfied. I think it would significantly improve the manuscript if the authors would study this phenomena in a more "functional" way -- is this a uniform behavior for all types of tasks/prompts, or are there some variations that we can understand? Is coupling "important" to either learning or inference, or is it an epiphenomena?
- Another possible way of making the findings more meaningful is a more mechanistic understanding of the phenomena (for example, does it happen in small models, trained on synthetic data? and so on; although this might be further away from the paper's overall approach).

Other than that, there are some "local" issues:

Line 35: The explanation for what a "discrete non-linear dynamical system" is, is pretty confusing. Almost every term seems to be at least inaccurate:
- "The term discrete reflects the network finite depth" (line 38): the term discrete is typically referring to the fact that *time* is discrete in these models, and that the dynamics are described in terms of finite differences rather than rate of change (as would be in a differential equation). It would still be "discrete" even if the network depth would be infinite (i.e., in discrete-time RNNs). The converse is also not true, a continuous-time dynamical system can be evaluated/evolved for finite time.
- "coupled refers to the interdependent token trajectories enabled by self-attention": this is inaccurate, the dynamics will be coupled even in a simple MLP (in fact even in a linear network -- consider the dynamics $\mathbf{x}^{(t+1)}=\mathbf{A}\mathbf{x}^{(t)}$ : for a non-diagonal $\mathbf{A}$, the dynamics is coupled because there are interactions between $x_i$ and $x_j$ through $\mathbf{A}$.)
- "dynamical is due to the residual connections spanning various layers": this is also an odd statement. the term "dynamics" is typically understood simply to refer to the fact that the state is changing with time. Consider the dynamical system from the previous bullet (as an example for dynamical system without the residual connections).
If the authors find it crucial to use non-standard interpretations of standard terms, they should at least note that, and hopefully provide a reasoning for their choice.



Figure 2 and section 5.3: The fact that the rise in "coupling" is so fast seems like a strong indication that the correlation between coupling and performance (as shown in Fig 1) is due to some other variables/confounders, and that coupling in itself has no "causal" importance. The reason is that, presumably, performance decreases gradually and doesn't plateau after 10-20K training steps, like coupling. It could still be the case that coupling early-on facilitates further training (even if not a necessary component of inference in a trained model). In that sense, the authors statement  "this finding suggests that the development of training methods to amplify coupling across transformer blocks may lead to favourable model performance" (Line 431) seems to be a strong claim, as the authors haven't demonstrated any evidence for a causal relationship between coupling and performance.


Section 6.2:
I find the entire explanation rather confusing.
The linearization of $F^l$ in and by itself doesn't mean one can approximate the entire dynamical system as a linear dynamical system. The system is linearized *at* the particular input $x$, so we can't read out much about the dynamics of the same input -- one has to introduce a perturbation in $x$ and *then* use the linearized dynamics to study the evolution of this perturbation (and while it's not clear how to perturb input tokens, it is easy to perturb their embeddings in the first layer).
I would also suggest to the authors to re-write their A.4 appendix in terms of a discrete dynamics (the same arguments can be made, but will be more relatable to the actual content/method of the paper)
Yes
No ethics review needed.
In summary, what can we concretely learn from the result? If we don't consider the quality, it's not very difficult to develop metrics that show a correlation but are ultimately not useful. Playing devil's advocate for a moment, for example, one of the most boring metrics could be the performance of a single downstream task like GSM8K. It may be correlated with the performance on the LLM Leaderboard. But, of course, this finding is not helpful. Compared with this or other candidates with such correlations, how can we claim that the Coupling metric is more interesting, helpful, or convincing?

The main claim may be based on Figure 1. But, the result looks very consolidated. Can we believe the presence of the correlation over various models or a limited number of models? For example, what happens if you only examine models that score above 50? Is there still a strong correlation, or does the relationship break down at higher performance levels? For example, if one observes phi-1 and phi-2, the Coupling rule may be broken because they may have similar settings, similar Couplings, and very different performances.

I'm not sure why linearity and exponential growth are discussed and considered as a goodness metric. Playing devil's advocate for a moment again, if we make Transformer blocks, which just scale their inputs by large positive constants, it seems to satisfy the linearity and exponential growth but would fail to achieve the high performance as they the Transformer is just a scaler. Could you provide more discussion and reasons for choosing or designing the metrics for analysis?

Figure 3 (c). Does this show that Linearity disappears at the end of training?

Figure 6. MPT -> Mistral?

Can we see Figure 6-style results of untrained models?

Figure 8. Can we see the results of untrained models?
To the best of my knowledge, this work is the first to explore the Coupling measures and indicate a correlation between them and the LLM performance on downstream tasks.
It also analyzes linearity in intermediate embeddings over the depth and exponential growth of the embedding norms.
The paper proposes Coupling, a metric of a Transformer model.
It computes the similarity of two Jacobians, followed by a normalization; each is the sum of non-residual terms by Attention and FFN.
The main claim of this paper is that the Couping metric of a model, among a set of public pre-trained LLMs, is correlated with the average performance of a set of downstream tasks from the HuggingFace Open LLM Leaderboard (R^2 value of 0.7214).
Additionally, the paper explores the properties of the Transformer models in terms of the line-shape score (LSS), defined by Gai & Zhang (2021), and the exponential spacing (expodistance) of the hidden trajectories.
The correlation between the Coupling metric and performance may not be a helpful or exciting finding. First of all, the paper does not provide any reasonable explanation or hypothesis about (1) why they investigated the Coupling, (2) why the Coupling could reflect performance, and (3) whether the correlation is from a causality (of a direction) or not.
Additionally, the paper does not perform intervention experiments, e.g., training a model with a regularization of increasing the Coupling and demonstrating the improvement. In total, it is difficult to conclude that this finding is helpful for the community so far. There are many experimental results, but they are not always well-organized and explained. I would like to know the real contributions of the results.

The correlation between the Coupling metric and performance might be very consistent. It may disappear in trends of recent high-performance models.

The contribution of the analysis of linearity and exponential growth is unclear. In my understanding, the analysis is conducted independently from the Coupling. Results are not specific to the models of interest and the Coupling. Many things might be obvious from Li & Papyan (2023). It would be great to add connections between them and Coupling (and the performance, if possible).
We discover coupling in transformer block Jacobians and show a relationship with generalization.
@inproceedings{
anonymous2024transformer,
title={Transformer Block Coupling and its Correlation with Generalization in {LLM}s},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=kvLenbZZgg},
note={under review}
}
Large Language Models (LLMs) have made significant strides in natural language
processing, and a precise understanding of the internal mechanisms driving their
success is essential. In this work, we trace the trajectories of individual tokens as they pass through transformer blocks, and linearize the system along these trajectories through their Jacobian matrices. By examining the relationships between these Jacobians, we uncover a **transformer block coupling** phenomenon in a variety of LLMs, characterized by the coupling of their top singular vectors across tokens and depth. Our findings reveal that coupling *positively correlates* with model performance, and that this relationship is stronger than with other hyperparameters, namely parameter budget, model depth, and embedding dimension. We further investigate the emergence of these properties through training, noting the development of coupling, as well as an increase in linearity and layer-wise exponential growth in the token trajectories. These collective insights provide a novel perspective on the interactions between token embeddings, and prompt further approaches to study training and generalization in LLMs.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
large language models transformers hidden representations
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/c24fb2e72aa0bcfad10ef13b5ab80b7797fdaafc.pdf
interpretability and explainable AI
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/59326fcc4babfabc8374a1c8f2bda0ef608a2d7b.zip
Transformer Block Coupling and its Correlation with Generalization in LLMs
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission