Cross-Entropy¶

The supervised learning loss is plain cross-entropy — the negative log-likelihood of the target tokens under the learner (p_\theta):

[ \mathcal{L}(\theta) = -\mathbb{E}x[\log p\theta(x)] ]

weights (0 or 1 per token) selects which positions contribute — typically the completion tokens, with the prompt masked out.

Equivalent code:

# weighted, per-token NLL, then sum
elementwise_loss = -target_logprobs * weights
loss = elementwise_loss.sum()

Inputs

target_tokens: array[(N,), int] (or (N, K) for top-K) — the IDs to predict.
weights: array[(N,), float] (or (N, K)) — per-token loss weight.

Outputs

logprobs — the learner's log-prob for each requested target token.

Diagnostics

loss:sum — total weighted cross-entropy.

Top-K distillation¶

When target_tokens has shape (N, K), the loss extracts the learner's log-probs for K candidate tokens per position and computes:

[ \mathcal{L}{\text{CE-topK}}(\theta) = -\sum) ]}\sum_{k=1}^{K} w_{t,k} \cdot \log p_\theta(x_{t,k

Soft targets (distillation): set (w_{t,k}) to the teacher's probability renormalised over its top-K tokens.
Hard targets: assign all the weight to one candidate and none to the others.

A distillation run samples a completion from a teacher, recovers the teacher's top-K token distribution for each sampled position (via a teacher-forced pass with include_prompt_logprobs=True, topk_prompt_logprobs=K), renormalises those log-probs with logsumexp, and trains the student to match:

import torch
from transformers import AutoTokenizer
import mint
from mint import types

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-235B-A22B")
K = 20

sc = mint.ServiceClient()
teacher = sc.create_sampling_client(base_model="Qwen/Qwen3-235B-A22B").result()
student = sc.create_lora_training_client(base_model="Qwen/Qwen3-4B").result()

messages = [{"role": "user", "content": "Write an efficient Fibonacci function."}]
prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

# 1. Sample a completion from the teacher.
resp = teacher.sample(
    prompt=types.ModelInput.from_ints(prompt_tokens),
    num_samples=1,
    sampling_params=types.SamplingParams(max_tokens=512, temperature=0.7),
).result()
sampled = list(resp.sequences[0].tokens)

# 2. Teacher-force the completion to get top-K logprobs at each position.
tf_resp = teacher.sample(
    prompt=types.ModelInput.from_ints(prompt_tokens + sampled),
    num_samples=1,
    sampling_params=types.SamplingParams(max_tokens=1),
    include_prompt_logprobs=True,
    topk_prompt_logprobs=K,
).result()
rows = tf_resp.topk_prompt_logprobs[len(prompt_tokens):]

teacher_tokens = torch.tensor([[tid for tid, _ in r] for r in rows])      # [seq, K]
teacher_logprobs = torch.tensor([[lp for _, lp in r] for r in rows])      # [seq, K]
teacher_logprobs -= torch.logsumexp(teacher_logprobs, dim=-1, keepdim=True)

# 3. Build the student Datum (shift by one for teacher forcing).
student_input = prompt_tokens + sampled[:-1]
gen_start = len(prompt_tokens) - 1
target_tokens = torch.zeros(len(student_input), K, dtype=torch.long)
weights = torch.zeros(len(student_input), K)
target_tokens[gen_start:gen_start + len(rows)] = teacher_tokens
weights[gen_start:gen_start + len(rows)] = teacher_logprobs.exp()

datum = types.Datum(
    model_input=types.ModelInput.from_ints(student_input),
    loss_fn_inputs={"target_tokens": target_tokens, "weights": weights},
)
fb = student.forward_backward([datum], loss_fn="cross_entropy").result()