Loss functions

Loss functions

Pass the loss name as loss_fn to forward_backward, and any loss-specific arguments as keyword arguments, e.g. model.forward_backward(data, loss_fn="cispo", eps_max=8.0).

The cross-entropy and RL losses below are summed over tokens (not averaged). Each token's contribution is scaled by its per-token field — weights for cross_entropy, advantages for the RL losses. Setting a token's value to 0.0 removes it from the loss; larger magnitudes weight it more. (This is why masking the prompt with 0.0 restricts learning to the completion, and why the loss magnitude grows with the number of contributing tokens.)

What each one is, at a high level:

• cross_entropy — standard supervised next-token loss. Trains the model to imitate the target tokens; weights choose which tokens count (e.g. 1.0 on the completion, 0.0 on the prompt). This is the SFT loss.
• importance_sampling — the basic off-policy policy gradient. Each token's update is scaled by its advantage and by the importance ratio between the current policy and the sampler that produced it (old_logprobs). Simplest RL loss; no clipping, so it can be noisy if the policy drifts far from the sampler.
• ppo — importance sampling with the ratio clipped to [1 − clip_low, 1 + clip_high]. The clip removes the incentive to move any single token's probability too far in one step, which keeps updates stable.
• cispo — clipped importance-sampling policy optimization. Instead of clipping the objective, it caps the importance ratio at eps_max while letting the gradient keep flowing through log π. Robust when samples are well off-policy (the approach used in ScaleRL-style training).
• dro — direct reward optimization. Adds a quadratic anchor (strength beta) pulling the policy toward the sampler's logprobs while it chases reward, trading a bit of reward for stability.

All four RL losses consume the same per-datum fields (old_logprobs + advantages) and differ only in how they shape the update.

Distillation objectives

For sampled reverse-KL feedback, use an RL loss with a detached per-token advantage of teacher_logprob - student_logprob. old_logprobs still records the student sampling policy. The distillation chapter explains the estimator and provides examples for a larger teacher, a self-teacher, and multiple specialists.

cross_entropy also accepts two-dimensional target_tokens and weights per datum, both shaped [sequence_length, K]. At each prediction position, supply K candidate token IDs and their teacher probabilities as weights. This trains against soft targets. Renormalizing over only the top-K tokens changes the teacher distribution; it is an approximation, not full-vocabulary KL. The returned per-position log probability is the weight-normalized mean over candidates, not a vector of their individual log probabilities.

Native self-distillation with opsd

opsd computes full-vocabulary forward KL, KL(teacher || student), from a hint-conditioned teacher using the same model weights. The teacher stream has no gradient. It can also include supervised cross-entropy on separate student positions. It accepts text-only data, not image chunks or an external teacher.

Each stream uses its own sequence length. Targets are next-token IDs aligned with prediction positions. Binary kl_mask fields select paired positions; the active kl_ids must be exactly 0..K-1 in each stream and pair identical target tokens. All masks must be zero at the trailing no-target position. Student loss_mask and kl_mask must be disjoint.

Unlike the token-summed losses above, opsd optimizes the sum of per-row means: each row contributes ce_coef × CE_sum / N_CE plus kl_coef × KL_sum / N_KL, with empty components contributing zero. A positive ce_coef requires SFT tokens in every row. Set ce_coef=0.0 for KL-only training; every row must then contain KL tokens. Coefficients must be finite, nonnegative, and at least one must be positive.

For a mean over B rows, apply gradient_scale=1 / B once at the optimizer step on endpoints supporting gradient_scale_v1. Do not also divide by token count. opsd_objective_sum reports the optimized scalar; opsd_kl and ce_loss report aggregated token means for the components. See the aligned example.

Custom loss functions aren't supported yet. Use one of the built-in losses above; if you need a custom objective, reach out and we'll help.