Loss Functions

MinT ships built-in losses for supervised learning and reinforcement learning. You pick one by passing its name to forward_backward:

future = training_client.forward_backward(data, loss_fn="cross_entropy")
result = future.result()

How a loss reads its data

Every training example is a Datum — one token sequence plus everything the loss needs to compute a scalar. A Datum carries:

• model_input — the token sequence the model sees.
• loss_fn_inputs — a dict of per-token tensors the loss consumes (targets, weights, advantages, sampling log-probs, …).

The design rule is: the loss only reads what is in the Datum. There is no batch-level state and no external lookup, so each example is self-contained.

SFT cross-entropy only needs targets and weights:

import mint
from mint import types

datum = types.Datum(
    model_input=types.ModelInput.from_ints(input_tokens),
    loss_fn_inputs={
        "target_tokens": target_tokens,   # what to predict at each position
        "weights": weights,               # 0 on the prompt, 1 on the completion
    },
)
future = training_client.forward_backward([datum], loss_fn="cross_entropy")
result = future.result()
print(f"Loss: {result.loss}")

RL losses additionally carry the rollout's sampling log-probs and the per-token advantages:

rl_datum = types.Datum(
    model_input=types.ModelInput.from_ints(tokens),
    loss_fn_inputs={
        "target_tokens": target_tokens,
        "weights": weights,
        "logprobs": sampling_logprobs,   # recorded by the sampler
        "advantages": advantages,        # reward signal per token
    },
)
future = training_client.forward_backward([rl_datum], loss_fn="importance_sampling")

forward_backward returns a ForwardBackwardOutput; per-token outputs land in result.loss_fn_outputs (for example the learner's log-probs).

Choosing a loss

Loss	Use case	Key idea
cross_entropy	Supervised learning	Maximise log-prob of the target tokens
importance_sampling	RL (policy gradient)	Reweight by the (p/q) ratio to correct off-policy samples
ppo	RL (clipped)	Clip the (p/q) ratio to bound the update size
cispo	RL (clipped coefficient)	Clip the ratio but apply it as a gradient coefficient
dro	RL (off-policy)	Add a quadratic penalty on policy divergence
forward_backward_custom	Anything else	Write an arbitrary loss over log-probs

Notation

Throughout these pages: (p_\theta) is the learner policy, (q) is the sampling policy, and (A) is the per-token advantage. We write the full completion sequence as (x).

All losses run at the token level. Tensors have shape (N,) where N is model_input.length, and may be numpy.ndarray or torch.Tensor — the return value uses whichever type you passed in.

On RL losses generally:

• The formulations are intentionally general; you generate the data and compute advantages in your own code. A typical setup samples several rollouts per prompt and centers their advantages (the GRPO pattern, Shao et al., 2024).
• The built-in policy-gradient objectives do not add a standalone KL term (the original GRPO KL term has been noted to be mathematically inconsistent; see Zhang et al., 2025). If you want KL regularisation, fold it into the reward — that is mathematically correct.
• Token-level losses are summed over the sequence. Want a different aggregation? Bake it into the advantage tensor.