Knowledge Distillation

How to train smaller SameDiff language models from larger teacher models using cached teacher logits, KL loss, hard-label loss, feature loss, and attention loss.

Knowledge distillation trains a smaller student model to match a larger teacher. In the LLM stack, this can be a practical way to move behavior from a GGUF-imported teacher into a smaller SameDiff graph that is cheaper to serve or fine-tune further.

The main runnable examples are:

• QwenToStudentDistillationExample.java for a real Qwen3.5 teacher imported from GGUF, WikiText-2 token windows, offline FP16 teacher-logit caching, and in-graph Hinton distillation loss.

• DistillationTrainingPipelineExample.java for logit, feature, attention, combined, and self-distillation training loops.

• KnowledgeDistillationExample.java and KnowledgeDistillationConfigExample.java for focused API coverage.

When to Use Distillation

Use distillation when:

• The teacher is too large or slow for your deployment target.

• You want to preserve teacher behavior in a smaller student architecture.

• You have unlabeled text but can run the teacher to produce soft targets.

• You want a student that can later be SFT-tuned or aligned with lower cost.

Do not use distillation as a substitute for data quality. The student learns the teacher distribution on the text you provide. Bad prompts, bad windows, or missing masks produce bad students.

Distillation Modes

Mode	Signal	Use case
Logit KD	Teacher logits or probabilities	Standard LLM distillation; student matches next-token distribution
Hard-label CE	Ground-truth next token IDs	Keeps student anchored to the dataset labels
Feature KD	Teacher/student hidden states	Useful when architectures expose compatible intermediate features
Attention KD	Teacher/student attention maps	Useful for transformer students with attention outputs
Self-distillation	EMA or prior student checkpoint	Stabilizes a model without a separate teacher

QwenToStudentDistillationExample uses the common LLM objective:

loss = alpha * temperature^2 * KL(softmax(teacher / temperature), softmax(student / temperature))
     + (1 - alpha) * CE(labels)

Import the Teacher

A real teacher can be imported from GGUF and driven through SameDiff. The example downloads Qwen3.5-0.8B and imports it for inference:

File ggufFile = LLMModelDownloader
    .download(LLMModel.QWEN35_0_8B, QuantType.Q4_K_M)
    .getModelFile();

GGMLMetadata metadata = GGMLModelImport.inspectModel(ggufFile);
SameDiff teacher = GGMLModelImport.importModel(
    ggufFile.getAbsolutePath(),
    ConversionOptions.forInference());

Tokenizer tokenizer = HuggingFaceTokenizer.fromDirectory(ggufFile.getParentFile());
String teacherLogitsName = teacher.hasVariable("logits") ? "logits" : "lm_logits";

Imported decoder graphs may require more than input_ids. The example builds a complete prefill input map with DecoderInputBuilder, position scalars, KV-cache placeholders, and architecture-specific recurrent state placeholders. When adapting this recipe, inspect the imported graph placeholders and feed all state inputs explicitly.

Build Training Windows

The real example uses WikiText-2 through PerplexityEvaluator.loadWikiText2(), tokenizes text with the teacher tokenizer, and cuts fixed windows:

String wikiText = PerplexityEvaluator.loadWikiText2();
int[] corpusIds = tokenizer.encode(slice, false).getIds();

List<int[][]> trainWindows = cutWindows(corpusIds, 0, trainBatchCount);
List<int[][]> evalWindows = cutWindows(
    corpusIds,
    trainBatchCount * BATCH * (SEQ_LEN + 1),
    evalBatchCount);

Each window should produce:

Batch value	Purpose
input_ids	Student input tokens
labels	Hard next-token labels
loss_mask	1.0 for valid token positions
teacher_logits	Teacher soft targets for the same positions

Cache Teacher Targets

Running the teacher inside every student update is expensive. The production pattern is to precompute teacher outputs once and train the student from cached soft targets.

List<INDArray> trainTeacherLogits = teacherLogitsFor(
    teacher,
    teacherLogitsName,
    trainWindows,
    vocabSize,
    metadata.getHiddenSize());

The example stores logits in FP16 to reduce memory. For larger vocabularies or longer corpora, store top-K logits per token instead of dense [batch, seq, vocab] arrays.

Train the Student

The student graph owns the distillation loss. That means training still uses the standard SameDiff.fit() path:

Adam adamW = Adam.builder()
    .learningRateSchedule(CosineWarmupSchedule.fromRatio(
        3e-4, 3e-5, 0.1, totalSteps))
    .beta1(0.9)
    .beta2(0.999)
    .epsilon(1e-8)
    .build();

student.setTrainingConfig(new TrainingConfig.Builder()
    .updater(adamW)
    .dataSetFeatureMapping("input_ids", "teacher_logits", "loss_mask")
    .dataSetLabelMapping("labels")
    .build());

History history = student.fit(
    new SimpleListMultiDataSetIterator(trainBatches),
    epochs);

This is different from helper APIs that only compute or report distillation loss. If you want the student weights to update, put the distillation loss in the graph and call fit().

DistillationConfig

For reusable trainer-style setups, DistillationConfig describes the loss mixture:

Parameter	Meaning
temperature	Softmax temperature for teacher and student logits
klLossWeight	Weight for logit KL divergence
hardLabelWeight	Weight for ground-truth cross-entropy
featureLossWeight	Weight for hidden-state matching
attentionLossWeight	Weight for attention-map matching

Use higher temperatures to expose more of the teacher distribution over non-argmax tokens. Use hard-label CE when you want the student anchored to the corpus rather than only the teacher.

Feature and Attention Distillation

DistillationTrainingPipelineExample demonstrates additional losses:

• DistillationKLLoss for logit matching.

• FeatureDistillationLoss for hidden representation matching, with projection when teacher and student dimensions differ.

• AttentionDistillationLoss for attention-map matching.

• Combined KD with multiple weighted loss terms.

Feature and attention losses require named intermediate outputs in both models. If the teacher and student dimensions differ, add a projection layer rather than forcing incompatible shapes.

Evaluation

Evaluate both student quality and teacher agreement:

Metric	What it catches
Student held-out masked NLL	Whether the student improved on real held-out text
Student perplexity	Standard LM quality comparison
Teacher top-1 agreement	Whether the student is copying teacher choices
Teacher reference NLL	Whether the teacher targets are sane on the same windows
Side-by-side generation	Qualitative regressions not visible in aggregate metrics
Compression ratio	Parameter reduction compared with teacher

The Qwen example computes before/after student metrics, teacher-student top-1 agreement, teacher reference metrics from cached logits, and teacher/student greedy continuations.

Memory Notes

Dense soft targets can dominate memory. A [batch, seq, vocab] tensor is large for modern tokenizers. The Qwen example keeps batch and sequence length small because a 248k-token vocabulary makes full soft targets expensive.

Use these reductions for larger jobs:

• Cache teacher logits as FP16 or BF16.

• Store only top-K logits and reconstruct sparse KD targets.

• Precompute teacher targets offline and stream them from disk.

• Keep the teacher out of the training process after target generation.

• Combine with Quantized and Memory-Efficient Training for the student optimizer and activation memory.

Run the Examples

cd ../deeplearning4j-examples/samediff-examples
mvn exec:java \
  -Dexec.mainClass="org.nd4j.examples.samediff.quickstart.training.QwenToStudentDistillationExample" \
  -Dexample.distill.batches=2 \
  -Dexample.distill.epochs=1

mvn exec:java \
  -Dexec.mainClass="org.nd4j.examples.samediff.quickstart.training.DistillationTrainingPipelineExample"

Related Pages

• LLM Fine-Tuning Overview

• Instruction SFT, LoRA, and QLoRA

• Preference and RL Alignment

• Quantized and Memory-Efficient Training

• PEFT and RL Alignment Training