The Generalization Ridge: Information Flow in Natural Language Generation

## Abstract

Abstract Transformer-based language models have achieved state-of-the-art performance in natural language generation (NLG) tasks, yet their internal mechanisms for synthesizing task-relevant information remain insufficiently understood. While prior studies suggest that intermediate layers often yield more generalizable representations than final layers, how this generalization ability emerges and propagates across layers during training remains unclear.
To address this gap, we propose InfoRidge, an information-theoretic framework, to characterize how predictive information—the mutual information between hidden representations and target outputs—varies across depth. Estimating this quantity enables us to trace the flow of task-relevant information throughout the model during training. Our experiments across various models and datasets reveal a consistent non-monotonic trend: predictive information peaks in upper-middle layers—forming a generalization ridge —before declining in final layers, reflecting a transition between generalization and memorization. To further investigate this phenomenon, we introduce residual scaling coefficients—trainable scalar parameters applied to each residual block—which serve as functional probes for assessing the relative importance of individual transformer layers. These coefficients reveal that, under distribution shift, models downweight final layers and increasingly rely on ridge layers, highlighting their role in generalization. Together, these findings offer new insights into the internal mechanisms of transformers and underscore the critical role of intermediate layers in supporting generalization.

## 1 Introduction

Transformer-based language models have achieved remarkable performance in natural language generation (NLG) tasks such as summarization, machine translation, and dialogue generation vaswani2017attention; dong2022survey. Nevertheless, we lack a rigorous understanding of how these models acquire and synthesize task-relevant information during training.

A growing body of research has shown that intermediate layers in deep neural networks often surpass final layers in terms of representational quality and generalization performance liu2019linguistic; voita2019bottom; ansuini2019intrinsic; ahrens2023read; uselis2025intermediate. In language models, intermediate layers often encode richer semantic and more robust features than final layers fan2024not; jin2024exploring; skean2025layer. However, questions still remain: In NLG, how does information evolve across layers during training, and how are different layers of the network functionally organized to support generalization versus memorization?

To investigate these questions, we propose InfoRidge, an information-theoretic framework, to analyze information flow in language models. Building on matrix-based mutual information estimation giraldo2014measures, our approach quantifies how predictive signals transform across layers during training.
We center our analysis on two complementary quantities. The first is the predictive information at each layer, defined as the mutual information $I(Z_{\ell};Y)$ between the hidden representation $Z_{\ell}$ at layer $\ell$ and the next-token label $Y$, which reflects how much task-relevant information is preserved. The second is the incremental information gain, denoted as $I(\Delta Z_{\ell};Y)$, where $\Delta Z_{\ell}$ is the residual changes between successive layer $\ell$ and $\ell-1$. This measures the additional predictive information introduced by each transformer block.

Using these tools, we uncover a non-monotonic trend: predictive information rises through early and middle layers, peaks in the upper-middle, and then declines in later layers. We term this phenomenon the generalization ridge, indicating a structural division of labor across depth. Intermediate layers appear to specialize in encoding abstract, generalizable features, while later layers increasingly reorient representations toward task-specific alignment and memorization. By further analyzing incremental information gain, we find that certain middle layers introduce the largest increases in predictive information, marking them as key contributors to the emergence of the ridge. Such layer-wise specialization offers insight into how generalization and memorization are distributed within the network.

To further validate this interpretation, we introduce residual scaling coefficients—learnable scalar parameters $\beta_{\ell}$ applied to each residual block—while keeping all other model weights fixed, drawing inspiration from prior work on layer-wise adaptation liu2019self; menghani2024laurel. A higher $\beta_{\ell}$ value indicates greater reliance on the corresponding layer’s output during prediction. These coefficients act as functional probes, revealing how the model redistributes layer-wise importance under different data distribution. Under in-distribution training, deeper layers retain higher residual weights, reflecting the model’s reliance on memorized, task-specific features.
When evaluated under distribution shift, models reduce reliance on late layers and increase reliance on the ridge, further supporting its role in generalization.

To understand the formation of the generalization ridge, we analyze both attention patterns and model capacity. Attention flow analysis shows that ridge layers attend more directly to semantically informative tokens, suggesting that the ridge is not only a statistical peak but also functionally aligned with meaningful computation. In contrast, deeper layers increasingly focus on positional tokens, consistent with a shift toward memorization. Beyond functional behavior, we also investigate the conditions under which the ridge emerges. Our depth ablation results show that the ridge only emerges beyond a certain depth threshold. Below this threshold, predictive information increases monotonically, indicating that sufficient capacity is a prerequisite for generalization ridge to emerge.

#### Contributions.

Our work provides a unified perspective on how predictive information is structured across depth in transformer-based language models for natural language generation tasks:

- 1.
We propose InfoRidge, an information-theoretic framework, to quantify both predictive information and incremental information gain across layers, offering a dynamic view of how task-relevant signals emerge and evolve during training.
- 2.
Through this framework, we identify the generalization ridge—a mid-layer peak in predictive information—revealing a transition between generalization and memorization.
- 3.
We propose residual scaling coefficients as functional probes for layer importance, uncovering a shift in reliance from final layers to intermediate layers under distribution shift, consistent with the ridge’s role in supporting generalization.

These findings suggest that transformer depth reflects functional stratification: intermediate layers encode core abstractions for generalization, while final layers increasingly tailor representations for memorization—often at the cost of robustness. This insight informs the design of more interpretable and adaptable architectures under distribution shift.

## 2 Related Work

Understanding how information is encoded and transformed across layers in deep models has been widely studied through probing classifiers alain2016understanding, attention flow analysis vig2019analyzing, and information-theoretic tools such as the information bottleneck shwartz2017opening, mutual information estimation goldfeld2019estimating, and matrix-based entropy methods giraldo2014measures. These methods provide various lenses to quantify representational capacity, abstraction, and invariance across layers.

A growing body of research has shown that intermediate layers in deep networks often outperform final layers in terms of representational quality and task performance ansuini2019intrinsic; yosinski2014transferable; uselis2025intermediate; ahrens2023read; ando2023use. In language models, mid-depth layers tend to capture richer semantic or robust features than output layers liu2019linguistic; voita2019bottom; jin2024exploring; fan2024not. These findings challenge the assumption that deeper layers always yield better representations, revealing instead a peak in representational utility at intermediate depths. This pattern holds across settings such as transfer learning yosinski2014transferable, continual learning ahrens2023read, and out-of-distribution generalization uselis2025intermediate. Recent work by Skean et al. skean2025layer proposes a unified framework that evaluates representation quality using entropy, curvature, and invariance. Collectively, these results highlight the importance of intermediate layers for abstraction and generalization, motivating our analysis of layer-wise information dynamics through an information-theoretic lens.
Differently, we focus on information flow for next token generation, specifically examining mutual information between the last input token and the target label. Our goal is to understand how generalization and memorization dynamics evolve during training from information-theoretic perspective.

## 3 InfoRidge: Information Estimation Framework

Figure: Figure 1: Overview of InfoRidge. (1) Extract internal representations at each layer and compute residual changes between successive layers. (2) Extract the target token embedding. (3) Estimate predictive information and incremental information gain.
Refer to caption: x1.png

We propose InfoRidge, an information-theoretic framework that uses mutual information to quantify how predictive information propagates through transformers layers during training in NLG.

#### Motivating Insight.

Prior work has shown that internal representations in deep neural networks tend to align most closely with the true data distribution at an intermediate layer he2024information. By employing the Wasserstein distance villani2008optimal, this alignment is shown to reach a minimum at a specific depth—referred to as the *generalization funnel layer*. At this point, the *Min Wasserstein Generalization Bound* he2024information ensures that the upper bound on the generalization gap—defined as the expected difference between the population and empirical risks—is minimized. This highlights the critical role of intermediate layers in supporting generalization.

#### Hypothesis.

Building on this insight, we hypothesize that there exists a specific intermediate transformer layer that encodes the most generalizable representations for next-token prediction, characterized by maximal mutual information with the target label.

Hypothesis: Generalization Ridge.
There exists an intermediate layer $\ell^{*}\in\{1,\dots,L\}$ such that the mutual information between the hidden state and the target label peaks at that layer:

$$ $\ell^{*}=\arg\max_{\ell}I(Z_{\ell};Y).$ $$

We refer to this layer as the *generalization ridge*. This ridge layer aligns most strongly with generalizable features and serve as robust predictors under distribution shift, whereas later layers increasingly specialize in memorization.

#### InfoRidge Overview.

To empirically investigate this hypothesis, we propose InfoRidge, an information estimation framework, that characterizes how predictive information evolves across transformer layers. Specifically, we estimate two key quantities:

- •
Predictive Information $I(Z_{\ell};Y)$: the mutual information between the hidden state at layer $\ell$ and the target token. This quantity measures how much information about the true next token is contained in the layer’s full representation. A high value indicates that the layer encodes a strong and direct signal relevant to the prediction task.
- •
Incremental Information Gain $I(\Delta Z_{\ell};Y)$: the information introduced by the residual transformation at layer $\ell$, where $\Delta Z_{\ell}=Z_{\ell}-Z_{\ell-1}$. This captures the additional predictive signal gained through the residual transformation at layer $\ell$, isolating how much new task-relevant information is introduced on top of the previous layer’s representation.

Together, these metrics allow us to track both the accumulation and transformation of task-relevant information throughout the network.

#### Notation and Setup.

We consider a transformer model with $L$ residual blocks. Given an input sequence $x_{1:T}$ of length $T$, $z^{(\ell)}_{i}\in\mathbb{R}^{d}$ denotes the hidden state at the last token position of the $i$th input sequence in layer $\ell$, for $\ell=1,\dots,L$. For next-token prediction, the ground-truth label is denoted by $y_{i}\in\mathbb{R}^{d}$, corresponding to the embedding of the true next token from the vocabulary $\mathcal{V}$. The residual transformation introduced at layer $\ell$ is defined as $\delta z_{i}^{(\ell)}=z_{i}^{(\ell)}-z_{i}^{(\ell-1)}$.

Across a batch of $N$ sequences, we collect the representations:

$$ $Z_{\ell}=\{z^{(\ell)}_{i}\}_{i=1}^{N},\quad Y=\{y_{i}\}_{i=1}^{N},\quad\Delta Z _{\ell}=\{\delta z_{i}^{(\ell)}\}_{i=1}^{N},$ $$

where all vectors are $\ell_{2}$-normalized.

#### Computational Flow Overview.

Figure [1](https://arxiv.org/html/2507.05387v1#S3.F1) illustrates the workflow used to extract intermediate representations for information analysis. In the first forward pass, a batch of $N$ input sequences $x_{1:T}$ is fed into the transformer to obtain hidden states at each layer. From these, we extract the final-token representations $Z_{\ell}$ and compute the residual changes $\Delta Z_{\ell}$ by differencing consecutive layer outputs. In the second forward pass, each input is concatenated with its ground-truth next token $y$, and we extract the corresponding label embedding $Y$ from the output of the embedding layer. These representations are then used to compute two information-theoretic quantities: the Predictive Information $I(Z_{\ell};Y)$, and the Incremental Information Gain $I(\Delta Z_{\ell};Y)$. We estimate both $I(Z_{\ell};Y)$ and $I(\Delta Z_{\ell};Y)$ using Equation [1](https://arxiv.org/html/2507.05387v1#S3.E1) and [2](https://arxiv.org/html/2507.05387v1#S3.E2), detailed below.

#### Matrix-Based Mutual Information.

We apply the matrix-based framework giraldo2014measures to estimate mutual information. Let $\mathcal{U}$ be a random variable, from which we draw a set of vectors $U=\{\mathbf{u}_{i}\}_{i=1}^{N}\subset\mathbb{R}^{d}$. A positive-definite Gram matrix $G_{U}\in\mathbb{R}^{N\times N}$ is computed using a Gaussian kernel $\kappa$ with bandwidth set to 1 and the matrix is then trace-normalized to satisfy $\operatorname{tr}(G_{U})=1$. The matrix-based Rényi entropy (with order $\alpha=1$) is then given by:

$$ $H(\mathcal{U})\approx H(U)=-\operatorname{tr}(G_{U}\log G_{U}).$ (1) $$

Specifically, $G_{U}$ is constructed with entries
$(G_{U})_{ij}=\exp\left(-\frac{\|u_{i}-u_{j}\|^{2}}{2\sigma^{2}}\right)$ and then trace-normalized.

Let $G_{\mathcal{U}}$ and $G_{\mathcal{V}}$ be the trace-normalized Gram matrices for two random variables $\mathcal{U}$ and $\mathcal{V}$, respectively. The mutual information between them is computed as:

$$ $I(\mathcal{U};\mathcal{V})\approx I(G_{\mathcal{U}};G_{\mathcal{V}})=H(G_{ \mathcal{U}})+H(G_{\mathcal{V}})-H(G_{\mathcal{U}}\circ G_{\mathcal{V}}),$ (2) $$

where “$\circ$” denotes the Hadamard (elementwise) product. Mathematical details are in Appendix [A](https://arxiv.org/html/2507.05387v1#A1).

## 4 Experimental Setup

#### Models.

We evaluate four models: GPT–2 Small (117M) radford2019language, GPT–2 Medium (345M) radford2019language, Qwen–2.5 0.5B qwen2, and LLaMA 3.1 8B meta_llama_3.1_8b. All models are fine-tuned on NLG tasks. Each model shares weights between the input token embedding and the output language modeling head.

#### Datasets.

We assess model behavior across three tasks reformulated as NLG problems: CLUTRR (sinha2019clutrr,), a relational reasoning benchmark; ECQA (aggarwal2021explanations,), a commonsense QA benchmark; and Synthetic Arithmetic, a controlled dataset designed to disentangle task-relevant signal from noise. Dataset statistics and implementation details are provided in Appendix [B](https://arxiv.org/html/2507.05387v1#A2) and [C](https://arxiv.org/html/2507.05387v1#A3).

Synthetic Arithmetic Dataset Construction. We construct a synthetic dataset to separate signal learning from noise memorization. Each input is a sequence of 10 elements, where the signal follows an arithmetic progression modulo $K$, computed as $s_{t}=(s_{0}+t\cdot d)\bmod K$ with $s_{0}\in[0,K{-}1]$ and $d\in[1,K{-}1]$. Each element takes the form S{signal}_N{noise}, where noise is sampled from $\mathcal{U}_{\text{int}}(0,\texttt{noise_range-1})$ ($\mathcal{U}_{\text{int}}$ denotes the uniform distribution). The model is trained to predict the signal value of the final (10th) element using the preceding elements as input context. For example, with $K=5$, $s_{0}=1$, and $d=2$, a sample input might be S1_N42 S3_N77...S2_N37, with the target signal being 4. By varying $K$, we can induce structured distribution shifts.

## 5 Generalization Ridge: Layer-wise Mutual Information Training Dynamics

Understanding how layers encode task-relevant information is key to uncovering the internal mechanisms that support generalization in deep language models. In this section, we trace the evolution of two complementary forms of mutual information—Predictive Information and Incremental Information Gain—across transformer depth and training time, revealing a consistent structure in information flow and highlighting the generalization–memorization trade-off.

### 5.1 Predictive Information: Information Peaks at Intermediate Layers

We investigate how predictive information—defined as the mutual information between hidden representations and target labels—evolves across transformer layers. Specifically, we calculate the matrix–based mutual information $I(Z_{\ell};Y)$ at each transformer layer $\ell$, capturing the amount of label-relevant information encoded during forward propagation.

Figure: (a) GPT-2 Small on Synthetic
Refer to caption: x2.png

Figure [2](https://arxiv.org/html/2507.05387v1#S5.F2) tracks the trajectory of the predictive information $I(Z_{\ell};Y)$ between the hidden representation at depth $\ell$ and the target label $Y$ throughout training, while Table [1](https://arxiv.org/html/2507.05387v1#S5.T1) reports the downstream accuracy obtained when we early exit after a given layer.(^1^11The in distribution split is generated with $K_{\text{id}}{=}13$; out of distribution splits use a uniformly–sampled $K_{\text{ood}}\in\{5,\dots,25\}\setminus\{13\}$.) Additional results are presented in Appendix [E](https://arxiv.org/html/2507.05387v1#A5).

#### Three-phase information dynamics.

Across models and tasks, predictive information curves exhibit a consistent *three-phase* pattern:

Progressive Accrual (early layers). In the initial layers, $I(Z_{\ell};Y)$ gradually increases, corresponding to basic feature extraction without substantial task-level comprehension. This aligns with the near-zero accuracy observed in Table [1](https://arxiv.org/html/2507.05387v1#S5.T1) for these layers.

Table: Table 1: Information Dynamics and Layer-wise Performance (GPT-2 Small, Synthetic). OOD performance declines beyond the generalization ridge.

Information Peak (intermediate layers). $I(Z_{\ell};Y)$ continues to rise through the mid-to-upper layer, typically peaking before the final few blocks. For the GPT-2 Small on Synthetic Arithmetic dataset, the peak reaches $I\!\approx\!0.32$ in layers 10-11, coinciding with a jump from $\approx 0\%$ to $72\%$ ID accuracy and $53\%$ OOD accuracy. These layers appear to play a critical role in synthesizing abstract features that are essential for generalization.

Representational Compression (final layers). Beyond the peak, $I(Z_{\ell};Y)$ decreases, even as in-distribution accuracy approaches $100\%$. The simultaneous drop in OOD accuracy indicates that the final layers tend to memorize training patterns, sacrificing generalization ability.

Notably, this three-phase progression emerges consistently across both GPT-2 Small and GPT-2 Medium, indicating that the observed information dynamics are robust to architectural scale within the same model family.

#### Generalization Ridge: Memorization-Generalization Trade-off.

The pronounced “information funnel” around intermediate layers reflects a key trade-off between generalization and memorization, which we term the “generalization ridge”. These layers maximize task-relevant information for generalization, while deeper layers increasingly compress and specialize representations, enhancing in-distribution memorization but reducing robustness. This positions intermediate layers as critical control points for managing this trade-off.

#### Discussion on Model Capacity Threshold

Figure: (a) 8 layers: Monotonic Increase
Refer to caption: x6.png

We empirically vary the depth of GPT-2 Small models fine-tuned on the CLUTRR dataset to examine how model capacity shapes information dynamics, shown in Figure [3](https://arxiv.org/html/2507.05387v1#S5.F3). When the model is truncated to 8 layers, $I(Z;Y)$ increases monotonically—indicating insufficient capacity to develop an information peak. Upon increasing the depth to 9 layers, a shallow peak emerges, signifying the threshold at which the model begins to distinguish generalized features from memorized signals. The full 12-layer model exhibits a clear peak followed by a decline, confirming that generalization ridge emerges only beyond a certain capacity threshold. These results underscore the role of architectural depth in information dynamics and sufficient capacity is required for generalization ridge to emerge.

#### Discussion on Overfitting Scenario

Figure: Figure 4: $I(Z_{\ell};Y)$ rises in final layers during overfitting (LLaMA, ECQA).
Refer to caption: x9.png

To probe overfitting dynamics, we intentionally fine-tuned the model beyond the optimal point. In Figure [4](https://arxiv.org/html/2507.05387v1#S5.F4), we observe that $I(Z_{\ell};Y)$ rises again in the final layers—departing from the typical compression phase. This reversal may suggest that the model begins to memorize shortcuts or redundant label-specific noise rather than learning useful information.

### 5.2 Incremental Information Gain: Information Concentrates at Intermediate Layers

Figure: Figure 5: Middle transformer blocks are key to encoding generalizable task-relevant information (GPT-2 Small, Synthetic).
Refer to caption: x10.png

To understand how information accumulates across the network, we compute Incremental Information Gain ($I(\Delta Z_{\ell};Y)$)—the mutual information between each residual transition and the target label embedding. Figure [5](https://arxiv.org/html/2507.05387v1#S5.F5) shows the layer-wise gains, revealing that intermediate layers yield the highest information increases. This concentration of information gain highlights their central role in encoding task-relevant features essential for generalization. A case study decoding $\Delta z$ via the LM head is provided in Appendix [F](https://arxiv.org/html/2507.05387v1#A6), highlighting token-level shifts across layers. The Incremental Information Gain analysis reveals a clear pattern: middle transformer blocks are key to encoding generalizable task-relevant information, thereby forming a generalization ridge. Conversely, later layers contribute little additional predictive signal, and in some cases, actively reduce alignment with the target embeddings. This diminishing contribution in later layers may suggest a shift from general reasoning to memorization. This trend underscores a fundamental trade-off in transformer training dynamics.

## 6 Residual Scaling Dynamics via Learnable β 𝛽 \beta italic_β Coefficients

To deepen our understanding of the Generalization Ridge hypothesis, we examine how modulating the contribution of individual transformer blocks affects information flow, and propose a corollary that links layer-wise contribution to generalization performance under distribution shift.

###### Corollary (Residual Scaling) .

Post-peak residual blocks (i.e., layers beyond the generalization ridge) may encode memorized signals. Suppressing these layers’ contribution via learned residual scaling improves out-of-distribution (OOD) generalization, while amplifying them degrades it.

We introduce a residual scaling mechanism with learnable scalar coefficient parameters (Algorithm [1](https://arxiv.org/html/2507.05387v1#alg1)), inspired by prior work on adaptive residual modulation liu2019self; menghani2024laurel. Transformer architectures inherently employ residual connections to iteratively refine representations. To isolate and quantify the contribution of each transformer blocks, we scale these residual connections with layer-specific scaling factors $\beta_{\ell}\in\mathbb{R}_{\geq 0}$, with definition below. Each $\beta_{\ell}$ controls the strength of the residual contribution from layer $\ell$, enabling the model to adaptively emphasize or suppress specific blocks:

$$ $z^{(\ell)}=z^{(\ell-1)}+\beta_{\ell}\cdot\text{block}^{(\ell)}(z^{(\ell-1)}), \qquad\beta_{\ell}\in\mathbb{R}_{\geq 0}.$ $$

###### Definition (Residual Scaling Coefficient) .

$\beta_{\ell}$ is a learnable scalar parameter associated with transformer layer $\ell$, which modulates the contribution of that layer’s residual output to the model’s forward pass.

Figure: Algorithm 1 Residual Scaling: Probing the Contribution of Transformer Blocks to the Generalization–Memorization Trade-off.

We freeze model weights and optimize only the residual scaling coefficient parameters $\beta_{\ell}$, which are initialized to 1. These scalars are trained on the (a) in-distribution (ID) split and (b) out-of-distribution (OOD) split. Since no other parameters are updated, the learned $\beta_{\ell}$ values serve as a direct diagnostic of the extent to which each layer should be amplified or attenuated to suit the data regime.

As shown in Table [2](https://arxiv.org/html/2507.05387v1#S6.T2), optimizing the residual scaling coefficient parameters on in-distribution results in $\beta_{12}>1$, indicating that in-distribution performance benefits from amplifying the contribution of the final transformer layers—suggestive of their role in memorizing features specific to the training distribution. In contrast, when trained on out-of-distribution (OOD) data, the final layer’s contribution is downweighted ($\beta_{12}<1$), implying that suppressing memorization-heavy components improves generalization. This pattern holds consistently across architectures such as Qwen and LLaMA (Figure [6](https://arxiv.org/html/2507.05387v1#S6.F6)) and corresponds to improved OOD performance, as shown in Table [3](https://arxiv.org/html/2507.05387v1#S6.T3). Together, these findings provide empirical support for the functional specialization of transformer layers: intermediate layers tend to encode generalizable representations, whereas deeper layers increasingly capture memorized signals.

**Table 2: $I(Z;Y)$ and residual scaling coefficients $\beta_{\ell}$ for GPT-2 Small. Optimizing on in-distribution data yields $\beta_{12}>1$, while OOD training suppresses the final layer’s contribution with $\beta_{12}<1$.**
|  | Layer 1 | Layer 2 | Layer 3 | Layer 4 | Layer 5 | Layer 6 | Layer 7 | Layer 8 | Layer 9 | Layer 10 | Layer 11 | Layer 12 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $I(Z;Y)$ (CLUTRR) | 0.0526 | 0.0474 | 0.0407 | 0.0372 | 0.0319 | 0.0337 | 0.0837 | 0.2688 | 0.3430 | 0.3901 | 0.4514 | 0.0851 |
| $\beta_{\ell}$ (CLUTRR ID) | 0.9998 | 0.9999 | 1.0002 | 1.0002 | 1.0003 | 1.0004 | 0.9993 | 0.9997 | 0.9998 | 0.9993 | 0.9999 | 1.0003 |
| $\beta_{\ell}$ (CLUTRR OOD) | 1.0005 | 1.0005 | 1.0000 | 1.0000 | 0.9999 | 0.9997 | 1.0002 | 1.0003 | 1.0002 | 0.9995 | 0.9996 | 0.9995 |
| $I(Z;Y)$ (Synthetic) | 0.0508 | 0.0513 | 0.0619 | 0.0822 | 0.0894 | 0.0989 | 0.1292 | 0.2431 | 0.2810 | 0.3209 | 0.3209 | 0.1402 |
| $\beta_{\ell}$ (Synthetic ID) | 0.9993 | 1.0007 | 0.9995 | 1.0005 | 1.0000 | 0.9999 | 1.0000 | 0.9999 | 0.9995 | 1.0018 | 1.0003 | 1.0001 |
| $\beta_{\ell}$ (Synthetic OOD) | 1.0015 | 1.0003 | 0.9974 | 1.0006 | 1.0010 | 0.9980 | 0.9987 | 0.9974 | 0.9985 | 0.9982 | 0.9981 | 0.9956 |

**Table 3: Loss and accuracy (%) before and after residual scaling on out-of-distribution (OOD) split. Performance improves following the optimization of residual scaling coefficient parameters.**
| Dataset | Clutrr | Synthetic |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Metric | Loss${}_{\text{before}}$ | Acc${}_{\text{before}}$(%) | Loss${}_{\text{after}}$ | Acc${}_{\text{after}}$(%) | Loss${}_{\text{before}}$ | Acc${}_{\text{before}}$(%) | Loss${}_{\text{after}}$ | Acc${}_{\text{after}}$(%) |
| GPT-2 Small | 6.9655 | 21.24 | 6.9636 | 21.24 | 2.6533 | 70.53 | 2.6485 | 70.54 |
| Qwen-2.5-0.5B | 4.2494 | 31.01 | 4.2465 | 31.01 | 1.8132 | 74.31 | 1.7217 | 74.55 |
| LLaMA-3.1-8B | 1.3209 | 57.23 | 1.3193 | 57.13 | 1.3927 | 87.39 | 1.3155 | 87.83 |

Figure: Figure 6: Residual scaling coefficients $\beta_{\ell}$ for Qwen-2.5-0.5B and LLaMA-3.1-8B. A similar pattern consistently holds across architectures, with $\beta_{\text{final layer}}$ decreasing under OOD training.
Refer to caption: x11.png

###### Corollary (Residual Scaling) .

###### Definition (Residual Scaling Coefficient) .

## 7 Case Study: Attention Dynamics Across Layers

To verify whether generalization ridge layer indeed correspond to semantically meaningful processing, we visualize the attention patterns as a more interpretable signal of where the model focuses. As shown in Figure [7](https://arxiv.org/html/2507.05387v1#S7.F7), we visualize attention maps across layers.

Figure: Figure 7: Attention map across layers (GPT-2 Small, Synthetic). At the generalization ridge layers, attention becomes more targeted, focusing on predictive tokens.
Refer to caption: x12.png

In early layers (e.g., Layer 1), attention is diffuse and biased toward final tokens—reflecting reliance on position rather than true predict signals.

By generalization ridge layers (e.g., Layer 11), attention becomes more targeted, concentrating on predictive tokens. This shift marks a transition from positional attention to semantically meaningful focus, suggesting that intermediate layers are increasingly capable of isolating task-relevant information from irrelevant content.

In the final layers, attention regresses toward final tokens. This reversion is aligned with the observed decline in $I(Z_{\ell};Y)$. The resurgence of attention towards terminal tokens indicates a potential memorization scenario, where the model re-engages superficial positional strategies, possibly memorizing noise rather than further refining the generalized extraction of predictive signals.

This attention trajectory further supports the generalization ridge hypothesis, highlighting a trade-off between generalization and memorization in the model’s representational strategy.

## 8 Conclusion

We present InfoRidge, an information-theoretic framework, to analyze how information evolves across layers in transformer language models. By estimating both predictive information and incremental information gain, we systematically characterize the dynamics of task-relevant signal propagation throughout the network. Our findings reveal a consistent generalization ridge emerging in intermediate layers, where mutual information between the hidden representation and the target label reaches its peak. This phenomenon reflects a fundamental trade-off between generalization and memorization as information flows deeper into the model. Residual scaling experiments further corroborate this interpretation, demonstrating the functional specialization of different layers—where intermediate blocks play a key role in supporting generalization, while deeper layers increasingly focus on memorization. Attention flow provides additional insight into where the model concentrates information. Taken together, these results offer a unified perspective on information flow in transformers for natural language generation and point to opportunities for enhancing model architecture, improving interpretability, and increasing robustness to distributional shifts.

## 9 Limitation

Our analysis focuses on a select set of datasets and model architectures, which may limit the breadth of empirical coverage. The mutual information estimates rely on kernel-based approximations that may be sensitive to representation geometry. Additionally, the residual scaling approach assumes layer-wise independence and does not account for potential inter-layer dependencies during adaptation. Future work could address these limitations by extending the analysis to broader model families and tasks and exploring joint optimization strategies.

## References

- [1]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
Advances in neural information processing systems, 30, 2017.
- [2]
Chenhe Dong, Yinghui Li, Haifan Gong, Miaoxin Chen, Junxin Li, Ying Shen, and Min Yang.
A survey of natural language generation.
ACM Computing Surveys, 55(8):1–38, 2022.
- [3]
Nelson F Liu, Matt Gardner, Yonatan Belinkov, Matthew E Peters, and Noah A Smith.
Linguistic knowledge and transferability of contextual representations.
In Proceedings of NAACL-HLT, pages 1073–1094, 2019.
- [4]
Elena Voita, Rico Sennrich, and Ivan Titov.
The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives.
arXiv preprint arXiv:1909.01380, 2019.
- [5]
Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan.
Intrinsic dimension of data representations in deep neural networks.
Advances in Neural Information Processing Systems, 32, 2019.
- [6]
Kyra Ahrens, Hans Hergen Lehmann, Jae Hee Lee, and Stefan Wermter.
Read between the layers: Leveraging multi-layer representations for rehearsal-free continual learning with pre-trained models.
arXiv preprint arXiv:2312.08888, 2023.
- [7]
Arnas Uselis and Seong Joon Oh.
Intermediate layer classifiers for ood generalization.
arXiv preprint arXiv:2504.05461, 2025.
- [8]
Siqi Fan, Xin Jiang, Xiang Li, Xuying Meng, Peng Han, Shuo Shang, Aixin Sun, Yequan Wang, and Zhongyuan Wang.
Not all layers of llms are necessary during inference.
arXiv preprint arXiv:2403.02181, 2024.
- [9]
Mingyu Jin, Qinkai Yu, Jingyuan Huang, Qingcheng Zeng, Zhenting Wang, Wenyue Hua, Haiyan Zhao, Kai Mei, Yanda Meng, Kaize Ding, et al.
Exploring concept depth: How large language models acquire knowledge and concept at different layers?
arXiv preprint arXiv:2404.07066, 2024.
- [10]
Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv.
Layer by layer: Uncovering hidden representations in language models.
arXiv preprint arXiv:2502.02013, 2025.
- [11]
Luis Gonzalo Sanchez Giraldo, Murali Rao, and Jose C Principe.
Measures of entropy from data using infinitely divisible kernels.
IEEE Transactions on Information Theory, 61(1):535–548, 2014.
- [12]
Fenglin Liu, Meng Gao, Yuanxin Liu, and Kai Lei.
Self-adaptive scaling for learnable residual structure.
In Proceedings of the 23rd Conference on Computational Natural Language Learning (CoNLL), pages 862–870, 2019.
- [13]
Gaurav Menghani, Ravi Kumar, and Sanjiv Kumar.
Laurel: Learned augmented residual layer.
arXiv preprint arXiv:2411.07501, 2024.
- [14]
Guillaume Alain and Yoshua Bengio.
Understanding intermediate layers using linear classifier probes.
arXiv preprint arXiv:1610.01644, 2016.
- [15]
Jesse Vig and Yonatan Belinkov.
Analyzing the structure of attention in a transformer language model.
arXiv preprint arXiv:1906.04284, 2019.
- [16]
Ravid Shwartz-Ziv and Naftali Tishby.
Opening the black box of deep neural networks via information.
arXiv preprint arXiv:1703.00810, 2017.
- [17]
Ziv Goldfeld.
Estimating information flow in deep neural networks.
In International Conference on Machine Learning, 2019.
- [18]
Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson.
How transferable are features in deep neural networks?
Advances in neural information processing systems, 27, 2014.
- [19]
Atsushi Ando, Ryo Masumura, Akihiko Takashima, Satoshi Suzuki, Naoki Makishima, Keita Suzuki, Takafumi Moriya, Takanori Ashihara, and Hiroshi Sato.
On the use of modality-specific large-scale pre-trained encoders for multimodal sentiment analysis.
In 2022 IEEE Spoken Language Technology Workshop (SLT), pages 739–746. IEEE, 2023.
- [20]
Haiyun He, Christina Lee Yu, and Ziv Goldfeld.
Information-theoretic generalization bounds for deep neural networks.
arXiv preprint arXiv:2404.03176, 2024.
- [21]
Cédric Villani et al.
Optimal transport: old and new, volume 338.
Springer, 2008.
- [22]
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.
Language models are unsupervised multitask learners.
2019.
- [23]
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zhihao Fan.
Qwen2 technical report.
arXiv preprint arXiv:2407.10671, 2024.
- [24]
Meta AI.
Meta llama 3.1 8b.
[https://huggingface.co/meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B), 2024.
Accessed: 2025-05-12.
- [25]
Koustuv Sinha, Shagun Sodhani, Jin Dong, Joelle Pineau, and William L. Hamilton.
Clutrr: A diagnostic benchmark for inductive reasoning from text.
Empirical Methods of Natural Language Processing (EMNLP), 2019.
- [26]
Shourya Aggarwal, Divyanshu Mandowara, Vishwajeet Agrawal, Dinesh Khandelwal, Parag Singla, and Dinesh Garg.
Explanations for commonsenseqa: New dataset and models.
In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 3050–3065, 2021.
- [27]
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.
Parameter-efficient transfer learning for nlp.
In International conference on machine learning, pages 2790–2799. PMLR, 2019.
- [28]
Zhengxuan Wu, Aryaman Arora, Zheng Wang, Atticus Geiger, Dan Jurafsky, Christopher D Manning, and Christopher Potts.
Reft: Representation finetuning for language models.
Advances in Neural Information Processing Systems, 37:63908–63962, 2024.
- [29]
Jing Huang, Zhengxuan Wu, Christopher Potts, Mor Geva, and Atticus Geiger.
Ravel: Evaluating interpretability methods on disentangling language model representations.
arXiv preprint arXiv:2402.17700, 2024.
- [30]
Chunyuan Deng, Ruidi Chang, and Hanjie Chen.
Learning distribution-wise control in representation space for language models.
arXiv preprint arXiv:2506.06686, 2025.
- [31]
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov.
Locating and editing factual associations in gpt.
Advances in neural information processing systems, 35:17359–17372, 2022.
- [32]
Zhengxuan Wu, Atticus Geiger, Thomas Icard, Christopher Potts, and Noah Goodman.
Interpretability at scale: Identifying causal mechanisms in alpaca.
Advances in neural information processing systems, 36:78205–78226, 2023.

## Appendix A Mathematical Details for Matrix-Based Information Estimation and Theoretical Foundations

We employ the matrix-based Rényi entropy  [11] to estimate mutual information between model representations and target labels. This method captured sample similarity structure via kernel Gram matrices.

### A.1 Matrix-Based Entropy Estimation

Let $U=\{u_{i}\}_{i=1}^{N}\subset\mathbb{R}^{d}$ denote $\ell_{2}$-normalized representations obtained from a specific transformer layer, a residual update, or the embedding of the target label. A Gaussian kernel Gram matrix $G_{U}\in\mathbb{R}^{N\times N}$ is constructed as:

$$ $(G_{U})_{ij}=\exp\left(-\frac{\|u_{i}-u_{j}\|^{2}}{2\sigma^{2}}\right),$ $$

with bandwidth $\sigma=1$. The matrix is then trace-normalized to ensure $\text{tr}(G_{U})=1$.

The matrix-based Rényi entropy of order $\alpha=1$ is defined as:

$$ $H(U)=-\text{tr}(G_{U}\log G_{U}).$ $$

This expression can be interpreted in terms of the eigenvalue spectrum $\{\lambda_{k}\}$ of $G_{U}$, since $G_{U}$ is positive semi-definite and trace-normalized:

$$ $H(U)=-\sum_{k=1}^{N}\lambda_{k}\log\lambda_{k}.$ $$

The entropy thus reflects the dispersion of the eigenvalues. A more uniform spectrum (i.e., higher entropy) suggests more diversity in the representation space, while a sharply peaked spectrum (i.e., low entropy) indicates redundancy or compression.

### A.2 Mutual Information Estimation

To estimate the mutual information between two random variables $U$ and $V$, we compute their Gram matrices $G_{U}$ and $G_{V}$, and form the joint similarity matrix via element-wise (Hadamard) product:

$$ $G_{UV}=G_{U}\circ G_{V}.$ $$

After trace-normalization, mutual information is estimated by:

$$ $I(U;V)=H(U)+H(V)-H(U,V),$ $$

where $H(U,V)=-\text{tr}(G_{UV}\log G_{UV})$. The eigenvalue spectrum of $G_{UV}$ governs the joint entropy term; its shape reflects how much of the structure in $U$ and $V$ aligns. A more concentrated spectrum in $G_{UV}$ relative to $G_{U}$ and $G_{V}$ implies stronger dependence and thus higher mutual information.

### A.3 Min Wasserstein Generalization Bound [ 20 ]

We restate the generalization bound proposed by He et al. [20], which characterizes generalization in terms of the Wasserstein distance between internal representations.

Suppose that the loss function $\tilde{\ell}:\mathcal{Y}\times\mathcal{Y}\to\mathbb{R}_{\geq 0}$ is $\rho_{0}$-Lipschitz, and the activation function $\phi_{\ell}:\mathbb{R}\to\mathbb{R}$ is $\rho_{\ell}$-Lipschitz for each $\ell=1,\ldots,L$. Then:

$$ $\small\mathrm{gen}(P_{W|\mathcal{D}_{n}}\!,\!P_{X,Y})\!\leq\!\min_{\ell}\frac{ \rho_{0}}{n}\sum_{i=1}^{n}\!\mathbb{E}_{W}\!\left[\!\left(1\!\vee\!\prod_{j= \ell+1}^{L}\!\rho_{j}\|W_{j}\|_{\mathrm{op}}\!\right)\!\cdot\!\mathcal{W}_{1} \!\left(\!P_{T_{\ell},i,Y_{i}|W}(\cdot|W),\!\ P_{T_{\ell},Y|W}(\cdot|W)\! \right)\!\right]$ (3) $$

This result shows that generalization can be tightly controlled by the Wasserstein distance [21] between internal representations at a specific layer—referred to as the generalization funnel layer.

#### Connection to Our Work.

The Min Wasserstein Generalization Bound [20] provides a theoretical foundation for our study by characterizing generalization in terms of distributional alignment at an intermediate layer. It motivates our analysis of information flow by suggesting that information is peaked at a specific layer—the generalization funnel. Our InfoRidge framework builds on this insight by quantifying predictive information across layers, and reveals that a specific intermediate layer exhibits peak mutual information and correlates with better generalization performance.

## Appendix B Dataset Overview and Statistics

To evaluate information flow and generalization dynamics across model layers, we conduct experiments on three datasets with varying levels of complexity and structure: CLUTRR, ECQA, and a custom-designed Synthetic Arithmetic dataset. Table B.1 summarizes key dataset statistics.

### B.1 Dataset Overview

**Table 4: Dataset Statistics**
| Dataset | #Train | Train Seq. Len | #Val | Val Seq. Len | #Test | Test Seq. Len |
| --- | --- | --- | --- | --- | --- | --- |
| CLUTRR | 9,074 | 30 | 2,020 | 29 | 1,146 | 70 |
| ECQA | 7,598 | 21 | 1,090 | 21 | 2,194 | 21 |
| Synthetic Arithmetic | 10,000 | 9 | 2,000 | 9 | 2,000 | 9 |

### B.2 CLUTRR

CLUTRR (Compositional Language Understanding and Text-based Relational Reasoning) [25] is a diagnostic benchmark for evaluating relational inference in language models. Each example contains a story describing family relations, and the task is to infer the missing relationship between two entities. The distribution shift stems from clause lengths that are absent in the training set but present during evaluation. We use the task split “gen_train23_test2to10”, where the model is trained on clause lengths 2 and 3 and evaluated on lengths 2 through 10.

### B.3 ECQA

ECQA (Explanations for CommonsenseQA) [26] is a commonsense multiple-choice question-answering dataset, where each question is accompanied by 5 answer options.

### B.4 Synthetic Arithmetic Dataset

We construct a synthetic diagnostic dataset to disentangle task-relevant signal learning from spurious noise memorization in a controlled setting. Each sample consists of a sequence of 10 symbolic elements, where the signal component follows an arithmetic progression modulo $K$, and the remainder of each element is independently corrupted with random noise. By varying the modulus $K$, we systematically control task complexity and introduce structured shifts in the data distribution.

#### Synthetic Arithmetic Dataset Construction.

At each position $t$, the signal value is computed as:

$$ $s_{t}=(s_{0}+t\cdot d)\bmod K,\quad\text{with }s_{0}\in[0,K{-}1],\quad d\in[1, K{-}1].$ $$

Each element in the sequence is represented as a string of the form:

$$ $\texttt{S\{signal\}_N\{noise\}},\quad\text{where noise}\sim\mathcal{U}_{\text {int}}(0,\texttt{noise_range-1}).$ $$

Here $\mathcal{U}_{\text{int}}$ denotes the uniform distribution. The model is trained to predict the signal value of the final (10th) element, using the preceding elements as input context.

For example, with $K=5$, $s_{0}=1$, and $d=2$, a sample might look like:

> S1_N42 S3_N77 S0_N18 S2_N56 S4_N90 S1_N11 S3_N65 S0_N23 S2_N37

Each element encodes both a signal (the number following S) and a noise component (the number following N). The target is 4, corresponding to the signal of the final (10th) item in the sequence.

## Appendix C Implementation Details

This appendix outlines implementation details in our experiments.

### C.1 Prompt Construction

All tasks are cast into a next-token generation format. The model receives a prompt (context) and must generate the correct continuation token. Below are construction strategies and examples for each dataset:

#### CLUTRR

Each input example in CLUTRR consists of a short narrative describing a set of family relationships, along with a query involving a pair of entities. We construct prompts by concatenating the narrative and a structured natural language question derived from the query tuple. The model is trained to predict the correct relationship as the next token.

Prompt:

> Story: [Alice] is [Bob]’s mother. [Bob] is [Charlie]’s father. Query: What is the relationship between Alice and Charlie? Answer:

Target: grandmother

#### ECQA (Explanation-augmented Commonsense QA)

Each ECQA instance consists of a multiple-choice question with five candidate answers. We format the prompt by presenting the question followed by all five options (labeled A–E), and conclude with an explicit answer query. The model is trained to predict the correct answer letter as the next token.

Prompt:

> Question: What do people usually do at a birthday party? Options: A. Sleep B. Celebrate C. Cook D. Exercise E. Drive Answer:

Target: B

#### Synthetic Arithmetic

Each synthetic sample consists of a sequence of 10 symbolic elements, where each element is formatted as S{signal}_N{noise}. The signal values follow an arithmetic progression modulo $K$, and the noise values are independently sampled from a uniform distribution with a fixed range of 100. During training, the modulus $K$ is set to 13. For evaluation, test sequences are generated using values of $K$ from the range $[5,26]$ excluding 13 to simulate a distribution shift. In the residual $\beta_{\ell}$ analysis, we use $K=13$ for in-distribution (ID) training and $K=17$ for out-of-distribution (OOD) training, allowing for a controlled comparison between generalization and memorization behavior. The model receives the first 9 tokens as input and is trained to predict the signal component of the 10th token.

Prompt:

> S1_N42 S3_N88 S5_N20 S7_N10 S9_N65 S11_N43 S0_N99 S2_N38 S4_N77

Target: 6

This controlled format enables manipulation of distributional properties by varying the modulus $K$.

### C.2 Information Estimation

To estimate mutual information, we subsample between 50 and 200 test examples, with the exact number determined by the model and task, in order to avoid excessive memory usage from large Gram matrices. Full results and statistical significance are reported in Appendix [E](https://arxiv.org/html/2507.05387v1#A5).

### C.3 Residual Scaling with Learnable β ℓ subscript 𝛽 ℓ \beta_{\ell} italic_β start_POSTSUBSCRIPT roman_ℓ end_POSTSUBSCRIPT Parameters

We introduce a vector of learnable scalar weights $\beta=\{\beta_{1},\dots,\beta_{L}\}$ applied to residual connections in a frozen transformer:

$$ $z^{(\ell)}=z^{(\ell-1)}+\beta_{\ell}\cdot\text{Block}^{(\ell)}(z^{(\ell-1)}).$ $$

- •
All transformer weights are frozen; only $\beta_{\ell}$ parameters are trained.
- •
Each transformer block’s residual output is modulated by $\beta_{\ell}$ via forward hooks.
- •
$\beta$ values are initialized to 1 and updated using gradient descent without altering model architecture.

#### Training and Evaluation Settings

We perform residual $\beta_{\ell}$ training on top of a model that has already been fine-tuned on the target dataset, keeping all previously learned weights frozen and optimizing only the $\beta_{\ell}$ parameters using HuggingFace’s Trainer API. We sweep over batch sizes from 4 to 32 depending on GPU memory constraints, and tune the learning rate between $5\times 10^{-5}$ and $5\times 10^{-6}$. We train for 3–10 epochs based on model capacity and task complexity. Optimization is performed using the AdamW optimizer with a weight decay of 0.01. Evaluation is conducted by comparing the next predicted token with the next ground-truth target token.

## Appendix D Compute and Licensing Details

#### Computing Resources

All experiments were conducted on an NVIDIA RTX A6000 GPU with 48GB of memory.

#### Model Licenses

The GPT-2 Small [22] and GPT-2 Medium [22] models are released under the MIT License and are available via Hugging Face Transformers. The Qwen-2.5 0.5B [23] model is provided by Alibaba Group under the Apache 2.0 License. The LLaMA-3.1 8B [24] model is made available by Meta under a non-commercial research license and accessed via Hugging Face.

#### Dataset Licenses

CLUTRR [25] is released under a CC BY 4.0 license as part of EMNLP 2019. The ECQA [26] dataset is also released under a CC BY 4.0 license by its authors as part of EMNLP 2021. The Synthetic Arithmetic dataset is custom-designed by the authors and does not rely on any external or licensed data sources.

All assets were used in compliance with their respective licenses, and no proprietary or restricted resources were employed in our experiments.

## Appendix E Full Quantitative Results with Confidence Estimates

To provide a comprehensive view of model behavior across layers, we report the full experimental results for all models and datasets evaluated in this study. This includes predictive information trends ($I(Z;Y)$), incremental information gain ($I(\Delta Z;Y)$) and residual scaling effects ($\beta_{\ell}$), together with statistical significance reporting.

### E.1 Predictive Information

To estimate mutual information between hidden representations and target labels, we compute matrix-based mutual information following the methodology outlined in Appendix [A](https://arxiv.org/html/2507.05387v1#A1). Due to the high memory cost of computing $N\times N$ Gram matrices, we subsample between 50 and 100 test examples per experiment and report results with statistical significance based on multiple random seeds.

Figure: (a) GPT-2 Small on CLUTRR
Refer to caption: x13.png

### E.2 Incremental Information Gain

In addition to predictive information $I(Z^{(\ell)};Y)$, we compute the incremental information gain $I(\Delta Z^{(\ell)};Y)$ at each layer. This quantity the information contribution for each transformer blocks.

Figure: (a) GPT-2 Small on CLUTRR
Refer to caption: x23.png

### E.3 Residual Scaling

We introduce a residual scaling mechanism with learnable scalar coefficient parameters, inspired by prior work on adaptive residual modulation [12, 13]. Similar ideas of modulating internal computation have also been explored in parameter-efficient fine-tuning (PEFT) [27], Representation-Efficient Fine-Tuning (REFT) [28] and interpretability-driven control [29, 30, 31, 32]. We present the complete set of residual scaling results, detailing the learned $\beta_{\ell}$ values across all transformer layers. These values reflect the relative contribution of each layer after optimizing the residual scaling coefficients while keeping all other model parameters frozen. $\beta_{\ell}$ are trained on in-distribution (ID) and out-of-distribution (OOD) data. We observe that in the ID setting, later layers tend to receive higher weights, consistent with memorization behavior. In contrast, OOD training consistently downweights final layers and shifts importance toward the middle of the network, aligning with the generalization ridge identified through InfoRidge.

Figure: (a) GPT-2 Small on Synthetic Arithmetic
Refer to caption: x33.png

## Appendix F Incremental Information Gain Case Study

To further illustrate how intermediate layers contribute to generalization, we analyze the semantic content introduced by residual transitions at different depths. Specifically, we decode the residual transition $\delta z$ at each layer using the language modeling (LM) head, projecting the incremental representation back into token space.

This analysis allows us to inspect the linguistic shift introduced by each transformer block in isolation, and to assess whether the changes correspond to task-relevant predictions or superficial noises.

### F.1 Methodology

For a fixed input, we compute residual transitions $\delta z^{(\ell)}$ at each layer and pass them through the model’s final linear projection (LM head) followed by a softmax. We record the top predicted tokens and their probabilities, and it reflects the directional change applied by layer $\ell$.

### F.2 Observations

We decode residual transitions across layers and report the top shifted tokens by projecting $\Delta Z^{(\ell)}$ through the LM head. These shifts provide insight into how each layer modifies the model’s internal prediction trajectory.

Layer 9 → 10: The top shifted tokens include ĠNine, Ġ10, Ġ4, Ġ8, and 8, which are all numerically aligned with the target prediction space in the synthetic arithmetic task. This indicates that the model is beginning to refine task-relevant numerical features at this depth.

Layer 10 → 11: The shifted tokens become partially diluted, featuring punctuation and less informative symbols such as ,, Ġ, and ., alongside occasional task-relevant entries like Ġ12 and Ġ4. This indicates a transitional phase where the model continues to refine meaningful task-relevant features, yet begins to exhibit noise from frequent but semantically uninformative, such as punctuation, tokens.

Layer 11 → 12: At the final layer transition, the top 5 shifted tokens become largely uninterpretable, including Ġcanaver, soDeliveryDate, enegger, 76561, and ILA. Conversely, the most negatively shifted tokens—Ġ4, Ġ3, Ġ5, Ġ6, Ġ1—correspond to plausible numerical predictions that were actively suppressed. This supports the hypothesis that final layers may overwrite generalizable abstractions with memorized or noise signals.

These patterns are consistent with our mutual information analysis, which identifies intermediate layers as better semantically aligned with the prediction target—forming a ridge of generalization.

### F.3 Implication

These decoding results reinforce our interpretation of the generalization ridge: intermediate layers contribute the most semantically informative updates to the model’s representation. The residual transitions thus serve as a useful lens for understanding how and where semantic meaning is introduced during forward propagation.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: In the early exiting results shown in table 1, are different prediction heads trained over a frozen model, or is the final head simply applied to the selected layer? In the latter case, the mismatch between latent spaces in absence of a projection could lead to misleading results.
soundness: 3
strengths: 1.	Originality: The proposed adaptation of matrix-based mutual information to next-token prediction is an interesting idea to identify layer influence towards predictions.
2.	Quality: Authors conduct extensive experiments across model families & tasks, careful controls (depth ablation, overfitting, attention/decoder visualizations).
3.	Clarity: Hypotheses, metrics, and takeaways are explicitly stated; additional mathematical details are added to appendices.
4.	Significance: The proposed method based on mutual information can provide a principled lens to locate “where” generalization happens, immediately useful for early-exit, model pruning, or targeted fine-tuning strategies.
summary: This paper introduces InfoRidge, an information-theoretic framework to analyze how predictive information flows through transformer layers during natural language generation (NLG). Using matrix-based mutual information estimation, the authors track two key quantities: 1) mutual info between hidden states at layer $l$ and target token (predictive information) and 2) info added by the residual update at layer $l$ (incremental information gain). Across GPT-2, Qwen-2.5, and LLaMA models on CLUTRR, ECQA, and a synthetic arithmetic dataset, the authors consistently observe a non-monotonic “generalization ridge”: predictive info peaks in upper-middle layers then declines, while final layers appear to memorize. They corroborate this with learnable residual scaling coefficients, allowing the model to assign layer importances depending on the layer influence towards prediction. Authors find that under distribution shifts, models automatically down-weight final layers and up-weight ridge layers, providing further evidence to the existence of generalized intermediate representations. The ridge emerges only beyond a depth threshold, and attention/visualization analyses show ridge layers focus on task-relevant tokens whereas final layers revert to positional heuristics.
weaknesses: 1. **Scalability of MI estimation**: In Appendix C, the authors state that 50-200 examples are subsampled to avoid excessive memory usage from large Gram matrices. Provided authors experimented mainly with small language models, this raises doubts on the applicability of such a method to analyze larger state-of-the-art LLMs.
2.	**Lack of proper causal verification**: While the analysis shows that ridge layers correlate with better OOD accuracy, the use of residual scaling coefficient to test their importance in a task-directed fashion might lead to confusing results, since 1) $beta$ coefficients are trained after full fine-tuning with frozen weights; it is unknown whether jointly learning $beta$ from scratch preserves the ridge or changes training dynamics; and 2) subsequent layers depend from previous ones in their computation, creating an implicit incentive to upweight earlier layers to reduce information loss. A more convincing analysis could, for example, employ activation patching methods to localize the units that contribute most to model predictions in OOD settings, showing that these are mainly contained within ridge layers.
3.	**Findings are not particularly novel**: While the characterization of the middle-to-final layers as responsible for generalization is novel, there is ample literature showing that these layers perform best when used for probes aiming to extract semantic task-related information (see e.g. the early work "BERT Rediscover the Classical NLP Pipeline" by Tenney et al., showing a hierarchy of linguistic information within BERT peaking at around 2/3rds of the model). While the analyses conducted by the authors are undoubtedly of value in providing additional evidence for the semantic generality of such layers, this is unlikely to affect the current practices that already employ those layers for interpretability analysis on language models.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: 1) Is there a specific reason the authors use a Gaussian kernel for entropy estimation? Could they include an ablation with alternative kernels? I want to see the sensitivity of the results depending on the kernel choice.

2)  Although multiple model sizes are considered, the scaling trends are not clearly characterized. Can you clarify how layer importance evolves with width/depth. I want to see the evaluation on existing open source models, no training is needed.

3) It seems the models are fine-tuned from pretrained checkpoints, and the layer-wise mutual information “peakedness” seems to reflect pretrained biases. Do the claims about ridge layers, memorization, and generalization still hold when training from scratch?

4) Suggestion: Consider splitting Figure 2 into multiple panels/figures to improve readability.
soundness: 2
strengths: 1) The paper studies very important questions of how task-specific generalization abilities emerge in a transformer layer, which is largely not understood.
2) The paper is well written, easy to follow and the core points are well articulated.
summary: This paper investigates how Transformers operate internally and how task-specific generalization emerges during training and propagates across layers for natural language generation. The authors estimate mutual information between layer-wise representations and the outputs, and learn per-residual-block scaling coefficients to assess each layer’s contribution. Using these coefficients, they find that mid-layers become substantially more important under distribution shift, whereas later layers are more tied to memorization of in-distribution patterns.
weaknesses: 1) It’s unclear how “which layers are important” alone explains how generalization emerges. The study doesn’t identify the causal factors behind these abilities.

2) Recent work already reports that middle (and later) layers matter, and that final layers can be poor for OOD generalization (e.g., [1], which is cited but not discussed much). Can the authors discuss more how the new findings from these papers are different the existing ones to understand how generalization emerges and what are the practical usefulness of these findings.

3) The evaluation datasets are limited. It would be valuable to relate the identified “ridge” layers to varied task types—math, general reasoning, and multi-hop reasoning. Please add results on standard datasets such as MMLU, GSM8K, HotpotQA etc.

4) There is no accompanying theoretical analysis.

[1] Uselis, Arnas, and Seong Joon Oh. "Intermediate layer classifiers for ood generalization." arXiv preprint arXiv:2504.05461 (2025).

---

confidence: 4
contribution: 2
initial_rating: 2
presentation: 3
questions: n/a
soundness: 2
strengths: - The paper is a nice application of using information theory to identify layer importance
- The finding is self-contained and the paper fairly easily to follow
summary: This paper analyzes the mutual information between hidden layers and outputs for language models across a few tasks. The paper finds a spike in the upper middle layers which they coin generalization ridge.
weaknesses: At a high level, I was disappointed after reading the paper for few reasons. (1) I was missing the reason for why I should care about these results, especially since the results are not connected to the existing body of literature. There are many prior papers that investigate layer-wise importance and the result that upper middle layers are important is not new. (2) The key claims even in the title relate to shedding light into NLG. But, there is no actual experiment on NLG tasks and the experiments only look at a single output, not even multiple tokens. 

**Some related work**: Some discussions I was expecting to see were grounding in work like "BERT rediscovers the classical NLP pipeline" which relates layers to specific NLP tasks (plus an extensive body of literature around this). I was expecting a discussion of causal vs. correlational probing work, for example causal mediation analysis as an alternative way to identify layer/neuron importance. And since this was investigated, I was expecting more relation to work on training dynamics. 

Similarly, there is a lot of prior work on information-theoretic probing approaches that the method here does not warrant a new name. This is an implicit overclaim. 

**Lack of Relation to NLG**

The model selection is very surprising and limited to very small models. The older models here are not models that were ever used for actual NLG work. 
Moreover, none of the datasets are well known NLG datasets or classic NLG tasks (such as summarization, dialog response generation, or surface realization). The arithmetic dataset is not at all an NLG dataset. In addition, looking at a single output representation implicitly looks at only a single output and not a longer utterance as would be expected in NLG. 

**Relationship between token representation and output**
The extraction of $\Delta$Z (the change in token representation) implies that the information is represented in linear space. While this could be a helpful approximation of the ground truth, this is a very strong assumption to make, given the very much non-linear nature of the model. Moreover, page 4 states that delta Z is further $l2$ normalized which further changes the nature of the representation away from what the model represents

On a clarity note, I don’t understand why the Matrix-Based Mutual Information section does not continue to use the prior representation introduced above. This seems rather lazy as it is left to the reader to connect how it is applied. Not even the information in the appendix discusses details

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. For Table 1 accuracy results, how exactly is the early-exit implemented? I see that the language modeling head uses the same weights as for the input embedding, but still - maybe the results are influenced by a mismatch between the language modeling head and the hidden states of intermediate layers? AFAIU the embedding/unembedding matrix weights are only optimized for projecting the input and last layer output.
2. Is the approach for creating the synthetic dataset new? Or did other works use similar approaches?

Minor suggestions:
* I felt Figure 1 in its current form is not particularly easy to follow visually, and thus less informative for the reader. I suggest to think a bit about how to organize it visually and create a stronger connection between the text and the graphics.
* The transition to the part about "signal attention" (l. 343-345) can be improved, I found it a bit confusing. My suggestion is to replace that with the more explicit goal of this experiment (to quantify where in the network attention to semantically important information is concentrated, as currently stated in l. 350) and then describe the experiment.

Typos:

l. 125 reflects -> that reflects

l. 328/330 generalization ridge -> a/the generalization ridge

l. 401 blocks -> block

l. 921 quantity -> quantifies, blocks -> block
soundness: 4
strengths: 1. Some very interesting findings on model behavioral dynamics, that are consistent across models and tasks.
2. The conclusions are supported through multiple complementary experimental setups: the observational results are complemented by intervention experiments - controlling the contribution of each transformer layer via residual scaling coefficients; standard NLG data is used alongside a controlled synthetic dataset setting; mutual information results are complemented by an explicit analysis of attention to signal-bearing tokens.
summary: This work analyzes the representations across transformer layers in a scenario of fine-tuning for natural language generation tasks. The authors propose an approach based on mutual information between the internal model representations and the ground truth target token, and formalize metrics of "predictive information" and "incremental information gain", corresponding to a specific layer hidden state or a specific layer residual, respectively. These mutual information metrics rely on matrix-based Renyi entropy. In a series of analysis experiments the work demonstrates a behavior termed "generalization ridge", where the mutual information with the task label reaches a peak before the final transformer layer, and this point in the model also corresponds to better out-of-domain generalization performance.
weaknesses: 1. My feeling is that the interpretation of the data is probably a bit over-simplified, given that high mutual information between a representation and the target label can also be potentially consistent with memorization. Throughout the paper a rise in predictive information is interpreted as increased generalization, whereas in the overfitting experiment (l. 332) a similar rise is described as "memorizing shortcuts" and "redundant label-specific noise". This is interesting, and helps to give the reader a more complete picture, but also highlights that the meaning of the predictive information metric may not be as straightforward as the authors make it out to be.
2. The analysis is limited to fine-tuning scenarios. I think it would have been informative to also see some examples of these MI metrics in a zero-shot/few-shot scenario.


Decision: Reject