跳转至

自定义损失函数

当内置损失都不合适时，forward_backward_custom 让你传入一个对模型对数概率进行运算的任意函数。它比内置 forward_backward 更灵活，但也更慢。

用法

你的函数接收数据与每个 datum 的对数概率张量，返回一个标量损失加上一个可选的 metrics 字典：

def logprob_squared_loss(data, logprobs):
    loss = (logprobs ** 2).sum()
    return loss, {"logprob_squared_loss": loss.item()}

future = training_client.forward_backward_custom(data, logprob_squared_loss)
result = future.result()
print(f"Loss: {result.loss}, Metrics: {result.metrics}")

损失函数也可以同时处理多条序列，例如计算跨序列的方差：

def variance_loss(data, logprobs):
    flat = torch.cat(logprobs)
    var = torch.var(flat)
    return var, {"variance_loss": var.item()}

另一种用途是在偏好的回复（chosen）与不偏好的回复（rejected）之间计算 Bradley-Terry / DPO 损失。内置偏好训练的用法见 DPO。

多目标损失

forward_backward_custom 也支持形状为 (N, K) 的 target_tokens，可以只对特定位置的一组候选 token 定义损失。下面以一道选择题为例，只在四个候选答案上重新归一化。这个运算对目标对数概率是非线性的，因此用自定义损失比 cross_entropy 更直接：

import torch
import mint
from mint import types

messages = [{"role": "user", "content": (
    "What is the capital of France?\n"
    "A) London\nB) Paris\nC) Berlin\nD) Madrid\nAnswer Letter:"
)}]
prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

a = tokenizer.encode(" A", add_special_tokens=False)[0]
b = tokenizer.encode(" B", add_special_tokens=False)[0]
c = tokenizer.encode(" C", add_special_tokens=False)[0]
d = tokenizer.encode(" D", add_special_tokens=False)[0]

target_tokens = torch.zeros(len(prompt_tokens), 4, dtype=torch.long)
target_tokens[-1] = torch.tensor([a, b, c, d])
datum = types.Datum(
    model_input=types.ModelInput.from_ints(prompt_tokens),
    loss_fn_inputs={"target_tokens": target_tokens},
)

CORRECT = 1  # B) Paris

def masked_ce_loss(data, logprobs_list):
    # logprobs_list[0] has shape [N, K=4]
    answer_logprobs = logprobs_list[0][-1]
    answer_logprobs = answer_logprobs - torch.logsumexp(answer_logprobs, dim=-1)
    loss = -answer_logprobs[CORRECT]
    return loss, {"masked_ce": loss.item()}

future = training_client.forward_backward_custom([datum], masked_ce_loss)
result = future.result()

如果你的自定义损失也适用于其他任务，可以告诉我们。我们可能会将它加入内置损失。

底层工作原理

Mint Recursive 不会将你的函数通过 pickle 序列化后发送到服务器。服务端先执行一次前向计算，再对一个加权交叉熵替代损失执行 forward_backward。替代损失产生的梯度与你定义的损失相同。

给定一个非线性损失 loss = f(g(params))，其中 g 返回目标对数概率，Mint Recursive 构造一个在对数概率上线性、但对 params 梯度相同的替代损失：

logprobs = g(params)
surrogate = (logprobs * dL_dLogprobs).sum()

客户端与服务端通过两次计算配合完成训练：

  1. 准备数据（客户端）—— 构建 Datum 列表。
  2. 前向（服务端）—— 计算目标 token 的对数概率。
  3. 自定义损失（客户端）—— loss = custom_fn(logprobs)。
  4. 反向（客户端）—— loss.backward() → grad_outputs。
  5. 前向-反向（服务端）—— 替代损失 sum(logprobs × grad_outputs) → 权重梯度。

由于增加了一次前向计算，forward_backward_custom 的 FLOPs 约为单次 forward_backward 的 1.5 倍，实际耗时最多可达 3 倍。