Mega Context

Knowledge Distillation

7 min read

Knowledge Distillation (arXiv:1503.02531) — Report

PDF: Knowledge Distillation - 1503.02531.pdf

Overview

  • Title: Distilling the Knowledge in a Neural Network
  • Authors: Geoffrey Hinton, Oriol Vinyals, Jeff Dean (Google)
  • Year: 2015
  • Core Idea: Train a smaller “student” model to mimic a large “teacher” model (or ensemble) by learning from its softened output distributions rather than hard labels, transferring “dark knowledge” about class relationships that improve generalization.

Core Concepts

Temperature-Scaled Softmax

The key innovation is using a “temperature” parameter T when computing softmax during distillation:

q_i = exp(z_i / T) / Σ_j exp(z_j / T)
  • T = 1: Standard softmax (hard targets)
  • T > 1: “Softens” the distribution, revealing subtle similarities between classes
  • During training: High T for distillation loss, T=1 for hard label loss
  • At inference: Student uses T=1

Distillation Loss

Combined objective function:

L_total = α * L_soft + (1-α) * L_hard

where:
  L_soft = KL_divergence(student_softT || teacher_softT)
  L_hard = Cross_entropy(student, true_labels)

Typical α ≈ 0.9, heavily weighting the soft teacher targets.

Dark Knowledge

The “dark knowledge” transferred includes:

  • Inter-class similarities: Which wrong answers are “less wrong”
  • Implicit regularization: Smoothed distributions prevent overfitting
  • Learned structure: Patterns the teacher discovered but aren’t in hard labels

Specialist Models

Paper also introduces training specialists on confusable subsets, then distilling them into a single generalist model—relevant for handling class imbalance and focusing compute on hard examples.

Relevance to MegaContext

Hierarchical Compression Alignment

Knowledge distillation’s core principle—preserving output distributions while reducing model capacity—directly mirrors MegaContext’s GistNet objective:

  • Teacher = base model processing full token sequences
  • Student = gist embeddings replacing those sequences
  • Soft targets = continuation distributions at horizon H
  • Distillation loss ≈ ΔNLL@H (measuring distribution mismatch)

GistNet Training already implements a variant of this: we train gists to minimize divergence between base model predictions with/without compression. The temperature-scaling insight suggests we could:

  1. Soften ΔNLL targets during GistNet early training to expose subtle dependencies
  2. Curriculum schedule that gradually lowers temperature as compression improves
  3. Multi-scale distillation where LOD2 gists learn from LOD1’s “soft” representations

LensNet Training Signal

The paper’s observation that soft targets carry richer training signal than hard labels applies to LensNet:

  • Instead of binary “expand/collapse” labels, we compute distributional counterfactuals
  • LensNet Scoring could benefit from temperature-scaled utility distributions when the optimal action is ambiguous
  • Multi-task distillation (expansion utility + collapse utility + keep utility) mirrors the specialist framework

Ensemble Compression Pathway

The specialist→generalist pipeline maps to potential GistNet training strategies:

  1. Train multiple GistNet specialists on different domains (code, narrative, structured data)
  2. Distill them into a single unified GistNet using their soft predictions as targets
  3. Maintains domain-specific sensitivity while reducing deployment complexity

This aligns with Track B goals for domain-adaptive gist encoding.

Regularization Benefits

Distillation’s implicit regularization (preventing overconfident predictions) helps address GistNet failure modes:

  • Repetitive gist degeneracy: Softened targets discourage collapsed modes
  • Overfitting to training distribution: Inter-class similarities preserve generalization
  • Catastrophic forgetting: Temperature annealing during continual learning

What We Can Use

1. Temperature-Scheduled GistNet Training

Modify GistNet Training to use temperature-scaled ΔNLL:

def distillation_loss(gist_logits, token_logits, T=2.0):
    """
    Compute KL divergence between gist-conditioned and token-conditioned
    next-token distributions at temperature T.
    """
    gist_probs = F.softmax(gist_logits / T, dim=-1)
    token_probs = F.softmax(token_logits / T, dim=-1)
    return T**2 * F.kl_div(gist_probs.log(), token_probs, reduction='batchmean')

Schedule: Start T=3.0 (very soft), anneal to T=1.0 (hard) over training. Early phases focus on coarse structure; later phases refine precise substitutability.

2. Soft Counterfactual Labeling for LensNet

Extend Counterfactual labeling with distributional targets:

# Instead of single scalar utility:
u_expand = -ΔNLL_expand(entry)  # Hard target
 
# Use softened distribution over utilities:
def soft_utility_target(entry, T=2.0):
    """
    Returns probability distribution over {collapse, keep, expand}
    based on temperature-scaled ΔNLL measurements.
    """
    utilities = {
        'collapse': -ΔNLL_collapse(entry),
        'keep': 0.0,
        'expand': -ΔNLL_expand(entry)
    }
    return softmax(utilities, T=T)

Teaches LensNet to handle uncertain situations where multiple actions have similar utility.

3. Specialist Ensemble for Domain Coverage

Train domain-specific GistNet variants, then distill:

  1. Specialist A: Code-focused (trained on GitHub, StackOverflow)
  2. Specialist B: Narrative-focused (trained on BookSum, PG19)
  3. Specialist C: Structured data (tables, JSON, lists)
  4. Generalist: Unified GistNet distilled from A+B+C ensemble

Each specialist runs during distillation phase; generalist deploys solo. Captures diverse compression strategies without mixture-of-experts overhead.

4. Confidence Calibration via Temperature

Use learned per-block temperature to signal gist confidence:

  • High-confidence gists (low perplexity contexts) → T ≈ 1.0
  • Low-confidence gists (rare patterns, code switches) → T ≈ 3.0

This metadata feeds into LensNet as auxiliary conditioning, helping it prioritize expansion of uncertain regions.

Limitations & Risks

1. Teacher Quality Dependence

Distillation cannot recover information the teacher never learned. For MegaContext:

  • Base model is frozen → gist quality ceiling set by base model’s original training
  • Domain shift (e.g., specialized jargon) may produce poor teacher signals
  • Mitigation: Curate diverse distillation corpus, use multiple teacher checkpoints

2. Capacity Bottleneck

Student must have sufficient capacity to capture teacher’s soft distributions. For GistNet:

  • 32→1 compression is aggressive (96.875% reduction)
  • Single-vector gists may underfit complex 32-token semantics
  • Mitigation: Two-stage refinement (32→1→32→1), contrastive regularization

3. Mode Collapse Risk

Soft targets can over-smooth, leading to generic/repetitive outputs:

  • Gists might learn “average” representations that work everywhere but excel nowhere
  • Risk increases with very high temperatures (T > 5)
  • Mitigation: Entropy regularizers, contrastive loss to maintain diversity

4. Temperature Sensitivity

Optimal T varies by task and data:

  • Narrative text: T ≈ 2-3 (benefits from smoothing)
  • Code: T ≈ 1.5-2 (needs precision)
  • Structured data: T ≈ 1-1.5 (schema-sensitive)

Fixed temperature in POC Implementation may underperform. Need adaptive T per block type or learned per-domain temperatures.

Potential Follow-Up Reading

Extensions & Variants

  • FitNets (Romero et al., 2015): Distill intermediate representations, not just outputs → could guide GistNet internal layers
  • Attention Transfer (Zagoruyko & Komodakis, 2017): Match attention patterns between teacher/student → verify gist-replaced attention stays aligned
  • Born-Again Networks (Furlanello et al., 2018): Self-distillation improves same-capacity models → iterative GistNet refinement

Theoretical Foundations

  • Label Smoothing (Szegedy et al., 2016): Connects distillation to regularization theory
  • Dark Knowledge Analysis (Müller et al., 2019): Formal analysis of what soft targets encode

Applications to Compression

  • Pruning + Distillation (Polino et al., 2018): Combines structural and functional compression
  • Quantization-Aware Distillation: Distill to lower-precision models while preserving accuracy

Open Questions for MegaContext

1. Adaptive Temperature Scheduling

How to automatically tune T during GistNet Training?

  • Per-block adaptive T based on source entropy?
  • Learned temperature predictor as auxiliary head?
  • Domain-specific temperature lookup (code vs prose vs tables)?

2. Multi-Horizon Distillation

Current ΔNLL@H uses fixed horizon H. Could we distill at multiple horizons simultaneously?

L_multi = Σ_h w_h * ΔNLL@h

Where w_h weights balance short-term (h=32) vs long-term (h=128) fidelity. Teaches gists to preserve both local coherence and global structure.

3. Bidirectional Distillation

Paper focuses on forward predictions. For GistNet, could we also distill backward language model predictions?

  • Train gists to preserve both future and past context
  • Symmetric loss: L = ΔNLL_forward + ΔNLL_backward
  • Ensures gists capture sufficient information for bidirectional reasoning

4. Cross-Level Consistency

Should LOD2 gists distill from LOD1’s soft outputs, or directly from LOD0 tokens?

  • Cascaded: LOD2 learns from LOD1 (faster, compounds errors)
  • Direct: LOD2 learns from LOD0 (slower, more faithful)
  • Hybrid: Multi-task loss with both signals

5. Distillation for Focus Allocation

Could we distill an optimal Focus Allocator policy from expensive search procedures?

  1. Run exhaustive beam search over focus configurations offline
  2. Collect (context, optimal_focus) pairs
  3. Distill into fast LensNet + greedy allocator
  4. Captures complex utility landscapes without online search cost

Core Architecture

Training & Optimization

Metrics & Evaluation

  • Gist Tokens - Attention-masked compression (complementary to distillation)
  • LLMLingua-2 - Token importance via teacher distillation
  • LoRA - Parameter-efficient fine-tuning (potential student architecture)

Future Work

Graph View

Backlinks