From Per-Image Low-Rank to Encoding Mismatch: Rethinking Feature Distillation in Vision Transformers

Conference: ICML 2026
arXiv: 2511.15572
Code: Available (paper supplementary)
Area: Model Compression / Knowledge Distillation / Vision Transformer
Keywords: Feature Distillation, ViT Compression, Subspace Rotation, Spectral Energy Pattern, Latent Capacity Bottleneck

TL;DR

The authors use a three-perspective analysis—sample-wise SVD, dataset-level PCA, and token-level Spectral Energy Pattern (SEP)—to reveal a seemingly paradoxical geometry in ViT representations: "Each image's feature matrix is low-rank, but the cross-image shared subspace is nearly full-rank, and the spectral bandwidth of single tokens approaches 100%." They then propose two minimal patches, Lift (retaining the lifting projector at inference) and WideLast (widening only the last block to teacher width), which enable plain MSE feature distillation to boost DeiT-Tiny ← CaiT-S24 from 74.86% to 78.23%.

Background & Motivation

Background: In knowledge distillation (KD), "matching intermediate features" is a classic approach from the CNN era (FitNet/AT/KR). For ViT, "same-size representation transfer" (e.g., CLIP distillation) generally works. However, when compressing from a wide teacher to a narrow student, direct feature alignment becomes surprisingly fragile—often yielding only marginal improvements or even degradation; works like ViTKD/SpectralKD/VkD have reported similar phenomena.

Limitations of Prior Work: Existing solutions either switch to distillation tokens (DeiT approach), use contrastive/attention/manifold losses, or introduce complex translator modules. While effective, these methods sidestep the core question of "why naive feature KD fails," leaving ViT distillation without a simple, interpretable narrative.

Key Challenge: The authors observe a puzzling paradox. Sample-wise SVD shows that ViT features for each image are highly compressible—for CaiT-S24, the last-layer token matrix (196×384) requires only 61 singular directions to retain 95% energy for 99% of images. By the Eckart-Young-Mirsky theorem, this suggests that "a narrow student plus a linear projector" should theoretically match the teacher. But in practice, this does not work.

Goal: To clarify why "theoretically feasible" does not translate to "practically feasible," and to provide a minimal-cost remedy.

Key Insight: The authors suspect the sample-wise perspective overlooks a key point—each image's low-dimensional subspace is different, and each token actually uses high bandwidth within its own subspace. Thus, they introduce two complementary diagnostics: dataset-level PCA (how wide the shared subspace needs to be) and token-level SEP (how many spectral channels a single token occupies).

Core Idea: The failure mechanism is endpoint encoding mismatch—there is both "subspace rotation" (a fixed projector cannot adapt to input-dependent subspace directions) and "bandwidth capacity insufficiency" (the narrow student cannot support the high spectral occupancy of token encodings), and these two aspects are coupled. The remedy is to provide the student with "endpoint capacity close to the teacher's width plus input-dependent subspace adjustment ability."

Method

Overall Architecture

The paper is divided into two halves. The first half (§2) is diagnosis: using three complementary spectral/geometric tools to characterize the intrinsic structure of the ViT last-layer token feature matrix \(\mathbf{X} \in \mathbb{R}^{N \times D}\), comparing sample-wise SVD, dataset-level PCA, and token-level SEP, leading to the "encoding mismatch" diagnosis. The second half (§3) is remedy: based on the diagnosis, proposing two minimal patches, Lift and WideLast. Lift uses a fixed linear projector at inference to raise the student width to the teacher's width; WideLast widens only the student's last Transformer block to the teacher's width.

Key Designs

1. Three-Perspective Representation Geometry Diagnosis:

   • Function: Reveals the true contradiction between "low-rank compressibility" and "narrow student matching failure."
   • Mechanism: (i) Sample-wise SVD: For each image, compute \(\mathbf{X}_i = \mathbf{U}_i \boldsymbol{\Sigma}_i \mathbf{V}_i^\top\), and record the minimal rank \(d_i^\text{SVD}\) needed for 95%/99% energy; for CaiT-S24, the 99th percentile requires only 61/121 dimensions. (ii) Dataset-level PCA: Aggregate the channel second moment over the dataset \(\mathbf{C} = \frac{1}{T}\sum_i \mathbf{X}_i^\top \mathbf{X}_i\) and decompose to obtain a shared PCA basis \(\mathbf{V}_d\); then, for each image, compute the energy retention \(E_i(d) = \|\mathbf{X}_i \mathbf{V}_d\|_F^2 / \|\mathbf{X}_i\|_F^2\). To retain 95% energy for 99% of images requires 302/384 dimensions (CaiT-S24)—a 5× gap compared to sample-wise SVD's 61. (iii) Token-level SEP: For each token \(\mathbf{x}_t \in \mathbb{R}^D\), perform 1D DFT along the channel dimension, compute cumulative spectral energy \(\text{SEP}(d)\) and normalized bandwidth \(b_\alpha\); across 14 different ViT/DeiT/Swin/CLIP/MAE/DINO/DINOv2 backbones, SEP curves almost all fall on the 45° diagonal, requiring ~90% of spectral channels to capture 90% energy.
   • Design Motivation: Looking only at SVD leads to the mistaken conclusion that "a narrow student plus linear projector suffices"; only by combining PCA (revealing input-dependent subspace rotation) and SEP (revealing high single-token bandwidth usage) can one fully explain why a fixed narrow interface fails. These three perspectives are complementary and indispensable.

2. Lift: Inference-Time Retained Lifting Projector:

   • Function: Provides the student with "endpoint capacity" at the teacher's width without changing the backbone architecture.
   • Mechanism: The student's last-layer output \(\mathbf{X}_S \in \mathbb{R}^{N \times D_S}\) (narrow width \(D_S < D_T\)) is passed through a token-wise linear projector \(\mathbf{P} \in \mathbb{R}^{D_S \times D_T}\) to obtain \(\widehat{\mathbf{X}}_S = \mathbf{X}_S \mathbf{P}\). Crucially, this projector is retained at inference (unlike traditional KD, where the projector is discarded after training), so the classification head \(\mathbf{W}_\text{head} \in \mathbb{R}^{D_T \times C}\) operates directly on the lifted representation.
   • Design Motivation: Ablation shows that once the projector provides a teacher-width interface, even the simplest MSE feature alignment improves from 0.21% gain to +1.75% gain. This confirms the SEP diagnosis—the main failure is insufficient endpoint bandwidth, and once width is restored, feature matching is immediately effective. However, Lift is a fixed linear mapping and cannot address the subspace rotation problem revealed by PCA, so its effect is inferior to WideLast.

3. WideLast: Native Width Alignment (Widening Only the Last Block):

   • Function: Simultaneously addresses endpoint bandwidth and subspace rotation.
   • Mechanism: Replace the student's last Transformer block with a teacher-width \(D_T\) version (all previous blocks remain narrow \(D_S\)). Thus, the last block's attention and MLP operate in \(D_T\) dimensions, outputting \(\widetilde{\mathbf{X}}_S \in \mathbb{R}^{N \times D_T}\), and the classification head also operates on \(D_T\).
   • Design Motivation: The key difference from Lift—widening the block is an input-dependent nonlinear mapping, enabling different effective subspace directions for different images, directly corresponding to the "subspace rotation" phenomenon revealed by PCA; a fixed projector can only provide the same rotation for all images. Ablation shows WideLast (78.23%) outperforms Lift (77.53%) by 0.7 points, validating the extra benefit of subspace adaptivity.

Loss & Training

The overall objective is \(\mathcal{L} = (1-\lambda_\text{logit}) \mathcal{L}_\text{CE}(\mathbf{y}, \mathbf{p}_S) + \lambda_\text{logit} \mathcal{L}_\text{KD}(\mathbf{p}_S, \mathbf{p}_T; \tau) + \lambda_\text{feat} \mathcal{L}_\text{feat}\), where \(\mathcal{L}_\text{feat}\) can be simple MSE or SpectralKD. The training recipe strictly follows DeiT defaults (AdamW, 5e-4 lr, cosine, 5 epoch warmup, 300 epochs, batch size 2048).

Key Experimental Results

Main Results

ImageNet-1K, CaiT-S24 (384-dim) → DeiT-Tiny (192-dim) distillation:

Configuration	Distillation Loss	Top-1 (%)	Δ
DeiT-Tiny baseline	–	74.86	–
Baseline + SpectralKD (naive)	–	75.07	+0.21
Lift + MSE only	MSE	76.61	+1.75
Lift + SoftKD + SpecKD	SoftKD+SpecKD	77.53	+2.67
WideLast + MSE only	MSE	77.15	+2.29
WideLast + SoftKD + MSE	SoftKD+MSE	78.23	+3.37

Naive feature KD is almost ineffective (+0.21), but with Lift, MSE alone yields +1.75, and with WideLast, +2.29, confirming that "endpoint capacity" is the key. The strongest combination, WideLast + SoftKD + MSE, achieves 78.23%, 3.37 points above baseline.

Ablation Study

Configuration	Top-1 (%)	Notes
Default baseline (192-dim)	74.86	No projector
Projector 256	75.46	Slightly wider
Projector 320	75.53	Near teacher
Projector 384 (= teacher)	75.41	Retained at inference + no KD
Projector 448 (exceeds teacher)	75.23	Wider actually hurts
Lift standalone (no KD)	75.41	+0.55 vs baseline
WideLast standalone (no KD)	75.54	+0.68 vs baseline
Teacher → DeiT-Small	–	WideLast + SpecKD 75.73
Teacher → DeiT3-Small-21k	–	WideLast + SpecKD 76.50

Key Findings

• Exceeding teacher width actually hurts (448 vs 384): The goal is not "the wider the better," but to precisely align with the teacher's subspace; extra dimensions introduce redundancy and harm learning.
• Architecture modification alone improves performance (no teacher needed): Lift +0.55, WideLast +0.68, confirming that encoding mismatch is not just a "distillation interface problem," but a latent capacity bottleneck in architectures like DeiT-Tiny—insufficient last-layer bandwidth limits expressiveness even in independent training.
• Pattern holds across three teachers: Lift/WideLast consistently improve performance on three diverse teacher architectures (CaiT-S24 / DeiT-Small / DeiT3-Small-21k), indicating encoding mismatch is a general ViT-family property, not a CaiT-specific quirk.
• SEP consistency across architectures ~100%: SEP curves for 14 backbones (including SL/SSL/MM training, Tiny to Huge scales) almost coincide on the diagonal; capturing 90% energy requires ~90% spectral channels—this "spectral universality" is one of the paper's most striking findings.

Highlights & Insights

• The "sample-wise low-rank + dataset-wise high-dimensional + token-wise high-bandwidth" three-perspective framework is a reusable diagnostic: It decomposes the overused term "redundancy" into three distinct geometric concepts, avoiding the confusion of "low-rank ≠ distillable" that conflates global and local redundancy.
• SEP as a new diagnostic tool: Prior literature on ViT feature geometry mainly relied on attention map visualization or CKA comparison; the authors introduce 1D DFT along the channel dimension to measure spectral occupancy, providing a simple "single-token capacity" probe potentially useful for other transformer distillation/compression tasks.
• "Endpoint bandwidth insufficiency is an architectural bottleneck, not just a distillation bottleneck" is unexpectedly important: Traditionally, width is seen as a trade-off between parameter efficiency and expressiveness; this work argues that for ViT architectures concentrating all information in the last-layer tokens, last-layer bandwidth is an independent bottleneck, directly informing future compact ViT designs (e.g., mobile ViT)—keeping intermediate layers narrow but widening the last layer is a new design point.
• Minimal patch approach: Compared to approaches like ScaleKD that build complex translator modules, Lift/WideLast add minimal parameters yet achieve significant gains, validating that "understanding the core problem > stacking modules."

Limitations & Future Work

• Only last-layer encoding mismatch is analyzed; whether similar phenomena exist in intermediate layers and how to address them is not discussed, which is a clear limitation for multi-layer feature alignment.
• The comparison between fixed projector vs input-dependent projector is clear, but "adaptive lifting" (e.g., attention-routed projectors) is not further explored, which might narrow the gap with WideLast.
• All experiments are on ImageNet-1K classification; detection/segmentation/multimodal downstream tasks are not covered. How encoding mismatch behaves in dense prediction scenarios remains an open question.
• The widened last block in WideLast incurs significant parameter overhead (attention + MLP at \(D_T\) width), making it a trade-off rather than a pure win for true edge deployment.

Related Work & Insights

• vs FitNet/AT/KR: Classic feature KD for CNNs, but all assume similar teacher-student layer widths; this work specifically addresses the "wide-to-narrow" interface problem, adapting the lineage to ViT.
• vs ScaleKD: The latter uses a complex alignment module to bridge heterogeneous teachers; Lift/WideLast are lighter-weight versions, but ScaleKD can handle more extreme heterogeneity.
• vs VkD: Uses orthogonal projectors for stable distillation; related in approach, but this work provides a deeper explanation from the representation geometry perspective.
• vs SpectralKD: Also uses a spectral perspective but focuses on layer selection and spectral alignment loss; this work's SEP further reveals single-token spectral occupancy, and the two can be combined (WideLast + SpecKD is among the strongest combinations in experiments).
• vs Yu & Wu (low-rank features but weights not low-rank): Both observe low-rank ViT features, but Yu & Wu apply this directly to few-shot compression; this work further uses PCA + SEP to reveal the paradox of "low-rank but hard to distill" and provides a remedy.

Rating

• Novelty: ⭐⭐⭐⭐⭐ The three-perspective diagnostic framework, encoding mismatch concept, and SEP tool are all new; the narrative and approach are fresh.
• Experimental Thoroughness: ⭐⭐⭐⭐ Thorough on ImageNet, and the SEP universality experiment across 14 backbones is rigorous; but only classification and a single backbone (DeiT-Tiny) as student are evaluated, with limited downstream task support.
• Writing Quality: ⭐⭐⭐⭐⭐ The logical flow from paradox introduction → three diagnostics → dual remedies → experimental validation is very smooth, and Figure 1's four-panel illustration is especially clear.
• Value: ⭐⭐⭐⭐ Provides a clear explanation and simple remedy for the long-standing ViT distillation failure, and the standalone gain makes the WideLast approach instructive for compact ViT architecture design.

Related Papers

• [ICCV 2025] EA-ViT: Efficient Adaptation for Elastic Vision Transformer
• [ICML 2025] Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
• [ICCV 2025] Gain-MLP: Improving HDR Gain Map Encoding via a Lightweight MLP
• [ACL 2025] Revisiting LoRA through the Lens of Parameter Redundancy: Spectral Encoding Helps
• [ICCV 2025] ViT-Linearizer: Distilling Quadratic Knowledge into Linear-Time Vision Models