Research map / Architectures and efficiency / RNNs and normalization
Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that sample-wise cosine similarity between transformer layers tracks CKA, explains saturation events via a geodesic-curve theory, and proposes aligned training (shared classifier across layers) that yields early-saturating, single-classifier multi-exit transformers.
Problem
Understanding how individual inputs' representations progress across layers in transformers, and why shallow layers can already produce correct predictions (saturation events), while enabling efficient early exit without many per-layer classifiers.
Approach
Measure layer-wise similarity per sample with cosine similarity (not CKA), justify monotonic similarity increase under a geodesic-curve assumption for residual networks, and train with an aligned loss that applies one shared last-layer classifier to all layers with linearly increasing weights, enabling a single-classifier multi-exit model.
Key results
- On DeiT-S with ImageNet, the aligned multi-exit model reaches 77.96% accuracy with 1.36x speedup, comparable to multi-exit/classifiers (78.32%, 1.42x) and standard training (80.28%).
- Aligned training yields an epsilon-effective depth of about 7 on CIFAR-10 versus 12 for standard training (12-layer DeiT-S).
- Aligned training with a single shared classifier reduces parameters by 78.39% for GPT-2-sized multi-exit models compared with per-layer classifiers (per the parameter table).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply aligned training with a shared classifier and early exits | DeiT · ImageNet | accuracy: 80.28% -> 77.96% | Hurt |
| Apply aligned training with a shared classifier and early exits | DeiT · ImageNet | speedup: 1.36 × | Helped |
| Use a single shared classifier instead of separate classifiers at each layer | DeiT-S · Multi-exit model | parameters: 26.27M -> 22.05M; 16.07% | Helped |
| Use a single shared classifier instead of separate classifiers at each layer | DeiT-B · Multi-exit model | parameters: 95.02M -> 86.57M; 8.89% | Helped |
| Use a single shared classifier instead of separate classifiers at each layer | GPT-2 · Multi-exit model | parameters: 541.57M -> 117.35M; 78.39% | Helped |
| Use a single shared classifier instead of separate classifiers at each layer | GPT-3 · Multi-exit model | parameters: 233.67B -> 175.63B; 25.10% | Helped |
| Use a single shared classifier instead of separate classifiers at each layer | LLaMA-2 · Multi-exit model | parameters: 75.11B -> 70.35B; 6.81% | Helped |
| Replace linearly increasing layer weights with uniform weights (ablation) | Aligned training · ImageNet-1K | layer-wise accuracy | Mixed |
| Replace aligned loss with CE-reg cosine-similarity regularization (ablation) | Aligned training · Vision experiments | layer-wise accuracy | Hurt |
| Add cosine similarity as a regularization term to the final-layer CE loss (ablation) | Vision models · Vision experiments | layer-wise accuracy | No clear effect |
| Add aligned CE losses at intermediate decoder layers (ablation) | DETR · Object detection | AP50 | Helped |
Surprising findings
- Sample-wise cosine similarity between layers is nonnegative on average, while rare samples show negative cosine similarity between distant layers.
- Training only with the aligned loss degrades last-layer performance, so an alternating training schedule is used.
Limitations
- Theory relies on a geodesic curve assumption that may not hold exactly in practice, so some samples may not show strictly monotonic similarity.
- Aligned loss alone hurts final-layer accuracy, requiring an alternating training schedule.
- Method requires all layers to share the same feature dimension, so it does not directly extend to varied-dimension architectures.
Experimental setup
Datasets: CIFAR-10 ImageNet-1K GLUE Wikitext-103 CIFAR-10.2
Benchmarks: CIFAR-10 ImageNet-1K GLUE Wikitext-103 CIFAR-10.2
Metrics: layer-wise accuracy CKA cosine similarity saturation events perplexity coherence diversity speedup
Baselines: standard training multi-exit/classifiers CE-reg (cosine regularization) CKA
Builds on: Centered Kernel Alignment neural collapse logit lens multi-exit networks DeiT
Closest papers by meaning
- How Not to Stitch Representations to Measure Similarity: Task Loss Matching Versus Direct Matching (AAAI 2025 · similarity 0.88)
- When Representations Align: Universality in Representation Learning Dynamics (ICML 2024 · similarity 0.88)
- When Does Closeness in Distribution Imply Representational Similarity? An Identifiability Perspective (NeurIPS 2025 · similarity 0.88)
- Representational Similarity via Interpretable Visual Concepts (ICLR 2025 · similarity 0.88)
- ReSi: A Comprehensive Benchmark for Representational Similarity Measures (ICLR 2025 · similarity 0.87)
- Residual Stream Analysis with Multi-Layer SAEs (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.