SpectralKD for ViT Distillation

• The paper introduces SpectralKD, a framework that leverages Fourier-domain feature analysis with a frequency-alignment loss for effective knowledge distillation in ViTs.
• It utilizes channel-wise FFT and layer intensity curves to identify key 'hot' layers, guiding the optimal alignment between teacher and student models.
• Empirical results on ImageNet-1K demonstrate significant accuracy gains over baselines, effectively transferring both low- and high-frequency information.

SpectralKD is a unified analytical framework for interpreting and distilling Vision Transformers (ViTs) via spectral analysis. By integrating model-wise and layer-wise Fourier-domain feature analyses with a principled frequency-alignment loss, SpectralKD provides both interpretability for ViT architectures and a state-of-the-art recipe for knowledge distillation (KD). The framework formalizes the information complexity of feature maps, guides optimal layer selection for distillation, and demonstrates that matching teacher–student spectral “fingerprints” yields significant gains, as validated on ImageNet-1K benchmarks (Tian et al., 2024).

1. Unified Spectral–Distillation Framework

SpectralKD combines spectral analysis of ViT intermediate features with a distillation mechanism that leverages their frequency-domain characteristics. The approach comprises two central components:

• Spectral analysis of feature maps: Identifies “hot” layers—those with maximal averaged frequency intensity—and characterizes which frequency components encode the majority of task-relevant information in a pretrained teacher model.
• Frequency-alignment loss: Augments standard logits-based KD with an L2 loss directly in the 2D-Fourier space, encouraging the student to replicate the teacher’s spectral signatures at selected layers.

SpectralKD’s theoretical underpinning interprets the spectral magnitude (via Parseval’s theorem and information-theoretic rationale) of a feature map as a quantification of its information complexity. This informs both “where” (which layers) and “how” (which frequencies) alignment should be enforced for effective distillation. By spectrally aligning the student to the teacher, both global (low-frequency) and fine-grained (high-frequency) information are transferred.

2. Spectral Decomposition and Feature Analysis

SpectralKD characterizes the spectral content in ViT feature maps using precise mathematical procedure:

Given a feature tensor $X\in\mathbb{R}^{B\times C\times H\times W}$:

1. Channel-wise 1D FFT:

$$F(X) = \mathrm{FFT}_{\text{chan}}(X) \in \mathbb{C}^{B \times C \times H \times W}$$

Each spatial position $(h, w)$ is transformed along the channel dimension.

2. Magnitude Spectrum:

$$A(X) = |F(X)|^2 = \Re(F(X))^2 + \Im(F(X))^2 \in \mathbb{R}_{\ge 0}^{B \times C \times H \times W}$$

3. Channel-wise Average Spectrum:

$$S(X)_c = \frac{1}{B H W} \sum_{b=1}^B \sum_{h=1}^H \sum_{w=1}^W A(X)_{b, c, h, w} \in \mathbb{R}^C$$

4. Scalar Layer Intensity:

$$\ell(X) = \frac{1}{C} \sum_{c=1}^C S(X)_c \in \mathbb{R}_{\ge 0}$$

Aggregating $\ell(X^{(i)})$ across layers yields the “layer intensity curve,” crucial for interpreting information flow and guiding which layers are matched during distillation.

3. Frequency-Alignment Loss and Objective Formulation

SpectralKD aligns selected teacher–student layer pairs $(i_k \leftrightarrow j_k)$ in the Fourier domain:

• Channel Alignment: Channels are matched by adaptive average pooling:

$C = \min(C_t, C_s)$

with $F'_t$, $F'_s$ denoting channel-aligned features.

• 2D Real FFT Calculation:

$\widetilde F = \mathrm{RFFT2}(F')$

Real and imaginary parts are stacked.

• Spectral Alignment Loss:

$$\mathcal{L}_{\mathrm{FFT}} = \frac{1}{B C H (W/2+1) 2} \| F_{\mathrm{stack}}(F'_s) - F_{\mathrm{stack}}(F'_t) \|^2_2$$

• Total Objective:

$$\mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{KD}} + \beta \mathcal{L}_{\mathrm{FFT}}$$

where $\mathcal{L}_{\mathrm{KD}}$ is a weighted combination of cross-entropy and KL divergence on logits with hyper-parameters $(\alpha, T, \beta)$.

This composite loss enables SpectralKD to transfer both semantic alignments (logits) and structural feature complexity (spectra).

4. Model-Wise and Layer-Wise Spectral Insights

The application of spectral analysis across layered ViT models such as CaiT-S24 yields a characteristic U-shaped “layer intensity curve”:

• Early layers (1–2): Exhibit high spectral intensity, corresponding to localized, detailed edge features.
• Middle layers (3–20): Markedly reduced intensity, capturing abstracted, transformed representations.
• Final layers (21–24): Intensity increases again, indicative of semantic aggregation and fine detail reintegration.

Layer-wise spectra $S(X)$ manifest a transition: flat (all frequencies) in early and late layers, and a steep low-to-high frequency decay in the middle, reflecting shifts in information encoding throughout the network. Swin-Small, a hierarchical transformer, demonstrates a virtually identical evolution in spectral encoding patterns across its stages, despite architectural differences. This cross-architecture convergence underpins SpectralKD’s generic alignment strategy.

5. Implementation Protocol

The SpectralKD distillation protocol can be summarized as follows:

Inputs: - Teacher T, Student S - Selected layer pairs (i_k in T ↔ j_k in S, for k=1…m) - hyper‐params α, T, β for each minibatch {x,y}: # (A) Forward pass f_t, {F_t^{(i_k)}} ← T(x) f_s, {F_s^{(j_k)}} ← S(x) # (B) Logits distillation loss L_CE ← CrossEntropy(f_s, y) L_KL ← KL( softmax(f_t/T) || softmax(f_s/T) ) L_KD ← (1–α)·L_CE + α·T^2·L_KL # (C) Spectral alignment loss L_FFT ← 0 for k=1…m: F_t′ ← channel_align(F_t^{(i_k)}, C=min(C_t,C_s)) F_s′ ← channel_align(F_s^{(j_k)}, C) stack_t ← stack_real_imag(RFFT2(F_t′)) stack_s ← stack_real_imag(RFFT2(F_s′)) L_FFT += MSE(stack_s, stack_t) L_FFT ← L_FFT / m # (D) Total loss & backward L_total ← L_KD + β·L_FFT S.backward(L_total) optimizer.step()

Key implementation details include the use of AdaptiveAvgPool for channel alignment and FFT operations with PyTorch’s native routines.

6. Empirical Evaluation and Benchmarking

SpectralKD achieves state-of-the-art results for ViT distillation on ImageNet-1K across both uniform (DeiT, CaiT) and hierarchical (Swin) architectures.

Method	Student	#Params	Top-1 (%)
Baseline (no KD)	DeiT-Tiny	5M	72.2
Hard KD (CE+KL)	DeiT-Tiny	5M	74.5
FitNets/Manifold/…	DeiT-Tiny	5M	75.9–77.2
SpectralKD	DeiT-Tiny	5M	77.4

Method	Student	#Params	Top-1 (%)
Baseline (no KD)	Swin-Tiny	29M	81.3
Distill/RKD/…	Swin-Tiny	29M	81.2–82.3
SpectralKD	Swin-Tiny	29M	82.7

Relative gains: +5.2% for DeiT-Tiny and +1.4% for Swin-Tiny compared to baselines, without introducing additional trainable parameters. The code base leverages PyTorch and the timm library, with FFT modules via torch.fft.rfft2 (Tian et al., 2024).

7. Distillation Dynamics and Practical Recommendations

During training, SpectralKD students are observed to reproduce teacher-like U-shaped spectral intensity curves, even at layers where no explicit alignment loss is applied. Vanilla students (no KD) have lower or noisier layer intensities in semantic/late layers, while SpectralKD enforces coherence, suggesting that aligning a subset of layers in the spectral domain can propagate segmentation of information complexity across the entire network (“distillation dynamics”).

Practical guidelines include:

• For uniform ViTs (DeiT, CaiT), distill at layers $\{1,2,n-5,\ldots,n\}$; for Swin, use stage boundaries.
• Hyperparameters $(T=1,\;\alpha=0.9,\;\beta=0.2)$ show robust performance across models.
• AdaptiveAvgPool is sufficient for channel alignment; no learnable projector is required.
• Training DeiT-Tiny with SpectralKD on 2 × RTX 4090 requires $\sim$184 GPU-hours at batch size 256.
• Implementation is available at https://github.com/thy960112/SpectralKD.

SpectralKD provides a principled, interpretable, and implementation-friendly approach to ViT knowledge distillation, harnessing frequency-domain structure to yield improved accuracy and offering insights into the flow and hierarchy of information within deep transformers (Tian et al., 2024).