Transformer Block Coupling and its Correlation with Generalization in LLMs

## Abstract

Abstract 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.

## 1 Introduction

In recent years, many openly available Large Language Models (LLMs) have been released, achieving state-of-the-art results on task-specific benchmarks. The abundance of models, each with differing architecture and training methodology, motivates comparing the underlying mechanisms that drive generalization.

Transformers (Vaswani et al., ) can be represented as discrete, nonlinear, coupled dynamical systems, operating in high dimensions (Greff et al., ; Papyan et al., ; Haber & Ruthotto, ; Ee, ; Ebski et al., ; Chen et al., ; Bai et al., ; Rothauge et al., ; Gai & Zhang, ; Li & Papyan, ). The term discrete reflects the network’s finite depth; nonlinear refers to the model’s nonlinear components; coupled refers to the interdependent token trajectories enabled by self-attention, and dynamical is due to the residual connections spanning various layers.

Throughout this work, we investigate whether we can identify structural characteristics that consistently occur in pretrained LLMs that are related to their generalization, as measured on the HuggingFace OpenLLM Leaderboard benchmark. To this end, we study and quantify the trajectories of token embeddings and their interactions, through linearizations of LLM layers given by the Jacobians of the associated transformer blocks.

This perspective raises several questions:

Q1. What regularity properties do these trajectories exhibit individually, and what are their relations with one another? More concretely, what is the relation between the Jacobians across different tokens and transformer layers?

Q2. How do the properties of hidden representations and their relations emerge with training?

Q3. Are any of these properties related to the generalization capabilities of LLMs?

Figure: Figure 1: Relationship between Coupling and Benchmark Score. The plot illustrates the relationship between average coupling (negative mis-coupling) and benchmark score across various base LLMs (that are not fine-tuned). The data shows a noticeable trend: increased alignment correlates with improved benchmark scores. The regression fit produces an $R^{2}$ value of $0.7214$.
Refer to caption: extracted/5922317/figs/alignment/benchmark_vs_coupling.png

### 1.1 Contributions

We investigate the motivating questions across several openly-available LLMs, most having over 1B parameters, trained by 6 independent organizations, with varying training methods and data (Appendix A.1). Through our experiments, we identify consistent patterns that define the transformer block coupling phenomenon.

1. Coupling. The singular vectors of the Jacobians of transformer blocks couple across depth (Figures LABEL:fig:paris_all_8to16, LABEL:fig:paris_all_8to16_appendix) and tokens (Figures LABEL:fig:coupling_inoutsame, LABEL:fig:coupling_fixedin, LABEL:fig:coupling_fixedout, LABEL:fig:coupling_inoutsame_appendix, LABEL:fig:coupling_fixed_input_appendix, LABEL:fig:coupling_fixed_output_appendix). Further, coupling across Jacobians emerges with training (Figures 2, LABEL:fig:paris_all_8to16, LABEL:fig:coupling_inoutsame, LABEL:fig:coupling_fixedin, LABEL:fig:coupling_fixedout).

2. Generalization. The strength of coupling is correlated with benchmark performance on the HuggingFace Open LLM Leaderboard (Beeching et al., ) (Figures 1, LABEL:fig:alignment_vary_k). Additionally, coupling is more strongly correlated with generalization than parameter budget, model depth, and token embedding dimension (Figure LABEL:fig:correlation_hyperparams).

3. Regularity. Linearity in hidden trajectories emerges with training (Figures LABEL:figs:lss-training, LABEL:figs:pca, 13, 15, 23, 24). Exponential growth occurs in contiguous token representations as a function of depth (Figures 14, 22, 23, LABEL:figs:expo-training).

We provide a new perspective on token embedding interactions within LLMs by examining transformers layers through their Jacobians. Our results display the effect of training on transformer blocks, and suggest potential approaches for promoting generalization in LLMs.

Figure: Figure 2: Emergence of Singular Vector Coupling with Training. Normalized coupling (Section 3.1) evaluated for $J_{i,j}^{d_{1}},J_{i,j}^{d_{2}}$, with $|d_{1}-d_{2}|\leq 5$. Measurements are taken for checkpoints of Pythia 12B, Pythia 6.9B, Pythia 1B (Biderman et al., ) at training steps $\{1,16,128,13\text{k},23\text{k},33\text{k},\ldots,143\text{k}\}$. Coupling is low at early training iterations, and increases with training, with Pythia 12B and 6.9B displaying greater coupling than Pythia 1B in final training steps.
Refer to caption: extracted/5922317/figs/alignment/_decoupling.png

## 2 Background on Large Language Models

We describe LLMs as a deep composition of functions that iteratively transform token embeddings. In the input layer, $l=0$, textual prompts undergo tokenization and are combined with positional encodings to create an initial high-dimensional embedding, denoted by $x_{i}^{0}\in\mathbb{R}^{d_{\text{model}}}$ for the $i^{\text{th}}$ token. When these embeddings are stacked, they form a matrix:

$$ $X^{0}=(x_{1}^{0},x_{2}^{0},\ldots,x_{n}^{0})\in\mathbb{R}^{n\times d_{\text{ model}}}.$ (1) $$

The embeddings then pass through $L$ transformer blocks:

$$ $X^{0}\xrightarrow{F_{\text{block}}^{1}}X^{1}\xrightarrow{F_{\text{block}}^{2}} \cdots X^{L-1}\xrightarrow{F_{\text{block}}^{L}}X^{L}.$ (2) $$

$X^{l}=F_{\text{block}}^{l}(X^{l-1})$ denotes the embeddings after the $l^{\text{th}}$ block, consisting of causal multi-headed attention (MHA), a feed-forward network (FFN), and normalization layers (LN) with residual connections:

$$ $\displaystyle h^{l+1}(X^{l})$ $\displaystyle=\text{MHA}(\text{LN}(X^{l}))$ (3) $\displaystyle g^{l+1}(X^{l})$ $\displaystyle=\text{LN}(X^{l}+h^{l+1}(X^{l}))$ (4) $\displaystyle f^{l+1}(X^{l})$ $\displaystyle=h^{l+1}(X^{l})+\text{FFN}(g^{l+1}(X^{l}))$ (5) $\displaystyle F_{\text{block}}^{l+1}(X^{l})$ $\displaystyle=X^{l}+f^{l+1}(X^{l}),$ (6) where the MHA, LN, FFN are implicitly indexed by layer. Among many models (Appendix 1 ), an additional rotary positional embedding (RoPE, Su et al. ( 2023 ) ) is applied in the MHA layer. In the final representation, typically an additional layer normalization is applied: $\displaystyle F_{\text{block}}^{L}(X^{L-1})$ $\displaystyle=\text{LN}(X^{L-1}+h^{L}(X^{L-1})+\text{FFN}(g^{L}(X^{L-1}))).$ (7) $$

The output $X^{L}$ from the final block $F^{L}$ is passed into a bias-free linear layer $M\in\mathbb{R}^{d_{\text{vocab}}\times d_{\text{model}}}$, with $d_{\text{vocab}}$ denoting the size of the token vocabulary and $d_{\text{model}}$ is the dimension of the token embeddings. This layer $M$ computes final-layer logits for each token embedding, $\ell_{i}=Mx_{i}^{L}$. The prediction for the next token is then determined by selecting the maximal logit value:
$\arg\max_{v\in\text{tokens}}\ell_{v,n}.$

## 3 Methods

### 3.1 Coupling of Singular Vectors of Jacobians

Jacobians. Coupling is investigated through analyzing the linearizations of transformer blocks given by their Jacobians

$$ $J_{t_{1}t_{2}}^{l}=\frac{\partial}{\partial x_{t_{1}}^{l-1}}\left(f^{l}(X^{l-1 })\right)_{t_{2}},$ (8) $$

defined for each layer $l\in\{1,\ldots,L\}$, and pair of tokens $t_{1},t_{2}\in\{1,\ldots,n\}$. Note that this is the Jacobian matrix for each transformer block without the contribution from the skip connection from the input of the block, similar to the quantity measured by
Li & Papyan (), but extended for cross-token interactions. Also, due to the causal structure of the representations, $J_{t_{1}t_{2}}^{l}=0$ if $t_{1}>t_{2}$, so we restrict our attention to the case $t_{1}\leq t_{2}$.

Singular value decomposition. The SVD of the Jacobians $J_{t_{1}t_{2}}^{l}$ are computed, i.e., $J_{t_{1}t_{2}}^{l}=U_{l}S_{l}V_{l}^{\top}$ (with superscripts $t_{1},t_{2}$, indicating the tokens, omitted for clarity), where $U_{l}\in\mathbb{R}^{d_{\text{model}}\times d_{\text{model}}}$ and $V_{l}\in\mathbb{R}^{d_{\text{model}}\times d_{\text{model}}}$ are matrices of left and right singular vectors respectively, and $S_{l}\in\mathbb{R}^{d_{\text{model}}\times d_{\text{model}}}$ is the singular value matrix.

To measure coupling between Jacobians $J_{t_{1}t_{2}}^{l}$ and $J_{t_{1}^{\prime}t_{2}^{\prime}}^{l^{\prime}}$, define $A_{t_{1}t_{2}t_{1}^{\prime}t_{2}^{\prime}}^{ll^{\prime},K}\coloneq U_{l^{
\prime},K}^{\top}J_{t_{1}t_{2}}^{l}V_{l^{\prime},K}$ for $l,l^{\prime}\in\{1,\ldots,L\}$ and $t_{1},t_{2},t_{1}^{\prime},t_{2}^{\prime}\in\{1,\ldots,n\}$, where $U_{l^{\prime},K}$ and $V_{l^{\prime},K}$ are the sub-matrices with columns that are the top-$K$ left and right singular vectors of $J_{t_{1}^{\prime}t_{2}^{\prime}}^{l^{\prime}}$, respectively.

Figure: Figure 3: Transformer Block Coupling. Displayed is a visualization of the various types of transformer block coupling, along with brief instructions on computing both the Jacobians $J$ and coupling matrices $A$.
Refer to caption: x1.png

Depth-wise Coupling. To analyze coupling across transformer blocks, we fix $t$ (usually as the last token), and check alignment between $J_{tt}^{l}$ and $J_{tt}^{l^{\prime}}$ through $A_{t}^{ll^{\prime},K}$ (we write the single subscript $t$ for clarity) across layers $l,l^{\prime}\in\{1,\ldots,L\}$.

Token-wise Coupling. To study the coupling across tokens, we start by fixing two layers $l,l^{\prime}\in\{1,\ldots,L\}$, and analyze coupling across different cases of the $t_{1},t_{2},t_{1}^{\prime},t_{2}^{\prime}$. We start with the case where the input and output tokens are the same and compare $J_{tt}^{l}$ and $J_{t^{\prime}t^{\prime}}^{l^{\prime}}$ across $t,t^{\prime}\in\{1,\ldots,n\}$, which is the coupling across tokens for a token’s effect on its own trajectory. We also consider the context tokens’ impact on a trajectory by including coupling across cases with $t_{1}<t_{2}$. More specifically, we measure coupling between $J_{t_{1}t_{2}}^{l}$ and $J_{t_{1}t_{2}^{\prime}}^{l^{\prime}}$ across $t_{2},t_{2}^{\prime}\geq t_{1}$ (fixing the input token to be the same) and also between $J_{t_{1}t_{2}}^{l}$ and $J_{t_{1}^{\prime}t_{2}}^{l^{\prime}}$ across $t_{1},t_{1}^{\prime}\leq t_{2}$ (fixing the output token to be the same).

Coupling measurement. To quantify the coupling of singular vectors of the Jacobians, we measure how far the $A_{t_{1}t_{2}t_{1}^{\prime}t_{2}^{\prime}}^{ll^{\prime},K}$ are from being diagonal. More precisely, we quantify the mis-coupling of $A_{t_{1}t_{2}t_{1}^{\prime}t_{2}^{\prime}}^{ll^{\prime},K}$ using $m:\mathbb{R}^{K\times K}\to\mathbb{R}$ defined by

$$ $m(A)=\lVert A-\text{Diag}(A)\rVert_{F},$ (9) $$

where $\text{Diag}(A)$ is the matrix $A$ with all non-diagonal entries replaced by zero and $\lVert\cdot\rVert_{F}$ denotes the Frobenius norm. For normalized comparison between models across training in Figure 2, we normalize $A$ by their norms

$$ $\tilde{m}(A)=\frac{\lVert A-\text{Diag}(A)\rVert_{F}}{\lVert A\rVert_{F}}$ (10) $$

### 3.2 Linearity of Trajectories

Linearity in intermediate embeddings is quantified with the line-shape score (LSS), defined by Gai & Zhang () as

$$ $\text{LSS}_{i}^{0,\ldots,L}=\frac{L}{\left|\left|\tilde{x}^{L}_{i}-\tilde{x}^{ 0}_{i}\right|\right|_{2}},$ (11) $$

where $\tilde{x}_{i}^{0}=x_{i}^{0}$, i.e., the input embeddings passed to the LLM, and $\tilde{x}_{i}^{l}$ is defined recursively as

$$ $\tilde{x}^{l}_{i}=\tilde{x}^{l-1}_{i}+\frac{x^{l}_{i}-x^{l-1}_{i}}{\left|\left |x^{l}_{i}-x^{l-1}_{i}\right|\right|_{2}}\ \ \mbox{for}\ \ l=1,\ldots,L.$ (12) $$

Note that $\text{LSS}\geq 1$, with $\text{LSS}=1$ if and only if the intermediate representations $x_{i}^{0},\ldots,x_{i}^{L}$ form a co-linear trajectory.

### 3.3 Layer-Wise Exponential Growth

We measure the presence of exponential spacing (*expodistance*) of the hidden trajectories. Assuming exponential growth of the embedding norms as they flow through the hidden layers, we can estimate,

$$ $\displaystyle\|x_{i}^{l}\|$ $\displaystyle\approx e^{\alpha l}\|x_{i}^{0}\|=e^{\alpha}\|x_{i}^{l-1}\|$ (13) $$

for some fixed $\alpha\in\mathbb{R}$ over all layers $l=1,\ldots L$. We quantify the validity of the representation in Equation 13 by measuring the coefficient of variation of $\alpha_{i}^{l}$, given by

$$ $\alpha_{i}^{l}\approx\ln\left(\frac{\|x_{i}^{l}\|}{\|x_{i}^{l-1}\|}\right),$ (14) $$

for each layer $l$ and token $i$. Under exponential growth, it is expected that $\alpha_{i}^{l}$ is independent of depth. We therefore denote the expodistance (ED) of the trajectory of the $i^{th}$ token of a given sequence by

$$ $\text{ED}_{i}=\frac{\text{Var}_{l}\alpha_{i}^{l}}{(\text{Avg}_{l}\alpha_{i}^{l })^{2}}.$ (15) $$

This measurement is motivated by the parametrization discussed in Appendix A.4, as well as empirical evidence in Figure LABEL:figs:norm-expo-traininga, and serves as a method to test the validity of the linearization presented in Equation 17.

## 4 Evaluation

### 4.1 Suite of Large Language Models

Our study evaluates a total of 38 LLMs (24 base LLMs and 14 fine-tuned, see Appendix A.1) that were independently trained by various individuals and organizations. These models, provided through HuggingFace (Wolf et al., ), vary in terms of parameter budgets, number of layers, hidden dimensions, and training tokens. Moreover, we analyze the dynamics of each measurement throughout training by deploying the Pythia Scaling Suite (Biderman et al., ). A summary of the models under consideration is presented in Table 1 of Appendix A.1 and further details in Appendix A.5.

### 4.2 Prompt Data

We evaluate these LLMs using prompts of varying length, ambiguity, and context, sourced from the test set of ARC (Clark et al., ), GSM8K (Cobbe et al., ), HellaSwag (Zellers et al., ), MMLU (Hendrycks et al., ), Truthful QA (Lin et al., ), and Winogrande (Sakaguchi et al., ). This data sets the performance benchmarks on the HuggingFace Open LLM Leaderboard (Beeching et al., ) and provide a representative evaluation of performance on many language tasks.

## 5 Results

### 5.1 Coupling of Jacobians Across Depth

In trained LLMs, we observe coupling of the top singular vectors of the Jacobians across depth (Figure LABEL:fig:paris_all_8to16 bottom row), evident in the low non-diagonal values with a visible diagonal present in the matrix subplots. This is consistently observed across various LLMs considered. On the other hand, in untrained models (Figure LABEL:fig:paris_all_8to16 top row), there is no coupling of Jacobians across different depths. There is coupling along the diagonal, however, because each Jacobian is trivially diagonalized by its own singular vectors. This, in addition to Figure 2, suggests that coupling across depth emerges through training.

### 5.2 Coupling of Jacobians Across Tokens

We also analyze the coupling of singular vectors of the Jacobians across tokens. First, we observe coupling between Jacobians where the input and output tokens are the same, that is, between $J_{l}^{tt}$ and $J_{l^{\prime}}^{t^{\prime}t^{\prime}}$ across $t,t^{\prime}\in\{1,\ldots,n\}$ (Figure LABEL:fig:coupling_inoutsame). This suggests that a token’s interaction on its own trajectory is coupled with that of other tokens. We also consider the context tokens’ interaction on a trajectory when fixing the input token $t_{1}$ (Figure LABEL:fig:coupling_fixedin), through $J_{l}^{t_{1}t_{2}}$ and $J_{l^{\prime}}^{t_{1}t_{2}^{\prime}}$ across $t_{2},t_{2}^{\prime}\geq t_{1}$, and when fixing the output token $t_{2}$ (Figure LABEL:fig:coupling_fixedout), with $J_{l}^{t_{1}t_{2}}$ and $J_{l^{\prime}}^{t_{1}^{\prime}t_{2}}$ across $t_{1},t_{1}^{\prime}\leq t_{2}$. Although generally present, the coupling in these two cases are less consistent, with variation in the strength of coupling between different token pairs. Again in untrained models (Figure LABEL:fig:coupling_inoutsame top row), there is a lack coupling of Jacobians across tokens.

### 5.3 Emergence of Coupling with Training

Coupling emerges through training for the evaluated LLMs, including coupling across depth (Figure LABEL:fig:paris_all_8to16 ) and across tokens (Figures LABEL:fig:coupling_inoutsame, LABEL:fig:coupling_fixedin, LABEL:fig:coupling_fixedout). Further, we evaluate layer-wise coupling at intermediate training checkpoints of Pythia 12B, 6.9B, and 1B (Biderman et al., ) (Figure 2), and observe that coupling is generally low at initialization and increases significantly during early training steps.

### 5.4 Correlation with Generalization

For each LLM, we measure the average coupling across depth(which we define to be negative mis-coupling) across prompts in the 6 evaluation datasets (Section 4.2), where for each prompt, $m(A_{ll^{\prime},K}^{n})$ is averaged over layers $l,l^{\prime}\in\{1,\ldots,L\}$. We plot the coupling values against the benchmark scores across several LLMs (Figure 1). Our results reveal a positive correlation between coupling and performance benchmark scores, and is more significant than the relationship between other significant model hyperparameters (Figure LABEL:fig:correlation_hyperparams). This observation suggests a compelling relationship between stronger coupling of singular vectors of Jacobians $J_{l}^{t_{1}t_{2}}$ and improved generalization.

### 5.5 Regularity in Hidden Trajectories

We identify a considerable degree of linearity in the hidden trajectories of all featured LLMs. The LSS of the trajectories has an average value $4.25$ in trained models, while taking average values of $6.54$ at initialization, and is supported by the low variation across benchmark prompts (Figure 13). Linearity increases with training at varying depths of Pythia 12B (Figure LABEL:figs:lss-training). Linear and expansive behavior of the representations is demonstrated in Llama-3 70B, Mistral 30B, and NeoX 20B through low dimensional projections of embedding trajectories (Figure LABEL:figs:pca).

Figure: Figure 13: Linearity of Trajectories. The figure depicts the line-shape score (LSS) of embedding trajectories, as discussed in Section 5.5), computed on 1,200 prompts of the HuggingFace Open LLM Leaderboard (Section 4.2) for a variety of trained (black) and randomly initialized (red) LLMs (Appendix A.1). Plotted are the median values over all prompts, and are accompanied with uncertainty intervals depicting the inter-quartile range of the results for each model. Models are sorted by number of parameters.
Refer to caption: extracted/5922317/figs/lss/nparams_lss.png

### 5.6 Layer-Wise Exponential Growth

Across all LLMs considered, the hidden trajectories exhibit exponential growth that emerges with training (Figure 14). The majority of the models exhibit a low coefficient of variation across prompts, showing the robustness of this property across a variety of tasks. The expodistance measurements show little dependence on model parameters, and take on consistently low values across all featured LLMs (Figure 14). In contrast, measurements at initialization show equally (rather than exponentially) distanced trajectories, as reflected by the low coefficients of variation in the norms of their layer-wise differences (Figure 22). Under certain assumptions, exponential spacing is justified in Section 6.2.

Figure: Figure 14: Coefficient of variation of layer-wise exponent. The figure presents the coefficient of variation of the layer-wise exponent (Section 3.3) computed on 1,200 prompts of the HuggingFace Open LLM Leaderboard (Section 4.2) for trained LLMs (Appendix A.1). Plotted are the median values over all prompts, and are accompanied with uncertainty intervals depicting the inter-quartile range of the results for each model. LLMs are sorted by increasing benchmark performance, while the vertical axis has a base 10 logarithmic scale.
Refer to caption: extracted/5922317/figs-neurips/expodistance.png

## 6 Discussion

### 6.1 Correlation with Generalization

As discussed in Section 5.4, stronger coupling across Jacobian singular vectors is related to improved generalization (Figure 1). This relationship may assist in evaluating the training of transformers. In particular, this finding suggests that the development of training methods to amplify coupling across transformer blocks may lead to favourable model performance.

### 6.2 Emergence of Exponential Spacing with Coupling

The emergence of exponential spacing in many LLMs can be viewed as a result of coupling under some assumptions on the homogeneity of the singular values across depth. Strong coupling suggests that the linearizations of hidden trajectories satisfy a linear difference equation (see Appendix A.4), given by

$$ $x^{l+1}-x^{l}=Ax^{l}$ $$

Considering the singular value decomposition of $A$ in Equation 18 as $A=U\Sigma V^{T}$, and assuming that there is coupling ($U=V$), the difference equation can be diagonalized in the same basis across depth. Further assuming that the top singular values of the Jacobian blocks display low variation across depth, (Figure LABEL:fig:svalues_trained), we expect

$$ $||x^{l}||\approx e^{\sigma l}\,||x^{0}||,$ (16) $$

where $\sigma$ is the top singular value. Since we assume $\sigma$ is independent of depth, the above norm is expected to be exponential in $l$, and offers an explanation to the measurements in Figure 14.

## 7 Related Work

Residual Networks. ResNets (He et al., ) have been viewed as an ensemble of shallow networks (Veit et al., ), with studies delving into the scaling behaviour of their trained weights (Cohen et al., ). The linearization of residual blocks by their Residual Jacobians was first explored by Rothauge et al. (), who examined Residual Jacobians and their spectra in the context of stability analysis, and later by Li & Papyan () who discovered Residual Alignment. Coupling of Jacobian singular vectors in LLM transformers extends previous results for classifier Resnets (Li & Papyan, ).
We show coupling in $J_{l}^{t_{1}t_{2}}$ across various tokens, which is specific to tokenization in transformers, in addition to demonstrating coupling of $J_{l}^{t_{1}t_{2}}$ across $l$, which was also identified analogously in ResNets (Li & Papyan, ). Our measurements quantify the emergence with training of linearity and spacing, including at different depths, which was previously observed mostly qualitatively. Further comparison is included in Appendix A.3.

Neural ODE. Introduced by (Chen et al., ), Neural ODE view ResNets as a discretized dynamical process, with past work (Sander et al., ) showing the convergence of Residual Networks to a system of linear ODE. (Zhong et al., ; Li et al., ) studied neural ODE specific perspective to transformers and other LLM architectures. In our work, the presence of residual connections and emergence of Jacobian Alignment in transformers suggests that a discretization of a simple ODE emerges in LLMs.

Information Processing in LLMs. Previously, Katz & Belinkov (); Bietti et al. () have described memory and semantic flow in intermediate states of transformers while van Aken et al. () have examined BERT hidden trajectories. Investigation of LLMs for reasoning tasks and interpretability remains an important research focus (Huang & Chang, ).

In-context learning. Previous works have observed that LLMs are able to perform in context learning (von Oswald et al., ; Bai et al., ; Ahn et al., ; Akyürek et al., ; Xie et al., ; Hahn & Goyal, ; Xing et al., ): the model is capable of performing a particular machine learning task through a few examples provided in a single prompt or context. Their results consider a single layer of trained self-attention as implementing a single step of a gradient-descent-like learning rules across transformers depth, minimizing the MSE of a linear model

$$ $x_{\text{final}}=\min_{x}||Ax-b||^{2}$ $$

A gradient-descent learning rule would cause depth-wise updates to representations of the form

$$ $x_{t+1}=x_{t}+\epsilon A^{T}(Ax_{t}-b)=(I+\epsilon A^{T}A)x_{t}-{\epsilon}A^{T}b$ $$

We should expect similarity in the matrices $I+{\epsilon}A^{T}A$ across depth, which is suggested by coupling.

Hidden Representation Dynamics. Previous studies (Gai & Zhang, ; Haber & Ruthotto, ; Ee, ), have focused on interpreting deep neural networks from a dynamical systems perspective. Throughout training, numerical experiments have shown that both plain networks and ResNets are aptly approximated by a geodesic curve in Wasserstein space (Gai & Zhang, ). Other interpretations involve partitioning activation space into basins of attraction, with fixed points representing discrete thoughts (Nam et al., ). For further related works, see (Geshkovski et al., ; ; Tarzanagh et al., ; Valeriani et al., ).

Structure in Hidden Representations. Emergence of regularity in last-layer representations during training was noticed in Neural Collapse (Papyan et al., ). Subsequent works examined structure of hidden representations and presented theoretical results describing the emergence of Neural Collapse (Wang et al., ; Parker et al., ; Zangrando et al., ; Garrod & Keating, ; Wang et al., ; Hoyt & Owen, ; Arous et al., ; Zarka et al., ; Ben-Shaul & Dekel, ; Papyan, ; Súkeník et al., ). Within LLMs, recent works have identified uniform structure in token means (Wu & Papyan, ), and low dimensional structure across hidden trajectories (Sarfati et al., ). Our work studies the representations of all layers in LLMs in a local sense through the interactions of tokens through their Jacobians.

## 8 Limitations

A limitation of the scope of our analysis is that due to the significant compute resources required to train LLMs of the size studied in this paper, our exploration is restricted to pretrained LLMs and those fine-tuned on them. Several of these models are trained independently of each other, without carefully controlling for different experimental parameters. Consequently, direct comparisons between models become challenging, and so determining the cause of varying levels of regularity between LLMs becomes difficult. Nonetheless, this makes the consistent emergence of these properties interesting for further study.

## 9 Conclusion

Our primary goal was to contribute to the understanding of the mechanics underlying transformer architectures through an analysis of the trajectories of token embeddings and their interactions. Our research builds on the understanding of transformer architectures by revealing the coupling, across depth and token, of singular vectors in the Jacobians of transformer blocks for multiple LLMs trained by various organizations. We additionally observe and quantify the emergence of linearity and exponential spacing with training in the embedding trajectories of LLMs. Notably, we establish a correlation between the strength of this coupling and benchmark performance on the HuggingFace Open LLM Leaderboard, highlighting the significance of transformer block coupling for generalization. These findings open avenues for future research, encouraging deeper exploration into the connections between regularity of hidden representations, model specifications, and generalization.

## Appendix A Appendix and Supplementary material

### A.1 Suite of Large Language Models and Prompt Data

We evaluate 38 total LLMs (24 base LLMs and 14 fine-tuned) on 6 datasets from the HuggingFace Open LLM leaderboard (Beeching et al., ).

Falcon (40B, 7B) (Almazrouei et al., ), Llama-3 (70B, 8B) (AI@Meta, ), Llama-2 (70B, 13B, 7B) (Touvron et al., ), MPT (30B, 7B) (Team, ; ), Mistral v0.1 (Jiang et al., ), Gemma (7B, 2B) (Team et al., ), Gemma 1.1 (7B, 2B), NeoX (20B) (Black et al., ), Neo (2.7B) (Black et al., ; Gao et al., ), Pythia (6.9B).(Biderman et al., ), and GPT-2 (1.5B, 774M, 355M, 117M) (Radford et al., ).

CodeLlama (34B, 13B, 7B) (Rozière et al., ), CodeLlama Instruct (34B, 13B, 7B) (Rozière et al., ), Mistral-v0.1 Instruct (7.3B) (Jiang et al., ), Mistral-v0.2 Instruct (Jiang et al., ), CodeGemma (Team et al., ).

ARC (Clark et al., ), HellaSwag (Zellers et al., ), MMLU (Hendrycks et al., ), Truthful QA (Lin et al., ), WinoGrande (Sakaguchi et al., ).

**Table 1: LLMs featured in the experiments throughout paper. Included in the table is the parameter budget of each model, the embedding dimension, the number of training tokens, and the Open LLM leaderboard (Beeching et al., ) benchmark score.**
| Model | Param. | Layers ($L$) | Dim. ($d_{\text{model}}$) | Tokens | Score | Pos.Encoding |
| --- | --- | --- | --- | --- | --- | --- |
| Llama-3 | 70 B | 80 | 8192 | 15 T |  | RoPE |
|  | 8 B | 32 | 4096 | 15 T | 62.35 | RoPE |
| Llama-2 | 70 B | 80 | 8192 | 2 T | 66.05 | RoPE |
|  | 13 B | 40 | 5120 | 2 T | 55.69 | RoPE |
|  | 7 B | 32 | 4096 | 2 T | 50.97 | RoPE |
| Codellama | 34 B | 48 | 8192 | 2 T | 55.33 | RoPE |
|  | 13 B | 40 | 5120 | 2 T | 45.82 | RoPE |
|  | 7 B | 32 | 4096 | 2 T | 39.81 | RoPE |
| Codellama (IT) | 34 B | 48 | 8192 | 2 T | 43.0 | RoPE |
|  | 13 B | 40 | 5120 | 2 T | 37.52 | RoPE |
|  | 7 B | 32 | 4096 | 2 T | 40.05 | RoPE |
| Orca | 13 B | 40 | 5120 | 2 T | 58.64 | RoPE |
|  | 7 B | 32 | 4096 | 2 T | 54.55 | RoPE |
| Falcon | 40 B | 60 | 8192 | 1 T | 58.07 | RoPE |
|  | 7 B | 32 | 4544 | 1.5 T | 44.17 | RoPE |
| Falcon (IT) | 40 B | 60 | 8192 | 1 T | 43.26 | RoPE |
|  | 7 B | 32 | 4544 | 1.5 T |  | RoPE |
| MPT | 30 B | 48 | 7168 | 1 T | 66.98 | ALiBi |
|  | 7 B | 32 | 4096 | 1 T | 56.83 | ALiBi |
| Phi | 2 B | 32 | 2560 | 1.4 T | 61.33 | RoPE |
|  | 1.5 B | 24 | 2048 | 150 B | 47.69 | RoPE |
|  | 1 B | 24 | 2048 | 54 B |  | RoPE |
| Mistral-v0.1 | 7.3 B | 32 | 4096 |  | 60.97 | ALiBi |
| Mistral-v0.1 (IT) | 7.3 B | 32 | 4096 |  | 54.96 | ALiBi |
| Mistral-v0.2 (IT) | 7.3 B | 32 | 4096 |  | 65.71 | ALiBi |
| Gemma | 7 B | 28 | 3072 | 6 T | 64.29 | RoPE |
|  | 2 B | 18 | 2048 | 6 T | 42.75 | RoPE |
| Gemma-1.1 | 7 B | 28 | 3072 | 6 T | 30.0 | RoPE |
|  | 2 B | 18 | 2048 | 6 T | 60.09 | RoPE |
| Codegemma | 7 B | 28 | 3072 | 6 T | 56.73 | RoPE |
|  | 2 B | 18 | 2048 | 6 T | 32.19 | RoPE |
| Neo | 20 B | 44 | 6144 |  | 41.69 | RoPE |
|  | 2.7 B | 32 | 2560 | 0.42 | 36.20 | Sine |
| Pythia | 12 B | 36 | 5120 | 0.3 T | 58.9 | RoPE |
|  | 6.9 B | 32 | 4096 | 0.3 T | 39.30 | RoPE |
|  | 2.8 B | 32 | 2560 | 0.3 T |  | RoPE |
|  | 1.4 B | 24 | 2048 | 0.3 T |  | RoPE |
|  | 1 B | 16 | 2048 | 0.3 T |  | RoPE |
| GPT-2 | 1.5 B | 48 | 1600 |  | 34.12 | Sine |
|  | 774 M | 36 | 1280 |  | 32.07 | Sine |
|  | 355 M | 24 | 1024 |  | 29.87 | Sine |
|  | 117 M | 12 | 768 |  | 28.53 | Sine |

### A.2 Additional Metrics

#### A.2.1 Visualization of Trajectories with PCA

Each token, with initial embedding $x_{i}^{0}$, forms a trajectory $x_{i}^{0},x_{i}^{1},\ldots,x_{i}^{L}$ as it passes through the $L$ transformer blocks. The dynamics in high-dimensional space are visualized through a 2-dimensional principal component (PC) projection, $\text{PC}_{L}$, fitted to the last layer embeddings $X^{L}=(x_{1}^{L},x_{2}^{L},\ldots x_{n}^{L})$. The projected embeddings, $\text{PC}_{L}(x_{i}^{0}),\,\text{PC}_{L}(x_{i}^{1}),\,\ldots,\,\text{PC}_{L}(x
_{i}^{L})$, are plotted for each of the $i=1,\ldots,n$ trajectories.

### A.3 Comparison to ResNets

Coupling. Our results complement and build upon those of (Li & Papyan, ), who have observed the coupling of singular vectors of Residual Jacobians in classification ResNets. We observe coupling across depth in a wide range of LLMs, where Jacobians are evaluated at the sequence of embeddings, with respect to the current token, whereas in RA Jacobians are evaluated at a single representation. Additionally, with transformers we may analyze coupling not only across depth but also across tokens. We observe coupling in LLMs across tokens in a variety of ways. Further, we consider and analyze the relationship between coupling and generalization.

Linearity and Equidistance. Linearity in hidden trajectories, as observed in ResNets Li & Papyan (); Gai & Zhang (), also emerges with training in LLMs. The mean LSS value among the evaluated LLMs is $4.24$ (Figure 13), greater than LSS measurements observed for ResNets (Gai & Zhang (), page 18) which range between $2.0$-$3.0$ (due to varying trajectory length and hidden dimension). In both architectures, training induces improved linearity and regularity (Figure LABEL:figs:pca) in trajectories. In contrast to ResNets, trajectories are not equidistant, instead showing exponential growth between layers (Figure 14). We quantify this spacing through a low coefficient of variation, displaying the presence of exponential growth in token trajectories. In both classifier ResNets and LLM transformers, there is an evident level of regularity in hidden representations (Figure LABEL:figs:pca).

Rank of Jacobians. Li & Papyan () show that residual Jacobians have rank at most $C$, the number of classes. This analogous result automatically holds for LLMs since the vocabulary size is significantly greater than the embedding dimension of the transformer blocks.

Singular Value Scaling. Li & Papyan () observe that top singular values of residual Jacobians scale inversely with depth. In trained LLMs, however, top singular values do not show a consistent depth scaling across models (Figure LABEL:fig:svalues_trained), notably differing from classification ResNets. In addition, the distribution of singular values at each layer varies significantly between models. Singular value scaling is more present in untrained transformers (FigureLABEL:fig:svalues_untrained), likely caused by additional layer normalizations in residual blocks (Section 4.1).

### A.4 Dynamical Motivation

The equality of top left and right singular vectors suggests that the linearizations form a simple dynamical system that acts on representations. Consider a homogenous linear ODE:

$$ $x^{\prime}=Ax$ (17) $$

Its solution is given by

$$ $x(t)=\exp(tA)x(0)=\lim_{L\to\infty}\left(I+\frac{At}{L}\right)^{L}x(0)$ (18) $$

where the product is repeated $L$ times. Expanding the brackets shows that $x(t)$ can be thought of as a collection of many paths of various lengths, due to the binomial identity. This agrees with Veit et al. () which views ResNets as

$$ $x(t)=\left(I+\frac{tA_{1}}{L}\right)\left(I+\frac{tA_{2}}{L}\right)\dots\left( I+\frac{tA_{L}}{L}\right)x(0)$ (19) $$

However, Veit et al. () do not make any assumptions about the alignment of the various $A_{i}$ matrices. The coupling phenomenon suggests ResNets as implementing the simpler system

$$ $x(t)=\left(I+\frac{tA}{L}\right)\left(I+\frac{tA}{L}\right)\dots\left(I+\frac{ tA}{L}\right)x(0),$ (20) $$

where all the $A_{i}$ matrices are aligned. One benefit of this interpretation is that, as $L\to\infty$, we can write $x(t)$ in a simple closed form, namely $x(t)=\exp(tA)x(0)$.

Borrowing from the dynamical perspective, we quantify the similarity of hidden trajectories to the evolution of a linear ODE, in order to detect the emergence of a simple linearization to representations.

### A.5 Experimental and Implementation Details

The source code used to produce the results reported in this experiment has been included as supplemental material. Models with varying parameter sizes are loaded on GPUs with appropriate memory requirements: NVIDIA A40 ($n_{\text{param}}\geq 40B$), NVIDIA Quadro RTX 6000 for Gemma variants and when ($40B>n_{\text{param}}>13B$), and NVIDIA Tesla T4 when ($13B\geq n_{\text{param}}$) except Gemma variants. 1,200 prompts from the OpenLLM leaderboard were evaluated in variable batch sizes were queued on a SLURM cluster, with appropriate adjustments depending on the memory required to load the LLM.

- •
$13B\geq n_{\text{param}}$: 100 prompts per batch, except Gemma variants, which used 25 prompts per batch. The larger memory requirement for Gemma variants is likely due to the much larger vocabulary size in the model.
- •
$40B>n_{\text{param}}>13B$: 10 prompts per batch, except NeoX 20B which used 100 prompts per batch.
- •
$n_{\text{param}}\geq 40B$: 50 prompts per batch.

Due to the high memory requirement for computing Jacobians, for experiments involving the Jacobians, NVIDIA Quadro RTX 6000 was used additionally for $13B>n_{\text{param}}\geq 7B$ and corresponding models were quantized. Additionally, due to compute restrictions, alignment of singular vectors of Jacobians was computed on a smaller subset of the 1,200 prompts.

### A.6 Additional Experimental Results

Figure: Figure 15: Hidden Trajectories in LLMs. Principal components of the trajectories of the hidden representations through various LLMs (columns, decreasing in model size, see Table 1) in the prompt: What is the capital of France? The capital is. Top row: all layers. Middle Row: layers in shallower transformer blocks (layers specified above plot). Bottom Row: layers in deeper transformer blocks (layers specified above plot). Trajectories of each input token (last token ‘is’ is plotted in black) are plotted in latent space, visualized with a 2-dimension principal component projection. Representations proceed in distinct outward directions, especially in the second half of transformer blocks (lower row) during which the norm of representations increases, with possible abrupt change in the last layer (outer points in upper row). A clear direction of movement is visible in each token trajectory.
Refer to caption: extracted/5922317/figs/Paris_pca_all.jpg

Figure: Figure 21: Coefficient of variation of layer-wise equidistance. Variation of layer-wise equidistance (Section 3.3) computed over 1,200 prompts from the HuggingFace Open LLM Leaderboard datasets (Section 4.2) on a suite of untrained LLMs (Appendix A.1). Plotted are the median values over all prompts, and are accompanied with uncertainty intervals depicting the inter-quartile range of the results for each model. The models are sorted by increasing benchmark performance.
Refer to caption: extracted/5922317/figs-neurips/expodistance_untrained.png

Figure: Figure 22: Coefficient of variation of layer-wise expodistance. Variation of layer-wise expodistance (Section 3.3) computed over 1,200 prompts from the HuggingFace Open LLM Leaderboard datasets (Section 4.2) on a suite of untrained LLMs (Appendix A.1).
Refer to caption: extracted/5922317/figs/other/expodistance-training.png

Figure: Figure 23: Various Measurements of Representations. All measurements were made on 100 prompts taken from the WikiText 2 datsets. (Row 1) Norm of hidden representations as a function of layer depth. (Row 2) Line Shape Score (LSS) of the hidden trajectories as a function of layer depth. (Row 3) Mean equidistance of contiguous hidden trajectories as a function of depth. (Row 4) Entropy of logit vectors as a function of depth. Noted in most plots is a line where the behaviour of the measurement drastically changes.
Refer to caption: extracted/5922317/figs/other/all-llms.png

Figure: Figure 24: Linearity Emerges with Training. Two plots displaying the evolution of the linearity of the token trajectories through training. (Left) The LSS as a function of training checkpoint for the variants of the Pythia Scaling Suite Biderman et al. (). Here, the LSS is measured over each entire prompt. (Right) The mean LSS as a function of layer depth measured at various checkpoints throughout the Pythia 12B model. Here, the LSS is computed on a window of layers of width $11$, centred at the value given by the x-axis.
Refer to caption: extracted/5922317/figs/lss/lss-training.png

Figure: Figure 32: LSS sorted by LLM parameters.
Refer to caption: extracted/5922317/figs/other/lss_sorted.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses section above.
soundness: 3
strengths: - 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.
summary: 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.
weaknesses: - 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

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1) what is the value of fig 6/9?
2) what are your computational costs?
3) anything surprisingly found in the jacobians?
soundness: 3
strengths: - 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)
summary: 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.
weaknesses: 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.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 4
questions: 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.
soundness: 3
strengths: 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.
summary: 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.
weaknesses: 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)

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 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?
soundness: 3
strengths: 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.
summary: 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.
weaknesses: 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).

Meta-review: This work investigates the relationship between transformer block coupling and the generalization capabilities of large language models (LLMs). The authors claim that the interaction between different transformer blocks significantly influences the model's ability to generalize from training data to unseen examples. Through a series of experiments, they demonstrate that certain coupling configurations can enhance performance on various benchmarks, suggesting that optimizing block coupling could be a viable strategy for improving LLMs.

Reviewers acknowledged the paper's contributions, particularly its novel approach to understanding transformer architectures. The findings indicate a clear correlation between specific coupling strategies and improved generalization metrics, providing valuable insight for the community. On the other hand, reviewers also pointed out some concerns about the submission regarding the comprehensiveness and confidence of the analysis. The authors conduct thorough clarifications in the rebuttal period to address these concerns and improve the quality of this submission.

Based on the strengths outlined above and the effective responses to reviewer concerns during the rebuttal period, I recommend accepting this paper. The authors should carefully follow the reviewers' suggestions to revise the manuscript in the camera-ready stage.

Decision: Accept (Poster)