Custom Loss Functions

When none of the built-in losses fit, forward_backward_custom lets you pass an arbitrary function over the model's log-probs. It is more flexible but slower than the built-in forward_backward.

Usage

Your function takes the data and the per-datum log-prob tensors and returns a scalar loss plus an optional metrics dict:

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}")

You can also write a loss that spans several sequences at once — for example a variance-over-sequences objective:

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

A more practical pairwise case is a Bradley-Terry / DPO loss over chosen-vs-rejected response pairs (see DPO for the built-in preference path).

Multi-target loss

forward_backward_custom also works with (N, K) target_tokens to define a loss over a small candidate set at specific positions. The example below trains on a multiple-choice question by renormalising across just the four answer candidates, not the whole vocabulary — non-linear in the target log-probs, so awkward to express through 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()

If you build a custom loss that seems generally useful, tell us and we may promote it to a built-in.

How it works under the hood

You do not need this to use the API, but for the curious: MinT does not pickle your function or ship it to the server. Instead it splits the gradient into a forward pass plus a forward_backward on a specially weighted cross-entropy surrogate, which produces the exact gradient of your custom loss.

Given a non-linear loss loss = f(g(params)) where g returns the target log-probs, MinT builds a surrogate that is linear in the log-probs but has the same gradient w.r.t. params:

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

The client and server collaborate across two passes:

1. Prepare data (client) — build the Datum list.
2. Forward (server) — compute target-token log-probs.
3. Custom loss (client) — loss = custom_fn(logprobs).
4. Backward (client) — loss.backward() → grad_outputs.
5. Forward-backward (server) — surrogate sum(logprobs × grad_outputs) → weight gradients.

Because of the extra forward pass, forward_backward_custom costs roughly 1.5× the FLOPs and up to 3× the wall time of a single forward_backward.