Loss functions

Built-in losses

Pass the loss name as loss_fn to forward_backward, with loss-specific keyword arguments, for example: model.forward_backward(data, loss_fn="cispo", eps_max=8.0).

Cross-entropy and the four RL objectives below sum token losses. To exclude a token, set its weights value to 0.0 for cross_entropy, or its advantages value to 0.0 for an RL loss. Mask prompt tokens this way to train only on the completion. opsd uses a different reduction, described below.

• cross_entropy trains the model to predict target tokens, scaled by per-token weights.
• importance_sampling scales each token's policy-gradient update by its advantage and the current-to-sampling policy probability ratio. It does not clip the ratio, so updates can be noisy when the policies diverge.
• ppo uses the smaller of the unclipped and clipped advantage-weighted ratios, with clipping bounds [1 − clip_low, 1 + clip_high]. This limits the incentive to move token probabilities too far in one update.
• cispo caps the importance ratio at eps_max and detaches it from the gradient, while differentiating through the current token log probability. It is used in ScaleRL-style training.
• dro adds a quadratic penalty, controlled by beta, for moving away from the sampler's log probabilities.

Distillation objectives

For sampled reverse-KL feedback, use an RL loss with a detached per-token advantage of teacher_logprob - student_logprob. Keep the student sampling policy's log probabilities in old_logprobs. The distillation guide 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-datum means: each datum 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 datum. Set ce_coef=0.0 for KL-only training; every datum must then contain KL tokens. Coefficients must be finite, nonnegative, and at least one must be positive.

For a mean over B datums, 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 losses

Install river-client[custom-loss] and define a scalar PyTorch objective. Call model.prepare_loss() before collecting training data to capture the graph and validate its operations and reductions:

import torch.nn.functional as F
from river_client import CustomLoss


def weighted_ce(logits, inputs):
    nll = F.cross_entropy(logits, inputs["targets"], reduction="none")
    w = inputs["weights"]
    return (nll * w).sum() / w.sum().clamp_min(1.0)


loss = model.prepare_loss(CustomLoss(
    weighted_ce,
    inputs={"targets": "int64", "weights": "float32"},
))
data = [{
    "input_ids": [1, 2, 3, 4],  # Illustrative token IDs.
    "loss_data": {
        "targets": [2, 3, 4, 0],
        "weights": [0.0, 1.0, 1.0, 0.0],
    },
}]
result = model.forward_backward(data, loss_fn=loss)
model.optim_step(lr=1e-4)

Supply unpadded token IDs and the declared columns in loss_data. Targets align with each position's logits; mask unwanted positions explicitly, including the last position when no next-token target exists. This example normalizes by the total weight across the request, including when execution is chunked.

Custom losses with AsyncTrainer

AsyncTrainer prepares custom losses before rollouts and supplies declared targets, old_logprobs, advantages, and mask inputs.

The example below uses CustomLoss to implement SAPO (Soft Adaptive Policy Optimization), which replaces hard clipping with a smooth sigmoid-shaped objective. It averages within each response, then across responses, as in the paper. Use the single-turn engine and rows from the RL recipe.

import torch.nn.functional as F
from river_client import CustomLoss, rl


async def train_custom(engine, rows):
    temperature = engine.temperature

    def custom_loss_inputs(datum):
        inputs = dict(datum["loss_data"])
        mask = inputs.pop("mask")
        length = max(sum(mask), 1)
        return {**inputs, "weights": [m / length for m in mask]}

    def sapo(logits, inputs):
        logprobs = -F.cross_entropy(
            logits / temperature, inputs["targets"], reduction="none"
        ).float()
        ratio = (logprobs - inputs["old_logprobs"]).exp()
        advantages = inputs["advantages"]
        tau_pos, tau_neg = 1.0, 1.05
        positive = advantages.clamp_min(0) * (4.0 / tau_pos)
        negative = advantages.clamp_max(0) * (4.0 / tau_neg)
        surrogate = (
            positive / ((-tau_pos * (ratio - 1.0)).exp() + 1.0)
            + negative / ((-tau_neg * (ratio - 1.0)).exp() + 1.0)
        )
        weights = inputs["weights"]
        return -(surrogate * weights).sum() / weights.sum().clamp_min(1.0)

    trainer = rl.AsyncTrainer(
        engine=engine,
        optimizer=rl.Adam(lr=1e-5),
        completion=rl.GroupCompletion(mode="wait"),
        groups_per_step=2,
        group_size=8,
        max_staleness=2,
        advantage=rl.GroupCentered(standardize=True),
        normalize="sum",
        loss=CustomLoss(sapo, inputs={
            "targets": "int64",
            "old_logprobs": "float32",
            "advantages": "float32",
            "weights": "float32",
        }),
        custom_loss_inputs=custom_loss_inputs,
    )
    async for step in trainer.run(rows, steps=100):
        print(step.n, step.metrics)

Call await train_custom(engine, rows) from your async application. The weights mask non-response tokens and give each response equal weight. tau_pos and tau_neg control SAPO's smoothness; they are separate from sampling temperature.

• Use normalize="sum"; handle normalization, temperature, and clipping in the objective. Mask every loss term.
• Custom losses wait for the full training batch by default, as in this example. Streaming backward is opt-in via forward_backward_groups or forward_backward_batch for supported additive objectives. Server microbatching and loss chunking preserve global normalization.
• All-zero-advantage batches skip training by default. Set train_zero_advantage_batches=True for advantage-independent terms such as CE.
• For extra inputs, pass a custom_loss_inputs callback returning the datum's declared loss_data columns. It runs on a worker thread, so synchronize access to shared mutable state.