Diversity-Aware Reverse Kullback-Leibler Divergence for Large Language Model Distillation
Organizations: Rochester Institute of Technology, Rochester, NY, USA · Rowan University, Glassboro, NJ, USA
Abstract
Reverse Kullback-Leibler (RKL) divergence has recently emerged as the preferred objective for large language model (LLM) distillation, consistently outperforming forward KL (FKL), particularly in regimes with large vocabularies and significant teacher-student capacity mismatch, where RKL focuses learning on dominant modes rather than enforcing dense alignment. However, RKL introduces a structural limitation that drives the student toward overconfident predictions. We first provide an analysis of RKL by decomposing its gradients into target and non-target components, and show that non-target gradients consistently push the target logit upward even when the student already matches the teacher, thereby reducing output diversity. In addition, RKL provides weak supervision over non-target classes, leading to poor tail alignment. To address these issues, we propose Diversity-aware RKL (DRKL), which removes this gradient effect and strengthens non-target supervision while preserving the optimization benefits of RKL. Extensive experiments across datasets and model families demonstrate that DRKL consistently outperforms FKL, RKL, and other state-of-the-art distillation objectives, achieving better performance and a superior fidelity-diversity trade-off.
Figures & tables
| Method | Dolly | Self-Instruct | Vicuna | Super-NI | UnNI | Avg. ( ) |
|---|---|---|---|---|---|---|
| GPT-2 XL (Teacher) | 27.00 0.19 | 14.07 0.37 | 16.31 0.32 | 26.46 0.41 | 31.10 0.06 | 22.99 |
| GPT-2 XL (1.5B) GPT-2 Base (0.1B) | ||||||
| FKL | 23.80 0.55 | 9.68 0.49 | 14.53 0.31 | 16.27 0.24 | 19.03 0.09 | 16.66 |
| RKL ( Gu et al., 2024 ) | 24.67 0.13 | 11.25 0.35 | 15.82 0.43 | 21.03 0.21 | 23.94 0.12 | 19.34 |
| Sym-KL | 24.38 0.14 | 11.07 0.43 | 15.72 0.75 | 19.92 0.25 | 22.21 0.03 | 18.66 |
| DKD ( Zhao et al., 2022 ) | 22.88 0.23 | 10.28 0.29 | 16.12 0.31 | 16.57 0.25 | 19.64 0.13 | 17.10 |
| Case study # 1 | |
|---|---|
| Instruction | Given a paragraph with different names for the same thing, write a question that asks about this. |
| Input | Passage: The United States of America is composed of fifty states, each with its own state capital. California’s capital is Sacramento, while Florida’s state capital is Tallahassee and Texas’ capital city is Austin. |
| Ground-truth | What is the capital of the United States of America? |
| FKL | Yes, the United States of America is made of fifty states, each having its own state capital. The states that comprise the United States of America are California, Florida, Pennsylvania, and Texas. |
| RKL | A passage describing the United States of America includes the phrase ”each with its own state capital”. This questions implies that the United States is composed of anywhere between 50 and 100 states, with each state capital being its own distinct entity. |
| Sym-KL | The United States of America is comprised of fifty states, each of which has its own state capital. <n> California’s state capital is Sacramento, while Georgia’s state capital is Atlanta and Texas’ capital city, Austin. |
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
| Case study # 2 | |
| Instruction | In this task, you will be given a piece of text and asked to report what someone said. Your answer should contain all the necessary information about who said what, but it can be delivered in any format. |
| Input | ”I’m so excited for the party tonight!” cried Sarah. |
| Ground-truth | Sarah said ”I’m so excited for the party tonight!” |
| FKL | ”I’m so excited for tonight!” cried Sarah |
| RKL | Sarah is so excited for the party tomorrow! She lives in the city and is going to have a blast. The weather is finally good, so she can go out and play. She typically brings along her two little crushes, who are her little guards. She plans to buy them each a party band, in memory of her late boyfriend. |
| Sym-KL | Text: ”Brian! Brian, I’m so excited for the Party tonight!” cried Sarah. |
| Case study # 4 | |
| Instruction | Given a paragraph, write down all the proper nouns in the order they occur. A proper noun is a specific name for a person, place, or thing and is always capitalized. |
| Input | Passage: Toby thought that Mr. Krebs was being unfair when he gave out detentions for talking during class, but Mrs. Dvorak said that it was part of his job as a teacher. |
| Ground-truth | Toby, Mr. Krebs, Mrs. Dvorak |
| FKL | Toby, Ms. Dvorak, and Mr. Krebs are all proper nouns. |
| RKL | Toby, Mrs. Dvorak, Mr. Krebs, Class |
| Sym-KL | Toby, Mrs. Dvorak |
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