tf.keras.distillation.LogitsDistillation

Distillation loss that transfers knowledge from final model outputs.

Inherits From: DistillationLoss

tf.keras.distillation.LogitsDistillation(
    temperature=3.0, loss='kl_divergence'
)

This distillation loss applies temperature scaling to the teacher's logits before computing the loss between teacher and student predictions. It's the most common approach for knowledge distillation.

temperature Temperature for softmax scaling. Higher values produce softer probability distributions that are easier for the student to learn. Typical values range from 3-5. Defaults to 3.0.
loss Loss function to use for distillation. Can be:

  • String identifier (e.g., 'kl_divergence', 'categorical_crossentropy')
  • Keras loss instance
  • Nested structure of losses matching the model output structure
  • None to skip distillation for that output (useful for multi-output models where you only want to distill some outputs) At least one loss must be non-None. Defaults to 'kl_divergence'.

Examlpe(s):

# Basic logits distillation with KL divergence
distillation_loss = LogitsDistillation(temperature=3.0)

# With categorical crossentropy loss
distillation_loss = LogitsDistillation(
    temperature=4.0,
    loss="categorical_crossentropy"
)

# With custom loss instance
distillation_loss = LogitsDistillation(
    temperature=4.0,
    loss=keras.losses.CategoricalCrossentropy(from_logits=True)
)

# For multi-output models
distillation_loss = LogitsDistillation(
    temperature=3.0,
    loss=["kl_divergence", "categorical_crossentropy"]
)

# For multi-output models, only distill some outputs
distillation_loss = LogitsDistillation(
    temperature=3.0,
    loss=["kl_divergence", None]  # Skip second output
)

Methods

compute_loss

View source

compute_loss(
    teacher_outputs, student_outputs, **kwargs
)

Compute distillation loss using the configured loss function.

Arguments
teacher_outputs Logits from teacher model. Can be a single tensor, list/tuple of tensors, or dict of tensors.
student_outputs Logits from student model. Can be a single tensor, list/tuple of tensors, or dict of tensors.
**kwargs Additional arguments (ignored).

Returns
Distillation loss tensor.

from_config

View source

@classmethod
from_config(
    config
)

Create instance from configuration.

get_config

View source

get_config()

Get configuration for serialization.

validate_model_compatibility

View source

validate_model_compatibility(
    teacher, student
)

Validate that teacher and student models are compatible.

Arguments
teacher The teacher model.
student The student model.

Raises
ValueError If models are not compatible with this distillation loss.

validate_outputs

View source

validate_outputs(
    teacher_outputs, student_outputs
)

Validate that teacher and student outputs are compatible.

Arguments
teacher_outputs Outputs from the teacher model.
student_outputs Outputs from the student model.

Raises
ValueError If outputs are not compatible.