Knowledge Distillation › Worked Example: Distillation on Sentiment Classification

Michael BrenndoerferFebruary 24, 20264 min read
Language AI HandbookAll chapters
Part 41 of 67Model Compression
Chapter 307Knowledge Distillation › Worked Example: Distillation on Sentiment Classification
Reading section 5 of 10 · Worked Example: Distillation on Sentiment Classification · 4 minSection 5 of 10 · 4 min

Book contents

Section 5 of 10 · 4 min read

Knowledge Distillation › Worked Example: Distillation on Sentiment ClassificationOne section · Sections in this chapter5 / 10 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's make these concepts concrete with a worked example. We'll simulate the distillation process on a sentiment classification task: binary classification of movie reviews into positive and negative. This is not an LLM example, but it clearly illustrates the mechanics before we scale up to the language model setting. The core insight, that soft labels encode more information than hard labels, is identical in both cases.

Suppose the teacher produces these logits for three reviews:

Teacher predictions on three example reviews. The ambiguous review (row 3) receives logits close to zero, signaling uncertainty about the correct classification.
ReviewTeacher logit (neg)Teacher logit (pos)True label
"Brilliant film"-3.23.2Positive
"Utter garbage"2.8-2.8Negative
"It was fine, I guess"0.10.2Positive

The first two reviews are clear-cut. The teacher is confident in both directions. Its soft labels at temperature τ=2\tau = 2 will be similar to hard labels: most of the probability mass will sit on the correct class, with a small residual on the wrong class. The third review is ambiguous. "It was fine, I guess" has a slightly positive teacher logit of 0.2 versus 0.1 for negative, but the difference is tiny. The teacher is essentially saying "I cannot tell."

Now consider what happens at training time. Hard-label training penalizes the student for assigning any probability to the negative class on review 3. The student learns that "fine, I guess" is unambiguously positive. This is technically correct based on the label, but it is overfitting to a noisy label. The true sentiment of that phrase really is ambiguous. A student trained on the teacher's soft labels at τ=2\tau = 2 learns that assigning nearly equal probabilities to both classes is acceptable for that phrase. It will generalize better to similar hedged expressions.

In[7]:
Code
import numpy as np


def softmax_with_temperature(logits, temperature=1.0):
    logits = np.array(logits, dtype=float)
    scaled = logits / temperature
    scaled -= scaled.max()
    exp_vals = np.exp(scaled)
    return exp_vals / exp_vals.sum()


def cross_entropy(target, predicted):
    predicted = np.clip(predicted, 1e-12, 1.0)
    return -np.sum(target * np.log(predicted))


# Teacher logits: [neg, pos] for each review
teacher_logits = np.array(
    [
        [-3.2, 3.2],  # Clearly positive
        [2.8, -2.8],  # Clearly negative
        [0.1, 0.2],  # Ambiguous
    ]
)

hard_labels = np.array(
    [
        [0, 1],
        [1, 0],
        [0, 1],
    ]
)

tau = 2.0

# Compute soft labels (teacher distribution at temperature tau)
soft_labels = np.array(
    [softmax_with_temperature(logits, tau) for logits in teacher_logits]
)

# Simulate a student that has some uncertainty
# Student outputs closer to 50/50 on the ambiguous example
student_logits = np.array(
    [
        [-1.8, 1.8],
        [1.6, -1.6],
        [0.05, 0.15],
    ]
)

student_soft = np.array(
    [softmax_with_temperature(logits, tau) for logits in student_logits]
)

# Compute losses
hard_losses = [
    cross_entropy(hard_labels[i], softmax_with_temperature(student_logits[i]))
    for i in range(3)
]
soft_losses = [cross_entropy(soft_labels[i], student_soft[i]) for i in range(3)]
Out[8]:
Console
Review                          Hard loss  Soft loss (tau=2)   Teacher soft label
----------------------------------------------------------------------------------
Review 1 (clear positive)          0.0270             0.2235  [0.039, 0.961]
Review 2 (clear negative)          0.0400             0.2756  [0.943, 0.057]
Review 3 (ambiguous)               0.6444             0.6928  [0.488, 0.512]

For the two clearly labeled reviews, hard and soft losses are similar: the teacher is confident, and the soft label closely resembles the hard label. For the ambiguous third review, the hard loss penalizes the student's uncertainty, while the soft loss accepts it. This is exactly the desired behavior: distillation avoids forcing overconfidence on uncertain examples, and the student learns to be appropriately uncertain where the teacher is appropriately uncertain.

The teacher's soft label for review 3 is approximately [0.47,0.53][0.47, 0.53]: nearly equal probabilities for both classes. This is the teacher's honest assessment of that example. A student trained on this soft label learns that equal uncertainty is acceptable for hedged language. A student trained on the hard label [0,1][0, 1] is pushed toward false confidence.