Technical Summary: Multi-Granular Trajectory Alignment (MTA) for LLM Distillation

Problem Statement

Existing Knowledge Distillation (KD) methods for Large Language Models (LLMs) primarily focus on aligning output probability distributions (logits) or intermediate hidden states at fixed, static layers. While effective for transferring local signals, these approaches often employ a uniform, token-level alignment strategy across all network depths.

The paper argues that this uniform treatment is misaligned with the hierarchical nature of human language and the functional organization of Transformer models. Linguistic research and Transformer interpretability studies indicate that:

1. Lower layers emphasize surface-level, lexical, and factual retrieval features.
2. Higher layers progressively encode abstract semantic information and compositional reasoning.

By treating all tokens in isolation and applying the same alignment objective uniformly, current methods fail to capture the progressive evolutionary trajectory of representations as they transform from lexical units into abstract semantic structures. This limitation restricts the student model's ability to generalize to downstream tasks.

Methodology: Multi-Granular Trajectory Alignment (MTA)

To address these limitations, the authors propose MTA, a modular framework that aligns the representational trajectory of teacher and student models using a layer-adaptive strategy. MTA operates on the hypothesis that distillation should respect the depth-dependent semantic roles of Transformer layers.

Core Components

1. Layer-Adaptive Multi-Granular Spans:

   • Lower Layers: Alignments are performed at the word-level (Word Spans) to preserve lexical grounding and fine-grained dependencies.
   • Higher Layers: Alignments shift to phrase-level (Noun Phrases and Verb Phrases) to capture compositional semantics and higher-order structural relationships.
   • Span Construction: Semantic units are extracted using a syntactic parser (e.g., spaCy). Token importance weights are computed via a bidirectional self-attention mechanism (excluding self-loops) to aggregate token representations into span representations.

2. Dynamic Structural Alignment Loss ($L_{DSA}$):

   • This objective enforces geometric consistency between the teacher and student.
   • Instead of matching individual token vectors, $L_{DSA}$ minimizes the discrepancy in the pairwise relational distances between semantic spans (words or phrases) within a specific layer.
   • It prioritizes semantically salient spans by weighting the alignment based on the aggregated importance of constituent tokens.
   • Mathematically, it aligns the topological structure of the representation space:
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     where $d(\cdot)$ is cosine distance, and $U$ represents the weighted span representation.

3. Hidden Representation Alignment Loss ($L_{Hid}$):

   • To complement structural constraints, a feature-level loss aligns specific hidden states.
   • Since student dimensions are typically smaller than teacher dimensions, a learnable linear projector maps student states to the teacher's space.
   • The loss minimizes the weighted cosine distance between projected student states and teacher states, focusing on the most informative tokens.

4. Optimization:

   • MTA is designed as a plug-in module. The total loss combines the base distillation loss ($L_{Base}$) of existing frameworks (e.g., DistiLLM, FDD) with the proposed structural and hidden losses:
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

Key Contributions

1. Identification of Limitations: The authors identify that uniform intermediate alignment ignores the hierarchical evolution of representations, leading to suboptimal knowledge transfer.
2. Layer-Adaptive Framework: They introduce MTA, which dynamically transitions from word-level alignment in lower layers to phrase-level structural alignment in higher layers, mirroring the bottom-up compositional nature of language.
3. Generalized Module: MTA is demonstrated to be a generalized optimization layer compatible with state-of-the-art distillation methods (DistiLLM, DistiLLM-2, FDD) that share tokenization.
4. Empirical Validation: Extensive experiments show that MTA consistently improves performance across diverse model architectures (GPT-2, Qwen1.5, OPT) and benchmarks.

Experimental Results

The authors evaluated MTA on instruction-following tasks using three student-teacher pairs:

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

Performance Metrics:

• ROUGE-L: MTA consistently outperformed baselines (FDD, DistiLLM, DistiLLM-2) across all datasets (Dolly, SelfInst, Vicuna, S-NI). For instance, in the GPT-2 1.5B $\to$ 120M setting, adding MTA to DistiLLM improved the average ROUGE-L from 20.21 to 21.45.
• LLM-as-a-Judge: Using GPT-4o-mini to evaluate factual accuracy, coherence, and helpfulness, MTA-enhanced models showed consistent score improvements over baselines.

Ablation Studies:

• Loss Components: Both $L_{DSA}$ (structural) and $L_{Hid}$ (feature) contributed positively, with their combination yielding the best results.
• Granularity: The "Full-level" (layer-adaptive) strategy outperformed static strategies (Word-only or Phrase-only), confirming that aligning lexical grounding at lower layers and compositional semantics at higher layers is crucial.
• Layer Selection: The method is robust to different layer schedules, though a hybrid allocation (1 Word-level layer : 2 Phrase-level layers) was found to be optimal for the tested configurations.

Efficiency:

• MTA introduces training overhead primarily due to syntactic span extraction (reduced to ~0.27s/step via optimization) and geometric matching.
• Inference costs remain unchanged as span extraction is not used at test time.
• Time-matched comparisons confirmed that performance gains are due to the alignment objective, not increased training budget.

Significance and Claims

The paper claims that MTA offers a semantically faithful approach to knowledge distillation by explicitly modeling the hierarchical representational trajectory of LLMs. By moving beyond static, token-level matching to dynamic, multi-granular structural alignment, MTA enables student models to better replicate the teacher's internal information processing dynamics.

The authors position MTA not as a replacement for existing distillation frameworks, but as a complementary, plug-and-play module that enhances the generalization and robustness of compressed LLMs. The work underscores the importance of respecting linguistic hierarchy and depth-dependent feature evolution in the design of future model compression techniques.