> For the complete documentation index, see [llms.txt](https://deeplearning4j.konduit.ai/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://deeplearning4j.konduit.ai/en-1.0.0-rewrite/deeplearning4j/overview-4/knowledge-distillation.md).

# 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:

```java
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:

```java
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.

```java
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:

```java
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](/en-1.0.0-rewrite/deeplearning4j/overview-4/quantized-training.md) for the student optimizer and activation memory.

## Run the Examples

```bash
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](/en-1.0.0-rewrite/deeplearning4j/overview-4/fine-tuning.md)
* [Instruction SFT, LoRA, and QLoRA](/en-1.0.0-rewrite/deeplearning4j/overview-4/instruction-sft-lora.md)
* [Preference and RL Alignment](/en-1.0.0-rewrite/deeplearning4j/overview-4/preference-rl-alignment.md)
* [Quantized and Memory-Efficient Training](/en-1.0.0-rewrite/deeplearning4j/overview-4/quantized-training.md)
* [PEFT and RL Alignment Training](/en-1.0.0-rewrite/deeplearning4j/peft-and-rl.md)