Technical Summary: A Study on Hidden Layer Distillation for Large Language Model Pre-Training

Problem Statement

While Knowledge Distillation (KD) is a established paradigm for compressing Large Language Models (LLMs) by transferring knowledge from a large "teacher" to a smaller "student," the majority of current research relies exclusively on output logits. This approach neglects the semantic information contained within the teacher's intermediate hidden representations. Although Hidden Layer Distillation (HLD) has shown promise in encoder architectures and smaller-scale models, its application to decoder-only pre-training at scale remains largely unexplored. Furthermore, existing comparisons often fail to account for the computational overhead introduced by HLD (e.g., storing full activations, regressor costs) or do not rigorously control for hyperparameter tuning, making it difficult to determine if HLD offers genuine advantages over well-tuned logit-based KD.

Methodology

The authors conduct a compute-controlled evaluation of HLD applied to the pre-training of causal LLMs using the Gemma family of models.

• Models:
  • Teacher: A Gemma 3.4B model (34-layer decoder-only Transformer, 2560 embedding dimension) pretrained on 106B unique tokens from the C4 dataset.
  • Students: Two Gemma-based models with standard random initialization: a 123M parameter model (18 layers) and a 735M parameter model (26 layers).
• Distillation Methods Evaluated:
  1. Standard KD (Logit-based): The student minimizes a composite loss of Negative Log-Likelihood (NLL) and Kullback-Leibler (KL) divergence against the teacher's output logits.
  2. Sequential Optimization (HLDF): Adapted from FitNets, this involves a two-phase protocol. Phase 1 aligns the student's intermediate hidden state (at the median layer) with the teacher's using a learned regressor (one-layer dense perceptron) via a normalized Mean Squared Error (MSE) loss. Phase 2 performs standard KD.
  3. Joint Optimization (HLDC): A single-stage training approach that minimizes an aggregated objective function combining NLL, logit distillation, and a hidden layer alignment loss (embedding loss) simultaneously.
• Experimental Rigor:
  • Compute Matching: Comparisons are strictly based on equivalent computational budgets (FLOPs), not just token counts. The authors refine the standard $6ND$ compute approximation to include the significant cost of de-embedding layers and regressor mapping, which is non-negligible for sub-billion parameter models with large vocabularies.
  • Hyperparameter Control: The study avoids "tuning to win" by evaluating methods across a shared grid of hyperparameters (learning rate, temperature, NLL weight) and comparing pointwise results. It also identifies the best configuration for each method to ensure a fair "ceiling-to-ceiling" comparison.
  • Data: Training occurs on the English subset of the C4 dataset (up to 168B tokens for the 735M student).

Key Contributions

1. Systematic Evaluation: The paper provides the first open-source study of HLD specifically during the pre-training phase of decoder-only LLMs, scaling to models with hundreds of millions of parameters.
2. Rigorous Baselines: It establishes a framework that benchmarks methods based on strict computational equivalence (FLOPs), explicitly accounting for the non-negligible costs of de-embedding and loss computation in small models, which previous studies often overlooked.
3. Comprehensive Benchmarking: The evaluation spans multiple downstream tasks (Wikitext 103, HellaSwag, WinoGrande, LAMBADA, PIQA, ARC-E) and analyzes performance across a broad hyperparameter grid rather than relying on a single "best" configuration.

Results

The experiments yield nuanced findings regarding the efficacy of HLD compared to standard KD:

• Downstream Performance: HLD does not consistently outperform standard logit-based KD on downstream evaluation tasks. In the "best hyperparameter" comparison, no HLD variant dominates KD across all benchmarks; differences are marginal and split between methods.
• Perplexity Gains: Despite the lack of downstream superiority, HLDF (Sequential Optimization) yields a systematic, albeit small, improvement in C4 perplexity across all shared-hyperparameter configurations. The mean improvement is approximately 0.26% for both student sizes.
• Context of Gains: The perplexity gains from HLDF are most pronounced in regimes where standard KD performs poorly. In regimes where KD is already efficient, the HLD advantage vanishes.
• Downstream vs. Perplexity: The systematic perplexity gain does not translate into consistent improvements in downstream scores. The authors suggest this indicates that while a latent signal exists in the hidden layers, it may not fundamentally alter the coarse-grained capabilities captured by downstream benchmarks, or that the signal is only useful for fine-grained distributional fitting.

Significance and Claims

The paper concludes that while HLD is an attractive concept for capturing deep internal knowledge, it does not provide a consistent, reliable improvement over a well-tuned logit-based KD baseline when training budgets are strictly equalized.

• Resilience of Logit KD: The findings underscore the resilience of standard logit-based distillation, suggesting that for decoder-only pre-training, the additional complexity and hyperparameter sensitivity of HLD may outweigh its benefits in practical deployment.
• Methodological Discrepancies: The authors attribute the contradiction with prior work (which often reports HLD gains) to four factors:
  1. Compute Accounting: Prior studies often matched token counts rather than FLOPs, ignoring the extra cost of HLD.
  2. Architecture/Phase: Prior work focused on encoders or post-training, where representations are already shaped by pre-training objectives. Causal decoders trained from scratch present a different optimization landscape.
  3. Baseline Tuning: Previous studies often compared HLD against poorly tuned KD baselines.
  4. Scalability: Adapting HLD from small CNNs to billion-parameter transformers requires more profound design choices (losses, adapters) than simple linear regressors.
• Future Outlook: The paper posits that the systematic perplexity gain observed in HLDF suggests a "latent signal" can be extracted, but realizing a breakthrough would likely require significant design innovations or stronger statistical evidence across larger scales and different model families. The authors emphasize the necessity of exact FLOPs accounting when evaluating novel supervision techniques.