MTA: Multi-Granular Trajectory Alignment for Large Language Model Distillation

Linh Ngo Van; Pham Khanh Chi; Quoc Phong Dao; Thanh Hong Nguyen; Thuat Nguyen; Trung Le

arxiv: 2605.01374 · v2 · pith:DWT2XANHnew · submitted 2026-05-02 · 💻 cs.CL

MTA: Multi-Granular Trajectory Alignment for Large Language Model Distillation

Pham Khanh Chi , Quoc Phong Dao , Thuat Nguyen , Linh Ngo Van , Trung Le , Thanh Hong Nguyen This is my paper

Pith reviewed 2026-07-01 00:27 UTC · model grok-4.3

classification 💻 cs.CL

keywords knowledge distillationlarge language modelsrepresentation alignmentmulti-granular alignmenttrajectory alignmentTransformer representationsmodel compression

0 comments

The pith

MTA improves LLM distillation by aligning representations at word level in lower layers and phrase level in higher layers along their transformation trajectories.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper aims to show that fixed-layer or token-only alignment in distillation leaves the student's capture of the teacher's evolving internal relations incomplete. MTA instead follows the full layer-wise trajectory, switching granularity so lower layers preserve lexical details while higher layers match compositional units. A Dynamic Structural Alignment loss enforces this by equating relative geometry among those units at each depth, supplemented by direct hidden-state matching. If the approach works, distillation should transfer more of the teacher's relational knowledge without adding parameters to the student.

Core claim

MTA aligns teacher and student representations along their layer-wise transformation trajectory with a layer-adaptive strategy: lower layers at the word level to preserve lexical information and higher layers on phrase-level spans to capture compositional semantics, instantiated through a Dynamic Structural Alignment loss that matches the relative geometry among semantic units within each layer together with a Hidden Representation Alignment loss.

What carries the argument

Dynamic Structural Alignment loss, which matches relative geometry among semantic units at each layer while adapting the unit size from words to phrases as depth increases.

If this is right

The student receives stronger guidance to reproduce the teacher's internal relational structure at every depth.
Knowledge transfer improves over both fixed-layer and purely token-level distillation methods.
Standard benchmark scores rise consistently across tested models and tasks.
Ablations isolate measurable contributions from the multi-granular trajectory component and the hidden-state term.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The same depth-dependent granularity switch might reduce the performance gap in other compression methods such as quantization-aware training.
If the abstraction gradient is architecture-dependent, MTA would need re-tuning of the word-to-phrase transition point when moving beyond standard Transformers.
Phrase-level spans could be replaced by other linguistically motivated units without changing the overall loss formulation.
The method implies that distillation objectives should track how meaning is built compositionally rather than only matching final outputs.

Load-bearing premise

That Transformer representations grow more abstract with depth, so word-level matching suits early layers and phrase-level matching suits later layers better than any uniform scheme.

What would settle it

An ablation on the same benchmarks in which the layer-adaptive component is removed and performance falls to or below the fixed-alignment baselines.

Figures

Figures reproduced from arXiv: 2605.01374 by Linh Ngo Van, Pham Khanh Chi, Quoc Phong Dao, Thanh Hong Nguyen, Thuat Nguyen, Trung Le.

**Figure 1.** Figure 1: The correspondence between linguistic compositionality and the layer-wise evolution of representations in large language models. between teacher and student representations. Together, these objectives constrain both the internal structure of representations within each layer and their transformation across depth. 3.1 Motivation: The Hierarchical Representational Trajectory Most existing Knowledge Distilla… view at source ↗

**Figure 2.** Figure 2: Dynamic Structural Alignment (LDSA). This objective enforces geometric consistency. It calculates the pairwise relational distances between semantic spans (words or phrases) within a layer for both Teacher and Student. By minimizing the discrepancy between these two topological structures across network depths, the Student learns to replicate the Teacher’s representational trajectory. This layer-adaptive d… view at source ↗

**Figure 3.** Figure 3: Hidden Representation Alignment Strategy. The student learns to match teacher representations using a weighted cosine distance objective, ensuring accurate feature reconstruction at key layers. R dS×dT to map the Student’s representations into the Teacher’s space: H˜ S t,l = HS t,lWl (12) We then minimize the weighted cosine distance between the projected Student states and the Teacher states: LHid = X l∈… view at source ↗

**Figure 4.** Figure 4: GPT-4o-mini evaluation scores (1-100) for view at source ↗

**Figure 5.** Figure 5: GPT-4o-mini evaluation scores (1-100) for view at source ↗

**Figure 6.** Figure 6: Effect of the number of distilled intermediate view at source ↗

**Figure 7.** Figure 7: Prompt for GPT-4 evaluation using Ground view at source ↗

read the original abstract

Knowledge distillation is a key technique for compressing large language models (LLMs), but most existing methods align representations at fixed layers or token-level outputs, ignoring how representations evolve across depth. As a result, the student is only weakly guided to capture the teacher's internal relational structure during distillation, which limits knowledge transfer. To address this limitation, we propose Multi-Granular Trajectory Alignment (MTA), a framework that aligns teacher and student representations along their layer-wise transformation trajectory. MTA adopts a layer-adaptive strategy: lower layers are aligned at the word level to preserve lexical information, while higher layers operate on phrase-level spans (e.g., noun and verb phrases) to capture compositional semantics. We instantiate this idea through a Dynamic Structural Alignment loss that matches the relative geometry among semantic units within each layer. This design is motivated by empirical findings that Transformer representations become increasingly abstract with depth, and is also consistent with linguistic views in which higher-level meaning emerges through the composition of lower-level lexical units. We further incorporate a Hidden Representation Alignment loss to directly align selected teacher-student layers. Experiments show that MTA consistently outperforms state-of-the-art baselines on standard benchmarks, with ablations confirming the contribution of each component.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

MTA adds a depth-dependent switch from word-level to phrase-level alignment in distillation but gives no data to justify that specific choice.

read the letter

The main thing to know is that this paper adds a layer-dependent granularity switch to knowledge distillation for LLMs, aligning at the word level early and phrase level later, but the evidence for why that switch helps is missing from the abstract.

They motivate it by noting that transformer representations become more abstract deeper in the model, so lower layers should match lexical items while higher ones match compositional phrases. The new loss tries to align the relative geometry among these semantic units along the layer trajectory. They combine it with standard hidden representation alignment.

What works here is the recognition that one-size-fits-all alignment might not capture how the teacher's knowledge is structured at different stages. Treating the distillation as a trajectory rather than isolated layer matches is a reasonable step beyond prior fixed approaches.

The soft spot is the untested assumption about depth and granularity. The abstract cites empirical findings and linguistic ideas but shows no layer-wise probing or analysis to back up the specific word-to-phrase transition. That leaves open the possibility that any improvements come from the extra loss terms or hyperparameter choices instead of the multi-granular design. The dynamic structural alignment is also only sketched at a high level, without enough detail to see how it improves on simpler matching methods.

The claimed consistent outperformance and successful ablations can't be checked without the actual results, baselines, and effect sizes.

This paper is aimed at people working on LLM compression and efficient inference. Readers in that niche could get some value from the idea if the full experiments are solid. It deserves a serious referee to examine whether the granularity choice is supported by the data and whether the gains are reliable.

I would send it out for peer review.

Referee Report

2 major / 0 minor

Summary. The paper proposes Multi-Granular Trajectory Alignment (MTA) for LLM knowledge distillation. It aligns teacher-student representations along layer-wise trajectories using a layer-adaptive strategy (word-level alignment in lower layers to preserve lexical information; phrase-level spans in higher layers to capture compositional semantics). This is realized via a Dynamic Structural Alignment loss that matches relative geometry among semantic units per layer, plus a Hidden Representation Alignment loss for direct layer matching. The design is motivated by the increasing abstraction of Transformer representations with depth. Experiments claim consistent outperformance over SOTA baselines on standard benchmarks, supported by ablations on each component.

Significance. If the claimed gains are robust and attributable to the multi-granular trajectory component rather than unablated factors, the work could meaningfully advance distillation by incorporating depth-dependent linguistic structure, potentially yielding more efficient student models that better preserve the teacher's internal relational geometry. The combination of trajectory alignment with adaptive granularity offers a concrete alternative to fixed-layer or token-only methods.

major comments (2)

[Abstract] Abstract: the central claim that MTA outperforms baselines specifically because the layer-adaptive word/phrase switch better captures the teacher's relational structure rests on the premise that representations become 'increasingly abstract with depth,' yet the text provides no quantitative support (e.g., no layer-wise probing accuracies, no cited empirical studies, and no comparison of lexical vs. compositional information by depth). This motivation is load-bearing for the design choice over simpler fixed-granularity alternatives.
[Abstract] Abstract: the Dynamic Structural Alignment loss is described only at a high level as matching 'relative geometry among semantic units within each layer,' without equations, pseudocode, or explicit comparison to token-level or fixed-layer baselines. Consequently it is impossible to determine whether observed gains derive from this loss or from the separate Hidden Representation Alignment loss.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive feedback. We address the two major comments point by point below, providing clarifications based on the full manuscript while noting where revisions can strengthen the presentation.

read point-by-point responses

Referee: [Abstract] Abstract: the central claim that MTA outperforms baselines specifically because the layer-adaptive word/phrase switch better captures the teacher's relational structure rests on the premise that representations become 'increasingly abstract with depth,' yet the text provides no quantitative support (e.g., no layer-wise probing accuracies, no cited empirical studies, and no comparison of lexical vs. compositional information by depth). This motivation is load-bearing for the design choice over simpler fixed-granularity alternatives.

Authors: We acknowledge that the abstract itself does not include quantitative evidence such as layer-wise probing accuracies. The full manuscript motivates the layer-adaptive strategy by referencing established empirical observations on increasing abstraction with depth (discussed in the introduction and related work sections) and linguistic compositionality principles. To address the concern directly, we will revise the manuscript by adding specific citations to relevant empirical studies on layer-wise representation properties and, where feasible, include or reference probing results comparing lexical versus compositional information across depths. revision: yes
Referee: [Abstract] Abstract: the Dynamic Structural Alignment loss is described only at a high level as matching 'relative geometry among semantic units within each layer,' without equations, pseudocode, or explicit comparison to token-level or fixed-layer baselines. Consequently it is impossible to determine whether observed gains derive from this loss or from the separate Hidden Representation Alignment loss.

Authors: The abstract provides a concise overview by design. The full manuscript details the Dynamic Structural Alignment loss with explicit equations for matching relative geometry among semantic units (Section 3.2), includes pseudocode in Algorithm 1, and presents ablations in Section 4.3 that isolate its contribution from the Hidden Representation Alignment loss. These sections also compare against token-level and fixed-layer baselines. The experimental results and ablations therefore allow attribution of gains to the specific components. revision: no

Circularity Check

0 steps flagged

No circularity: method defined independently of performance claims

full rationale

The paper proposes MTA as a new framework with explicitly defined components (layer-adaptive word/phrase alignment via Dynamic Structural Alignment loss plus Hidden Representation Alignment loss). These are presented as design choices motivated by general observations about depth-dependent abstraction in Transformers, not derived from or fitted to the target benchmark metrics. Performance claims rest on external experiments and ablations rather than any self-referential reduction, fitted parameter renamed as prediction, or load-bearing self-citation chain. No equations or uniqueness theorems are shown that collapse back to inputs by construction.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Only abstract available; no free parameters, axioms, or invented entities are specified in the provided text.

pith-pipeline@v0.9.1-grok · 5754 in / 1047 out tokens · 24272 ms · 2026-07-01T00:27:56.175347+00:00 · methodology

Review history (2 revisions) →

discussion (0)

Reference graph

Works this paper leans on

4 extracted references · 2 canonical work pages

[1]

CoRR , volume =

Visualizing and measuring the geometry of BERT. InAdvances in Neural Information Process- ing Systems 32 (NeurIPS 2019), pages 8592–8600. Stephen Crain and Mineharu Nakayama. 1987. Struc- ture dependence in grammar formation.Language, pages 522–543. Daniel Gildea and Daniel Jurafsky. 2002. Automatic la- beling of semantic roles.Computational Linguistics, ...

work page arXiv 2019
[2]

Distillm-2: A contrastive approach boosts the distillation of llms.arXiv preprint arXiv:2503.07067,

Distillm-2: A contrastive approach boosts the distillation of llms.arXiv preprint arXiv:2503.07067. ICML 2025 Spotlight. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se- Young Yun. 2024. Distillm: Towards streamlined distillation for large language models.arXiv preprint arXiv:2402.03898. ICML 2024. Tue Le, Hoang Tran Vuong, Quyen Tran, Linh Ngo Van, Mehrtas...

work page arXiv 2025
[3]

Richard Socher, Cliff C

A primer in bertology: What we know about how BERT works.Transactions of the Association for Computational Linguistics, 8:842–866. Richard Socher, Cliff C. Lin, Andrew Y . Ng, and Christopher D. Manning. 2011. Parsing natural scenes and natural language with recursive neural networks. InProceedings of the 28th International Conference on Machine Learning ...

2011
[4]

pulling up

Super-naturalinstructions: Generalization via declarative instructions on 1600+ NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 5085–5109, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics. Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou. 2020. BERT-o...

2022

[1] [1]

CoRR , volume =

Visualizing and measuring the geometry of BERT. InAdvances in Neural Information Process- ing Systems 32 (NeurIPS 2019), pages 8592–8600. Stephen Crain and Mineharu Nakayama. 1987. Struc- ture dependence in grammar formation.Language, pages 522–543. Daniel Gildea and Daniel Jurafsky. 2002. Automatic la- beling of semantic roles.Computational Linguistics, ...

work page arXiv 2019

[2] [2]

Distillm-2: A contrastive approach boosts the distillation of llms.arXiv preprint arXiv:2503.07067,

Distillm-2: A contrastive approach boosts the distillation of llms.arXiv preprint arXiv:2503.07067. ICML 2025 Spotlight. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se- Young Yun. 2024. Distillm: Towards streamlined distillation for large language models.arXiv preprint arXiv:2402.03898. ICML 2024. Tue Le, Hoang Tran Vuong, Quyen Tran, Linh Ngo Van, Mehrtas...

work page arXiv 2025

[3] [3]

Richard Socher, Cliff C

A primer in bertology: What we know about how BERT works.Transactions of the Association for Computational Linguistics, 8:842–866. Richard Socher, Cliff C. Lin, Andrew Y . Ng, and Christopher D. Manning. 2011. Parsing natural scenes and natural language with recursive neural networks. InProceedings of the 28th International Conference on Machine Learning ...

2011

[4] [4]

pulling up

Super-naturalinstructions: Generalization via declarative instructions on 1600+ NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 5085–5109, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics. Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou. 2020. BERT-o...

2022