🤖 machine learning

Uncovering the Latent Potential of Deep Intermediate Representations

This paper introduces Layer-wise Optimal Embedding Selection (LOES) and Geometric Regularization Loss (GeoReg) to demonstrate that task-relevant information in deep foundational models is non-monotonically distributed across layers, showing that explicitly identifying and geometrically aligning task-discriminative intermediate representations significantly outperforms standard transfer learning approaches.

Original authors: Arnesh Batra, Arush Gumber, Aniket Khandelwal, Jashn Khemani, Anubha Gupta

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Arnesh Batra, Arush Gumber, Aniket Khandelwal, Jashn Khemani, Anubha Gupta

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, multi-story library (a "Foundational Model") that has read almost everything in the world. When you ask it a question, it doesn't just give you one answer; it processes your request through many different floors, or "layers," of its brain.

For a long time, scientists thought the top floor (the final layer) was the only place where the smart, useful answers lived. They assumed the information got better and better as it climbed the stairs, so they only ever looked at the very top.

This paper argues that this assumption is wrong. It's like assuming you only need the view from the penthouse to understand a city, ignoring the bustling markets on the ground floor or the quiet libraries on the middle floors. Sometimes, the most useful information for a specific task is hidden in the middle, or scattered across several different floors.

Here is how the authors solved this problem, using simple analogies:

1. The Problem: The "Penthouse" Trap

The authors found that if you only use the final layer, you often miss out on crucial details.

The Analogy: Imagine trying to identify a specific type of bird. The top floor of the library might only know "It's a bird." But the 5th floor might know "It has a red beak," and the 3rd floor might know "It sings at dawn." If you only look at the top, you miss the clues that actually help you identify the bird.
The Reality: In many AI models, the final layer gets too specialized for its original training (like predicting the next word in a sentence) and forgets the specific details needed for new tasks.

2. The Solution: LOES (The Smart Librarian)

The authors created a new method called LOES (Layer-wise Optimal Embedding Selection). Think of LOES as a super-smart librarian who doesn't just go to the top floor. Instead, she walks through the whole building to find the exact mix of floors that hold the clues you need.

How it works:
- The Search: The librarian looks at every floor. She doesn't just pick the "best" one; she picks a combination of floors that work well together without repeating the same information.
- The Geometry Check: She checks if the information on a floor is "well-organized." She avoids floors where the information is messy or crammed into a tiny corner (which she calls "anisotropy"). She prefers floors where the information is spread out evenly (which she calls "isotropy"), like a neatly arranged bookshelf rather than a pile of books.
- The Result: She fuses these selected floors together to create a custom "super-layer" specifically for your task.

3. The Safety Net: GeoReg (The Stabilizer)

Once the librarian picks the right floors, there's a risk that when you start training the AI on a new task, the building might get shaken up, and the information might collapse back into a messy pile.

The Analogy: Imagine you build a tower out of blocks using the best blocks from different floors. If you start shaking the table (training the model), the tower might fall.
The Fix: The authors added a second tool called GeoReg. This acts like a glue or a stabilizer. It gently holds the blocks in place, ensuring the tower stays tall and organized while you are building your new structure. It prevents the "smart" information from collapsing into a useless mess.

4. What They Found

The paper tested this on many different types of AI (for seeing images, reading text, and listening to speech) and found:

Depth Matters: The deeper the library (the more layers the model has), the more important it is to use this "multi-floor" approach. The gains get bigger as the models get bigger.
Training Style Changes the Map:
- If a model was trained on a huge, diverse mix of data (like CLIP, which saw millions of images and texts), the useful clues were spread out evenly across the middle floors.
- If a model was trained on a narrow set of data (like only pictures of cats and dogs), the useful clues were mostly stuck at the very top.
Language Matters: For languages that are rare or underrepresented in training data, this method helped the most. It was like giving a translator a dictionary that included not just the final translation, but also the grammar rules and cultural context from the middle floors, which are often lost in the final output.

Summary

The paper claims that AI models are like deep libraries where the best answers are often scattered across many floors, not just the top one.

By using LOES, we can automatically find the right mix of floors to build a better, more accurate AI for specific jobs. By using GeoReg, we make sure that this new, custom-built AI stays stable and doesn't fall apart while we teach it new things. This makes AI smarter, more efficient, and easier to understand without needing to rebuild the whole library from scratch.

Technical Summary: Uncovering the Latent Potential of Deep Intermediate Representations

Problem Statement

Foundational models, pretrained on massive datasets, learn hierarchical representations that evolve across network depth. The prevailing assumption in transfer learning is that semantic utility increases monotonically with depth, leading to the standard practice of utilizing only the final layer output for downstream tasks. However, recent empirical analyses challenge this view, demonstrating that intermediate layers often outperform final representations by up to 10% in downstream tasks. This phenomenon arises because final layers may over-specialize to pretraining objectives (e.g., next-token prediction) or suffer from representation collapse. Furthermore, task-relevant information is distributed non-monotonically across layers, meaning that naive aggregation or simple weighting of layers fails to recover the optimal task-specific subspace. The core problem addressed is how to systematically identify and fuse the specific subset of intermediate layers that encode task-discriminative structure without incurring prohibitive computational costs.

Methodology

The authors propose a unified framework comprising two coupled components: Layer-wise Optimal Embedding Selection (LOES) and Geometric Regularization (GeoReg).

1. Layer-wise Optimal Embedding Selection (LOES)

LOES is a constructive spectral framework designed to approximate the task-optimal subspace by iteratively selecting non-redundant layers. It formulates layer selection as a ridge-residual optimization problem constrained by orthogonality and isotropy. The algorithm proceeds in two stages:

Initialization: The first layer is selected by fitting a ridge probe against original targets and minimizing a score combining fit loss with an isotropy term.
Iterative Selection: Given a current set of selected layers $S$ , candidate layers are orthogonalized against the span of $S$ to isolate marginal explanatory power. Each candidate is scored using a composite objective (Eq. 7):
$\text{Score}(\ell) = \text{Loss}_\ell + \alpha(1 - \text{Iso}(X_\ell)) + \gamma \text{Red}_\ell - \eta \text{Tri}(\tilde{X}_\ell)$
Where:
- $\text{Loss}_\ell$ : Mean squared error of a ridge probe on the orthogonalized residual.
- $\text{Iso}(X_\ell)$ : An isotropy score derived from the eigenvalues of the empirical covariance, promoting well-conditioned embeddings.
- $\text{Red}_\ell$ : A redundancy term measuring the global alignment of the candidate with already selected layers.
- $\text{Tri}(\tilde{X}_\ell)$ : An optional term for classification estimating the average triangle area of class centroids to encourage high-volume simplex structures.
The layer minimizing this score is selected, refit against original targets, and added to the ensemble. This process repeats until a budget $K$ is reached.

2. Geometric Regularization (GeoReg)

To prevent the degradation of favorable geometric properties (such as isotropy and class separation) during fine-tuning, the authors introduce GeoReg. This auxiliary loss is applied during downstream adaptation to the fused multi-layer embeddings. It penalizes spectral imbalance in the embedding covariance (encouraging uniform eigenvalue distribution) and penalizes configurations where class centroids span low volumes (indicating collapse). Unlike pretraining regularizers (e.g., VICReg), GeoReg operates specifically during the transfer phase to preserve the geometry identified by LOES against competing task gradients.

Key Contributions

Empirical Demonstration: The paper demonstrates that optimal representations for downstream tasks are rarely the final layer alone but rather a task-specific subspace formed by a subset of intermediate layers.
LOES Algorithm: A spectral algorithm that identifies optimal layer combinations by minimizing residual error under geometric constraints (isotropy and orthogonality). It outperforms learnable weighting baselines by explicitly constructing subspaces with high effective rank and low anisotropy.
Interpretability Lens: By analyzing layer selection preferences, the method reveals where different foundational models encode specific semantic or structural attributes, showing that pretraining paradigms (e.g., contrastive vs. masked reconstruction) systematically influence information distribution.
GeoReg: A loss function that improves stability in fine-tuning settings, mitigating feature collapse often observed in large-scale models.
Scalability: The method shows performance gains that scale with model depth and are modality-agnostic, working across vision, language, speech, and multimodal architectures.

Results

The authors evaluated LOES across a diverse set of architectures (DINOv2, CLIP, BERT, ModernBERT, Wav2Vec 2.0) and datasets (ImageNet, MTEB benchmarks, ASVspoof, etc.).

Performance Gains: LOES consistently outperforms standard last-layer transfer and learnable weighting baselines. Gains are particularly pronounced in deeper models and diverse pretraining regimes. For instance, on ModernBERT, LOES improved accuracy from 78.07% to 94.48% on the MTOP task, whereas BERT-large saw a modest gain from 95.32% to 97.81%.
Pretraining Influence: Models pretrained on large, diverse corpora (e.g., CLIP, DINOv2) tend to distribute task-relevant information in early-to-mid layers, while models trained on narrower distributions (e.g., MAE, ViT trained on ImageNet) concentrate discriminative features in later layers.
Cross-Lingual and Cross-Modal: LOES provides significant benefits for underrepresented languages (e.g., +10.9% on Urdu) and adapts effectively to speech tasks (e.g., spoofing detection vs. emotion recognition requiring different layer depths).
Efficiency: The method introduces minimal computational overhead (typically <1% increase in wall-clock time) and parameter count (<1% increase), as it relies on closed-form ridge regression on a small calibration set.
Comparison: LOES outperforms exhaustive search over layer subsets (recovering near-optimal solutions at a fraction of the cost) and greedy selection baselines.

Significance and Claims

The paper claims that layerwise embedding geometry is not incidental but central to how deep models represent and transfer knowledge. By moving beyond the assumption of monotonic semantic utility, the work provides strong evidence that effective transfer depends on identifying which layers encode task-discriminative structure and how their embeddings are geometrically organized.

The authors position LOES not merely as a performance booster but as a tool for interpretability, offering a structured view of how semantic factors are distributed across layers. They claim that as foundation models scale and diversify, principled layer selection becomes increasingly valuable, particularly for low-resource and underrepresented domains where final layers alone inadequately capture task-relevant information. The work suggests that the geometry of intermediate representations is a critical, yet underutilized, resource for efficient and transparent model adaptation.