Decoder-based Sense Knowledge Distillation

Qitong Wang; Mohammed J. Zaki; Georgios Kollias; Vasileios Kalantzis

cs.CLFeb 25, 2026

Decoder-based Sense Knowledge Distillation

Authors: Qitong Wang, Mohammed J. Zaki, Georgios Kollias, Vasileios Kalantzis

Organizations: Rensselaer Polytechnic Institute Troy, New York, USA · IBM Research Yorktown Heights, New York, USA

Abstract

Large language models (LLMs) learn contextual embeddings that capture rich semantic information, yet they often overlook structured lexical knowledge such as word senses and relationships. Prior work has shown that incorporating sense dictionaries can improve knowledge distillation for encoder models, but their application to decoder as generative models remains challenging. In this paper, we introduce Decoder-based Sense Knowledge Distillation (DSKD), a framework that integrates lexical resources into the training of decoder-style LLMs without requiring dictionary lookup at inference time. Extensive experiments on diverse benchmarks demonstrate that DSKD significantly enhances knowledge distillation performance for decoders, enabling generative models to inherit structured semantics while maintaining efficient training.

Figures & tables

Figure 1. Distribution of selected word embeddings from Llama-3-8B-Instruct ( Dubey et al., 2024 ) on Wikipedia dump ( WikipediaCorpus, 2024 ) , with t-SNE ( Van der Maaten and Hinton, 2008 ) for visualization. Semantically related tokens form coherent clusters in the embedding space: synonyms of red (e.g., ruby , crimson , vermilion , scarlet ) appear in a compact region, while vehicle-related tokens such as bus and car occupy a separate neighborhood. For words that are decomposed into multiple subword tokens by the tokenizer (e.g., vermilion , scarlet ), a composed representation is used for visualization; details are provided in Section 3.1. This structured organization suggests that contextual embeddings encode discrete semantic regularities.

Original Pair	Original Relationship	New Pair	New Relationship
quit, discontinue	synonym	quit, continue	antonym
accurate, faultless	synonym	accurate, fault	antonym
opposite, dissimilar	synonym	opposite, similar	antonym
variable, unchangeable	antonym	variable, changeable	synonym
unkind, friendly	antonym	kind, friendly	synonym
disinterest, zeal	antonym	interest, zeal	synonym

Table 1. Examples of lexical relation transformation under morphological negation.

	LLaMA		Mistral
$m$	Synonym (%)	Antonym (%)	Synonym (%)	Antonym (%)
1	48.11	60.09	34.97	41.48
2	82.22	87.03	72.39	74.38
3	95.06	97.86	91.21	93.96
4	98.54	99.61	96.96	98.99
5	99.42	99.91	98.65	99.68

Table 2. Sensitivity analysis of the maximum span length

m

. Keep rate denotes the percentage of tokenized lexical relations whose span length does not exceed

m

Figure 2. Overview of DSKD training: Given an input corpus, the teacher’s last hidden layer $\mathbf{m}_{t}$ is used to retrieve sense embeddings of the token $t+1$ and its synonyms (positive candidates) and antonyms (negative candidates) from the sense dictionary. The $\kappa$ nearest positive and negative embeddings to $\mathbf{m}_{t}$ are selected, and the student’s last hidden layer $\mathbf{n}_{t}$ is optimized with a hinge MSE loss. This process aligns the student’s semantic space with the teacher’s sense dictionary.

Model	# of Layers	ARC	CSQA	MMLU	PIQA	SQuAD
LLaMA3-8B	32	81.45 $\pm$ 1.07	78.08 $\pm$ 0.41	67.22 $\pm$ 0.20	78.49 $\pm$ 1.33	72.90 $\pm$ 0.13
KD	16	77.89 $\pm$ 0.42	76.43 $\pm$ 0.45	61.26 $\pm$ 0.59	81.78 $\pm$ 0.34	68.58 $\pm$ 0.13
DSKD	16	79.66 $\pm$ 0.34	76.70 $\pm$ 0.17	64.38 $\pm$ 0.38	82.63 $\pm$ 0.27	69.11 $\pm$ 0.36
Mistral-7B	32	70.33 $\pm$ 0.57	68.34 $\pm$ 0.75	55.37 $\pm$ 0.27	73.72 $\pm$ 0.59	66.26 $\pm$ 0.13
KD	16	61.11±0.31	66.16 $\pm$ 0.78	52.04 $\pm$ 0.43	78.88 $\pm$ 0.22	57.82 $\pm$ 0.25
DSKD	16	64.78 $\pm$ 0.43	68.44 $\pm$ 0.70	54.44 $\pm$ 0.37	79.53 $\pm$ 0.45	58.09 $\pm$ 0.18

Table 3. Performance comparison between student models trained with KD and DSKD. Bold numbers indicate the best results among the student models.

Panel of Figure 3 — Figure 3. Performance of KD and DSKD distilled from Llama-3-8b-Instruct on ARC and MMLU across varying (a) total number of layers, and (b) number of trainable student layers.

$k$	LLaMA		Mistral
	ARC	MMLU	ARC	MMLU
5	77.33	62.18	64.78	54.44
10	79.66	64.38	64.33	54.25
15	79.00	64.57	65.00	54.91
20	81.33	64.68	64.00	53.68
25	80.00	65.23	65.33	54.99

Table 4. Effect of the number of clusters

k

on ARC and MMLU benchmarks. The selected value of

k

used in subsequent experiments is highlighted in bold.

$\kappa$	LLaMA		Mistral
	ARC	MMLU	ARC	MMLU
1	78.00	63.93	63.67	53.60
5	79.66	64.38	64.78	54.44
10	79.33	64.20	64.33	54.35
15	80.33	64.76	65.67	54.53
20	79.00	64.71	66.00	53.68

Table 5. Effect of

\kappa

on ARC and MMLU benchmarks.

$\beta_{p}$ , $\beta_{n}$	LLaMA		Mistral
	ARC	MMLU	ARC	MMLU
0.5	77.67	61.42	64.00	52.07
1.0	79.66	64.38	63.33	53.51
1.5	79.33	63.45	64.78	54.44
2.0	78.33	61.45	61.33	50.57

Table 6. Performance on ARC and MMLU under different values of

\beta_{p}

and

\beta_{n}

Model	# Layers	# Params	Memory	Train	Eval
				Time	Time
Llama-3-8B	32	8.0B	15.0GB	N/A	45min
KD	16	4.5B	8.5GB	5h02min	23min
DSKD	16	4.5B	8.5GB	5h17min	22min

Table 7. Comparison of model size, memory, training time per epoch, and MMLU evaluation time for Llama-3-8B and its student models.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

	Llama-3-8B-Instruct	Mistral-7B-Instruct
Tokenizer vocab size	128,256	32,000
# of Sense token keys	69,832	27,541
# of Word token keys	15,695	27,702
Sense dict size	7.08 GB	2.06 GB
# of synonym pairs	923,950	923,950
# of antonym pairs	75,949	75,949

Appendix

Table 8. Vocabulary and dictionary statistics for Llama-3-8B and Mistral-7B models.

$\alpha$	LLaMA		Mistral
	ARC	MMLU	ARC	MMLU
0	75.67	61.25	58.67	51.75
0.5	78.33	62.03	59.33	49.27
1	77.89	61.26	61.11	52.04
2	80.33	64.76	60.00	51.28

Appendix

Table 9. Effect of

\alpha

on ARC and MMLU benchmarks across LLaMA and Mistral backbones on KD. The selected value of

\alpha

used in subsequent experiments is highlighted in bold.

	Llama-3 KD	Llama-3 DSKD	Mistral KD	Mistral DSKD
Total layers ( $L$ )	16	16	16	16
Trainable layers	2	2	2	2
Learning rate	0.00005	0.00005	0.00002	0.00002
epoch	2	2	2	2
candidate sample size $\kappa$	N/A	5	N/A	5
$\alpha$	1.0	1.0	1.0	1.0

Appendix

Table 10. Full training hyperparameter configurations.

Model	# of Layers	# of Parameters	Memory	Training Time
Llama-3-8B-Instruct	32	8.0B	15.0GB	N/A
KD	16	4.5B	8.5GB	5h02min
DSKD	16	4.5B	8.5GB	5h17min
Mistral-7B-Instruct	32	7.2B	13.5GB	N/A
KD	16	3.7B	6.9GB	4h53min
DSKD	16	3.7B	6.9GB	5h04min

Appendix

Table 11. Comparison of space and training time.

Model	ARC	CSQA	MMLU	PIQA	SQUADV2
Llama-3-8B-Instruct	1min	2.5min	45min	4.5min	23min
KD	1min	1.5min	23min	3min	15min
DSKD	1min	1.5min	22min	3min	15min
Mistral-7B-Instruct	1min	2.5min	51min	5min	21min
KD	1min	1.5min	25min	3min	15.5min
DSKD	1min	1.5min	26min	3min	15min

Appendix

Table 12. Evaluation time comparison across models and tasks.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 512 papers

WASD: Wasserstein-based Knowledge Distillation for Large Language Models
Oct 6, 2026Byeonghu Na, Donghyeok Shin, Yeongmin Kim +2Language Model Distillation Knowledge Distillation
KAIST
Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch
Mar 23, 2026Stella Eva Tsiapali, Cong-Thanh Do, Kate KnillCross-Tokenizer Knowledge Distillation Language Model Distillation
University of Cambridge, Department of Engineering, Cambridge, United Kingdom · Toshiba Europe Limited, Cambridge, United Kingdom
KDFlow: A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models
Mar 2, 2026Songming Zhang, Xue Zhang, Tong Zhang +3Cross-Tokenizer Knowledge Distillation Language Model Distillation
School of Computer Science and Technology, Beijing Jiaotong University, Beijing, China · Key Laboratory of Big Data & Artificial Intelligence in Transportation, (Beijing Jiaotong University), Ministry of Education · Weixin AI

Oct 6, 2026J/K move · Enter open · S save

WASD: Wasserstein-based Knowledge Distillation for Large Language Models

Byeonghu Na, Donghyeok Shin, Yeongmin Kim, Mina Kang+1

KAIST

Language Model Distillation Knowledge Distillation

Open paper AskPDF arXiv

Autoregressive large language models (LLMs) have rapidly advanced in capability, but their increasing scale comes with substantial computational and memory costs at inference time. Knowledge distillation (KD) offers a practical solution by transferring knowledge from a large teacher model to a smaller student model via alignment of discrete probability distributions. However, existing KD methods for LLMs primarily rely on divergences that evaluate discrepancies through probability values at each vocabulary index, without explicitly leveraging token-level semantic information. We propose Wasserstein-based knowledge distillation (WASD) for LLMs, which incorporates token-level semantic information via the Wasserstein-based distance with a cost matrix derived from token embeddings. To ensure computational tractability, we adopt the Sinkhorn divergence and derive a gradient-equivalent objective that can be efficiently optimized without introducing additional networks. Experiments across multiple LLM families and scales show that WASD consistently improves distillation performance on diverse tasks, including instruction following, mathematical reasoning, and code generation. Our results highlight the importance of semantic information encoded in the token space for effective distribution alignment in LLM distillation. The implementation is publicly available at https://github.com/aailab-kaist/WASD .

Abstract

Language Model Distillation

Explore similar work

WASD: Wasserstein-based Knowledge Distillation for Large Language Models

Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch

KDFlow: A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models