🤖 AI

Scaling Laws for Task-Specific LLM Distillation

This paper establishes empirical scaling laws for domain-specific LLM distillation in quantitative finance, demonstrating that while compression predictably degrades in-domain performance, chain-of-thought supervision effectively recovers general knowledge that would otherwise collapse under structural pruning.

Original authors: Lavinia Ghita, Dhruv Desai, Ioana Boier

Published 2026-06-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lavinia Ghita, Dhruv Desai, Ioana Boier

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

Imagine you have a brilliant, world-class chef (the Large Language Model or LLM) who can cook anything from simple soup to complex molecular gastronomy. This chef is incredibly talented but also huge, expensive to feed, and takes a long time to prepare a meal. You want to hire a smaller, faster, cheaper apprentice (the Student Model) who can work in a busy, high-speed kitchen where speed and cost matter most.

The problem is: How do you train this apprentice to be almost as good as the master without losing their ability to cook anything (general knowledge), while still making them perfect at your specific menu (financial news)?

This paper is a recipe book for that training process. Here is what the authors discovered, explained simply:

1. The Training Methods: "Just the Answer" vs. "Show Your Work"

The researchers tested two main ways to teach the apprentice:

The "Just the Answer" Method (Label-only): The master chef says, "This dish is 'Spicy Ramen'." The apprentice just memorizes the name.
The "Show Your Work" Method (Chain-of-Thought): The master chef says, "This dish is 'Spicy Ramen' because I see the chili oil, the noodles, and the broth." The apprentice learns the reasoning behind the answer.

The Big Discovery:
If you only teach the apprentice the "Answer," they get good at your specific menu but forget how to cook anything else. They become narrow specialists who can't think outside the box.
However, if you teach them the "Reasoning" (Chain-of-Thought), something magical happens. Even as you shrink the apprentice's brain (compress the model), they retain their general cooking skills. The reasoning acts like an anchor, keeping their general knowledge from floating away.

2. The Shrinking Process: "Slicing the Cake"

To make the model smaller, the team used a technique called Pruning. Imagine the chef is a giant cake. To make it smaller, you have to slice off layers.

The Mistake: If you try to slice off 80% of the cake in one giant chunk, the cake collapses. The apprentice fails completely.
The Solution: You slice the cake in small, bite-sized pieces over several rounds. After each slice, you let the apprentice practice (distillation) to recover their skills before you slice again.
The Result: By slicing slowly and practicing in between, they managed to shrink the chef to just 16% of their original size while still keeping them surprisingly competent at the specific task.

3. The "Blended" Secret Sauce

The researchers found that simply asking for "Reasoning" wasn't enough; the training got messy. They invented a Blended Supervision method.
Think of it like grading a student's test. If you only grade the final answer, they might guess. If you only grade the long explanation, they might ramble.
The "Blended" method grades both the final answer and the reasoning steps, giving the answer extra weight. This stabilizes the training, ensuring the apprentice learns the right answer through the right reasoning.

4. The Trade-Off: Speed vs. Wisdom

The paper highlights a crucial trade-off:

If you want the fastest, most efficient apprentice for a specific job: Use the "Just the Answer" method with a lot of data. They will be great at your specific menu but might forget how to cook a different cuisine.
If you need an apprentice who stays smart and versatile: Use the "Show Your Work" (Blended Chain-of-Thought) method. It takes a bit more effort to train, but it saves the apprentice's general intelligence, preventing them from becoming a "one-trick pony."

5. The "Hidden Cost" of Training

One of the most surprising findings is that the act of training itself causes more damage than the shrinking does.
Imagine the apprentice is a perfect copy of the master. When you start training them on your specific menu, they naturally drift away from the master's style just by learning the new rules. The paper found that this "drift" accounts for about twice as much performance loss as the actual shrinking (pruning) does.

Takeaway: Even if you don't shrink the model at all, just training it on your specific data changes it. The shrinking is just the cherry on top of that change.

Summary for the Everyday Reader

If you want to shrink a giant AI to make it fast and cheap:

Don't rush the shrinking: Do it in small steps, not one big jump.
Teach the "Why," not just the "What": If you want the AI to stay smart and not forget general knowledge, teach it to explain its reasoning, not just give the answer.
Mix your grading: Weight the final answer heavily, but don't ignore the reasoning steps.
Accept the drift: The biggest change in the AI comes from teaching it your specific job, not from making it smaller.

The paper provides a clear map (scaling laws) for companies to decide exactly how small they can make their AI before it stops being useful, depending on how much data they have and how much "general smarts" they need to keep.

Technical Summary: Scaling Laws for Task-Specific LLM Distillation

Problem Statement

Large Language Models (LLMs) face significant deployment challenges in latency- and cost-sensitive applications, particularly in domains like quantitative finance. While LLMs achieve strong performance, their scale often precludes efficient deployment. This paper addresses the tradeoff between compressing LLMs for specific tasks and retaining their broad general-knowledge capabilities. The authors investigate how in-domain task performance and general-knowledge retention scale with dataset size, compression ratio, supervision format, and iterative pruning schedules. The study focuses on a 35-class financial news event classification task, a representative production workload, to derive empirical scaling laws that guide practitioners in making domain-specific compression decisions.

Methodology

Data and Setup

The authors constructed a synthetic dataset, FinHeadlineMix, containing 500,000 financial news headlines generated by NVIDIA's Nemotron-3-Nano-30B model. Labels were assigned by a Qwen3-32B teacher model in "thinking mode" (generating Chain-of-Thought traces before the final class). The dataset comprises approximately 2.6 billion training tokens. The teacher model was selected for its comparable performance to larger models (e.g., Llama-3.1-405B) on the in-domain task while maintaining a manageable inference footprint.

Core Techniques

The study evaluates compression via structured pruning (joint depth and width reduction) and knowledge distillation (KD).

Pruning: A lightweight Neural Architecture Search (NAS) selects optimal configurations of depth and width pruning to meet target size constraints while maintaining hardware efficiency (using standard dense kernels).
Distillation Methods: The paper compares four configurations based on two axes:
- Method Complexity: LoRA-based (low-rank adapters) vs. Logit-based (full weight updates via KL divergence).
- Thinking Depth: Label-only supervision vs. Blended Chain-of-Thought (CoT) supervision.
Blended CoT Supervision: A novel contribution is a blended supervision loss that independently weights label tokens and CoT reasoning tokens. This stabilizes KL-divergence distillation over reasoning traces, preventing the gradient from being diluted across diffuse CoT tokens. The loss is formulated as $L = \lambda \ell_{label} + (1-\lambda) \ell_{cot}$ , with $\lambda=0.5$ found to be most stable.
Iterative Pruning and Distillation: The authors employ an iterative strategy where the model is pruned in steps (e.g., 21 percentage points per step) and distilled from the original teacher at each stage. This contrasts with single-step pruning, which often leads to performance collapse at extreme compression ratios.

Evaluation

Models are evaluated on:

In-domain performance: Macro F1, Gold-label Negative Log-Likelihood (NLL), and Brier Score on the financial event classification task.
General-knowledge retention: MMLU and MMLU-Pro benchmarks.
Self-Distillation Baselines: To isolate the cost of compression from the cost of the training process itself, the authors run distillation on unpruned teachers. This establishes a "distribution shift" ceiling, revealing that training-induced performance gaps are roughly twice as large as those caused by compression alone.

Key Contributions

Empirical Scaling Laws: The paper derives scaling laws quantifying how performance scales with dataset size, compression ratio, and supervision format for domain-specific LLM compression.
Supervision Format as a Key Driver: The study identifies that the choice of supervision format (Label-only vs. Blended CoT) is the primary determinant of the tradeoff between in-domain performance and general-knowledge retention.
- Blended CoT KD: Actively recovers general knowledge lost during pruning and improves data efficiency.
- LoRA: Preserves the level of general knowledge surviving pruning but does not actively recover it; it excels in in-domain quality with sufficient data.
- Direct-Label KD: Degrades general knowledge significantly, often collapsing below the pruned baseline at larger data sizes.
Blended CoT Loss: Introduction of a weighted loss function that stabilizes distillation over CoT traces, enabling the use of reasoning supervision in decoder-only LLMs at multi-billion parameter scales.
Iterative Compression Schedules: Demonstration that iterative pruning with decayed schedules (e.g., linear, exponential, cosine) allows for more aggressive compression (down to 16% of original parameters) while maintaining task quality, compared to uniform or single-step pruning.

Results

Data Efficiency: Blended CoT KD is the most sample-efficient method, achieving strong performance with small datasets (e.g., 3k–25k examples). LoRA outperforms Blended CoT KD on in-domain metrics when large datasets (e.g., 400k examples) are available.
General Knowledge Retention:
- Blended CoT KD recovers general knowledge (MMLU scores of 54–55) significantly better than the pruned baseline (40.30).
- LoRA maintains general knowledge near the pruned baseline level.
- Direct-Label KD causes catastrophic forgetting of general knowledge, with MMLU scores collapsing to ~35 at 400k examples.
Compression Limits: Iterative pruning allows compressing the teacher to 16% of its parameters while retaining meaningful task quality (Macro F1 ≈ 0.77). Single-step pruning to 16% results in significant degradation (Macro F1 ≈ 0.52).
Pruning Schedules: Among decayed schedules, Linear decay achieves the best in-domain scores at the 16% target. Exponential decay offers the most stable trajectory and allows for further compression beyond 16% (e.g., to 10–15%) without immediate collapse.
Training vs. Compression Cost: Self-distillation baselines show that the performance gap introduced by the training process (distribution shift) is roughly twice the gap introduced by compression itself.

Significance and Claims

The paper claims to provide a reusable framework for domain-specific compression decisions, grounded in quantitative finance but applicable to other domains. Its primary significance lies in shifting the focus from simple model size reduction to a nuanced understanding of how supervision format and compression scheduling interact.

The authors modestly claim that their findings offer "practical guidance" for practitioners:

Use Blended CoT KD when data is scarce, reasoning is critical, or general-knowledge retention is required.
Use LoRA when large datasets are available and general-knowledge retention is secondary to maximizing in-domain task quality.
Avoid Direct-Label KD if general-knowledge preservation is a priority.
Employ iterative pruning with decay schedules (specifically linear or exponential) to achieve extreme compression (down to 16%) without catastrophic failure.

The study explicitly notes limitations, including the use of a single dense teacher, a specific short-sequence classification task, and the exclusion of quantization or cross-architecture distillation. It positions its work as a foundational step toward understanding scaling laws in task-specific LLM compression, highlighting that the "path" of compression (intermediate architectures) is as critical as the final model size.