Cross-Entropy

监督学习的损失就是普通的交叉熵——目标 token 在学习器 (p_\theta) 下的负对数似然：

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

weights（逐 token 取 0 或 1）决定哪些位置参与计算——通常是 completion token，prompt 被掩掉。

等价代码：

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

Inputs

• target_tokens: array[(N,), int]（top-K 时为 (N, K)）—— 要预测的 token ID。
• weights: array[(N,), float]（或 (N, K)）—— 逐 token 的损失权重。

Outputs

• logprobs —— 学习器对每个请求目标 token 的对数概率。

Diagnostics

• loss:sum —— 加权交叉熵总和。

Top-K 蒸馏

当 target_tokens 形状为 (N, K) 时，损失会取出学习器对每个位置 K 个候选 token 的对数概率，并计算：

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

• 软目标（蒸馏）： 把 (w_{t,k}) 设为教师模型在其 top-K token 上重新归一化后的概率。
• 硬目标： 把全部权重集中在一个候选上，其余为零。

一次蒸馏运行的过程是：从教师采样一段 completion，取回教师在每个采样位置上的 top-K token 分布（通过一次 teacher-forced 的前向，参数为 include_prompt_logprobs=True, topk_prompt_logprobs=K），用 logsumexp 重新归一化这些对数概率，再训练学生去匹配：

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()