💬 NLP

SRA: Span Representation Alignment for Large Language Model Distillation

This paper introduces SRA, a novel cross-tokenizer knowledge distillation framework that leverages a multi-particle dynamical systems perspective to align robust, attention-weighted span representations rather than individual tokens, significantly outperforming existing methods in challenging cross-architecture scenarios.

Original authors: Quoc Phong Dao, Hoang Son Nguyen, Pham Khanh Chi, Tung Nguyen, Linh Ngo Van, Nguyen Thi Ngoc Diep, Trung Le

Published 2026-05-05

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Quoc Phong Dao, Hoang Son Nguyen, Pham Khanh Chi, Tung Nguyen, Linh Ngo Van, Nguyen Thi Ngoc Diep, Trung Le

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Speaking Different Languages

Imagine you have a brilliant, giant professor (the Teacher Model) who knows everything. You want to teach a smaller, faster student (the Student Model) everything the professor knows so the student can do the job on a regular laptop instead of a supercomputer.

Usually, this works great if the professor and student speak the exact same language and use the same dictionary. But in the world of AI, they often use different "tokenizers."

The Tokenizer Problem: Think of a tokenizer as a way of breaking sentences into puzzle pieces. The Professor might break the word "unbelievable" into three pieces: un, believ, able. The Student might break it into two: un, believable.
The Old Way: Previous methods tried to force the Professor's three pieces to line up perfectly with the Student's two pieces. It's like trying to match a 3-piece puzzle to a 2-piece puzzle. It's brittle, confusing, and often leads to mistakes because the pieces don't fit.

The New Solution: SRA (Span Representation Alignment)

The authors of this paper say, "Stop trying to match the tiny puzzle pieces. Let's match the whole pictures instead."

They introduce a framework called SRA. Instead of focusing on individual words or fragments (tokens), they focus on Spans—chunks of text that make sense together, like a complete phrase or sentence.

Here is how SRA works, using a physics analogy:

1. The "Center of Mass" Analogy

Imagine a swarm of bees flying together.

Old Method: You try to track every single bee individually. If the Teacher's swarm splits differently than the Student's swarm, you lose track.
SRA Method: You don't track individual bees. You look at the Center of Mass of the whole swarm.
- In physics, the "Center of Mass" is the average position of a group of objects, weighted by how heavy (or important) they are.
- In SRA, a "Span" (a chunk of text) is the swarm. The "bees" are the individual tokens.
- The system calculates a Center of Mass for the text chunk. It pays more attention to the "heavier" bees (the most important words, like "not" or "crucial") and less attention to the "lighter" ones (like "the" or "a").
- This creates a single, stable, robust point that represents the meaning of that chunk, regardless of how the Teacher or Student broke the words up.

2. The "Longest Common Subsequence" (LCS) Bridge

How do they know which chunk of the Professor's text matches which chunk of the Student's text if their word breaks are different?

They use a method called LCS. Imagine you have two sentences written in different handwriting. You find the longest string of letters that appears in both, ignoring the spaces or breaks in between.
This allows SRA to say, "Okay, this chunk of the Professor's text corresponds to this chunk of the Student's text," even if the Professor split it into 5 words and the Student split it into 3.

3. Keeping the Shape (Geometric Regularizer)

When you move a group of objects, you don't just want them to end up in the right spot; you want them to keep their shape relative to each other.

If the Professor's text chunks are arranged in a specific pattern (like a triangle), the Student's chunks should also form a triangle, not a square.
SRA uses a special "geometric rule" to ensure that the relationships between the text chunks stay the same. This preserves the structure of the knowledge being transferred.

What Did They Find?

The researchers tested this by teaching powerful, large AI models (like Qwen and Mistral) to smaller, weaker models (like GPT-2 and TinyLLaMA) that used completely different dictionaries.

The Result: SRA consistently beat all the other methods. It was better at teaching the student to understand the teacher's logic, even when they "spoke" different token languages.
Efficiency: It didn't require massive amounts of extra computer power. In fact, it was often faster to train than the previous best methods.

Summary

SRA is like a translator that stops trying to translate word-for-word (which fails when dictionaries differ) and instead translates ideas and phrases. By treating groups of words as single, weighted "centers of gravity," it creates a stable bridge between two AI models that speak different technical languages, allowing the smaller model to learn effectively from the larger one without getting confused by the mismatched puzzle pieces.

Technical Summary: SRA (Span Representation Alignment for Large Language Model Distillation)

Problem Statement

The paper addresses the challenge of Cross-Tokenizer Knowledge Distillation (CTKD), where a large "teacher" language model and a smaller "student" model utilize different tokenizers and vocabularies. Existing CTKD methods primarily rely on token-level or logit-level alignment strategies (e.g., using edit distance, optimal transport, or projection layers). The authors argue that these approaches are often brittle and sensitive to discrepancies between tokenizers because they attempt to align discrete, atomic units that may not correspond semantically across different vocabulary systems. Furthermore, current methods often fail to capture the underlying dynamics of feature evolution within the Transformer architecture when tokenizers differ.

Methodology: SRA Framework

The proposed SRA framework reframes CTKD by shifting the fundamental unit of alignment from individual tokens to robust, tokenizer-agnostic spans. This approach is grounded in the physical analogy of Multi-Particle Dynamical Systems (MPDS), where the Transformer is viewed as a system of interacting particles (tokens) evolving through layers.

1. Span Construction and Mapping

Span Definition: Instead of aligning tokens directly, SRA groups tokens into "spans" that correspond to the same textual units in the raw input.
Alignment via LCS: To match spans between teacher and student models, the authors compute the Longest Common Subsequence (LCS) of the character offsets produced by both tokenizers. This identifies matched segments of the input text, defining start and end indices for aligned spans in both models without requiring direct token-to-token mapping.

2. Span Representation via Center of Mass (CoM)

Drawing from the MPDS perspective, each span is treated as a cluster of particles. The state of this cluster is represented by its Center of Mass (CoM):

Weighted Pooling: The representation of a span is computed as an attention-weighted average of its constituent token hidden states.
Attention Weights: The weight $w_i$ for a token is derived from the attention scores received from the last token of the sequence across all attention heads. This mimics the physical concept where "heavier" (more semantically salient) particles dominate the cluster's center.
Formula: $C_i = \sum w_t H_t$ , where $H_t$ is the hidden state and $w_t$ is the attention-derived weight.

3. Distillation Objectives

SRA employs a combined loss function to transfer knowledge at the span level:

Span-Level Hidden State Transfer ( $L_{HS}^{Span}$ ):
- Aligns the CoM representations of student and teacher spans using a cosine similarity loss.
- Geometric Regularizer ( $L_{Geo}$ ): To preserve the structural integrity of the representation space, a regularizer is added to maintain the relative distances (cosine distance) between pairs of spans in the student model, matching those in the teacher model.
- Weighted Loss: The loss is weighted by the semantic salience of the spans (derived from the sum of token weights within the span), ensuring the student focuses on the most informative regions.
Aligned Span Logits Distillation ( $L_{KD}^{Span}$ ):
- Addresses vocabulary mismatch ( $V_{tea} \neq V_{stu}$ ) by restricting knowledge transfer to the shared subspace $\tilde{V} = V_{tea} \cap V_{stu}$ .
- Span-level logits are aggregated from token logits within the shared vocabulary, and a Kullback-Leibler (KL) divergence is computed between the teacher's and student's distributions over this shared space.
Total Objective: The final loss combines the standard cross-entropy loss on ground-truth labels with the span-level distillation losses (hidden states and logits), balanced by a hyperparameter $\alpha$ .

Key Contributions

Novel Conceptual Framework: The paper introduces a CTKD framework grounded in the Multi-Particle Dynamical System view of Transformers, modeling spans as particle clusters and utilizing their Center of Mass for stable, additive representations.
Span-Level Distillation: It proposes a method to distill knowledge at the span level, encompassing both hidden representations and logits. This includes an attention-derived weighting scheme to prioritize salient spans and a geometric regularizer to preserve the relative geometry of the representation space.
Empirical Validation: The authors demonstrate that SRA significantly outperforms state-of-the-art CTKD baselines (including ULD, DSKD, MinED, and MultilevelOT) across multiple teacher-student pairs and benchmarks.

Experimental Results

The authors evaluated SRA on diverse cross-architecture and cross-tokenizer pairs (e.g., Qwen1.5/2.5, Mistral-7B, GPT-2, OPT, TinyLLaMA) using instruction-following datasets (Dolly, Vicuna, SelfInst, S-NI, DialogSum).

Performance: SRA consistently achieved the highest ROUGE-L scores and the best average performance across all settings compared to strong baselines.
Semantic Quality: Semantic evaluations using GPT-4o-mini as a judge confirmed that SRA produces more semantically faithful and contextually aligned responses.
Ablation Studies:
- Removing the geometric regularizer or the weighted pooling mechanism resulted in performance degradation, confirming the necessity of preserving structural relationships and focusing on salient spans.
- The method showed robustness even in same-tokenizer scenarios, outperforming distance-based and other CTKD methods.
Efficiency: While SRA requires slightly more GPU memory than some baselines due to span matching and aggregation, it achieved the lowest average per-step training time among the compared methods.

Significance and Claims

The paper claims that SRA provides a robust, physically-grounded approach to knowledge distillation that overcomes the brittleness of token-level alignment. By shifting the alignment unit to spans and leveraging the Center of Mass concept, SRA naturally supports the transfer of dynamic, layer-to-layer feature evolution. The authors assert that their method validates the principle that distillation should focus on the dynamics of the system's geometric evolution rather than just final outputs, offering a scalable solution for deploying efficient models across heterogeneous tokenizer environments.

The work is presented as a step toward more stable and comprehensive knowledge transfer, though the authors acknowledge limitations such as the static nature of current logit mapping and the computational cost of on-the-fly teacher inference required for span alignment.