> ## Documentation Index
> Fetch the complete documentation index at: https://docs.baseten.co/llms.txt
> Use this file to discover all available pages before exploring further.

# Loss functions

> Loss functions supported by the Loops trainer, with data shapes and minimal snippets.

Select a loss function with the `loss_fn` argument on [`forward_backward()`](/reference/sdk/loops/training-client). Each loss reads a different set of per-position fields from `Datum.loss_fn_inputs`, and some take a `loss_fn_config` dict to tune hyperparameters.

## `cross_entropy`

Standard supervised fine-tuning. Each `Datum` carries per-token targets under `target_tokens`, with `-100` marking positions the loss should ignore (typically the prompt). `forward_backward` does not shift labels internally, so shift when you tokenize: the label at position `i` is the token at position `i+1`.

```python theme={"system"}
from baseten.loops import Datum, ModelInput, TensorData

datum = Datum(
    model_input=ModelInput.from_ints(tokens),
    loss_fn_inputs={
        "target_tokens": TensorData(
            data=targets, dtype="int64", shape=[len(targets)]
        ),
    },
)

fb = training_client.forward_backward(data=[datum]).result(timeout=600.0)
```

`cross_entropy` takes no `loss_fn_config`. See the [quickstart](/loops/quickstart) for the full tokenization step.

### Weight individual positions

Pass an optional `weights` tensor to scale the per-position loss. Each position contributes `-log(p(target)) * weight`, so a weight of `0.0` drops a position and a larger weight increases its influence. A `-100` target drops the position outright; use `weights` to scale what remains. Build `weights` alongside `target_tokens` so they align with the shifted labels:

```python theme={"system"}
from baseten.loops import Datum, ModelInput, TensorData

def build_weighted_datum(tokenizer, prompt, answer):
    p = tokenizer.encode(prompt, add_special_tokens=False)
    a = tokenizer.encode(answer, add_special_tokens=False)
    tokens = (p + a)[:-1]
    targets = [-100] * (len(p) - 1) + list(a)
    # One weight per target: ignore the prompt, train on the answer.
    weights = [0.0] * (len(p) - 1) + [1.0] * len(a)
    return Datum(
        model_input=ModelInput.from_ints(tokens),
        loss_fn_inputs={
            "target_tokens": TensorData(
                data=targets, dtype="int64", shape=[len(targets)]
            ),
            "weights": TensorData(
                data=weights, dtype="float32", shape=[len(weights)]
            ),
        },
    )

datum = build_weighted_datum(
    training_client.get_tokenizer(),
    prompt="What is the capital of France?\nAnswer:",
    answer=" Paris",
)
fb = training_client.forward_backward(data=[datum]).result(timeout=600.0)
```

Set `weights` on every `Datum` in the batch or on none of them.

### Keep loss inputs one-dimensional

Keep `target_tokens` and `weights` one-dimensional, one value per token position, with `shape` `[N]` where `N` is the number of tokens in `model_input`. Loops reads the flat `TensorData.data` buffer and ignores any second dimension, so it does not support the two-dimensional `(N, K)` form Tinker uses for top-K soft targets.

## `importance_sampling`, `ppo`, `dppo`, `cispo`, `dro`

On-policy reinforcement learning. Each `Datum`'s `loss_fn_inputs` must include per-position `logprobs` (from the behavior policy at rollout time) and `advantages` (from your reward model or return-to-go).

```python theme={"system"}
L = len(tokens)

datum = Datum(
    model_input=ModelInput.from_ints(tokens),
    loss_fn_inputs={
        "logprobs": TensorData(data=behavior_logprobs, dtype="float32", shape=[L]),
        "advantages": TensorData(data=advantages, dtype="float32", shape=[L]),
    },
)

fb = training_client.forward_backward(
    data=[datum],
    loss_fn="ppo",  # or "importance_sampling"
).result(timeout=600.0)
```

These losses take the same `loss_fn_inputs` shape and differ in clipping and regularization: `importance_sampling` doesn't clip the ratio; `ppo` clips it; `dppo`, `cispo`, and `dro` add further variants. Drive them from a rollout loop that samples from a paired [`SamplingClient`](/reference/sdk/loops/sampling-client) and scores each trajectory before calling `forward_backward`.

<Note>
  The trainer accepts `dppo` and `dpo`, but they aren't in the SDK's `LossFnType` literal yet, so type-checked code needs `# type: ignore[arg-type]` on those calls.
</Note>

## `dpo`

Direct Preference Optimization. Trains on chosen/rejected preference pairs against a frozen reference policy. DPO adds the following steps to SFT:

1. **Snapshot the reference before any `optim_step`.** The initial weights are the reference for offline DPO.
2. **Interleave the batch** as `[chosen_0, rejected_0, chosen_1, rejected_1, ...]`. Pairing is by position, and the SDK rejects odd-length batches before submitting anything.
3. **Score reference logprobs once** with [`attach_reference_logprobs`](/reference/sdk/loops/helpers) and reuse across every step.

Each `Datum` needs the following `loss_fn_inputs` fields. Build `target_tokens` and `weights` yourself when tokenizing; `ref_logprobs` comes from [`attach_reference_logprobs()`](/reference/sdk/loops/helpers):

* `target_tokens`: labels shifted by one, `-100` over the prompt positions.
* `weights`: `0.0` over the prompt, `1.0` over the response. The DPO log-ratio sums logprobs weighted by this mask.
* `ref_logprobs`: frozen reference logprobs, one per position. Populated automatically by `attach_reference_logprobs`.

```python theme={"system"}
from baseten.loops import AdamParams, Datum
from baseten.loops.helpers.datum import attach_reference_logprobs

# 1. Snapshot the reference policy. This blocks until the snapshot is
#    published and returns a SamplingClient directly.
reference = training_client.save_weights_and_get_sampling_client(name="dpo-ref")

# 2. Build interleaved chosen/rejected datums with target_tokens and
#    weights set as above, then score them once.
datums = attach_reference_logprobs(datums, reference)

# 3. Train.
fb = training_client.forward_backward(
    data=datums,
    loss_fn="dpo",
    loss_fn_config={"beta": 0.1},  # optionally: "label_smoothing": 0.0
).result(timeout=600.0)
training_client.optim_step(AdamParams(learning_rate=1e-5)).result(timeout=600.0)
```

**`loss_fn_config` keys**:

* `beta` (float, default `0.1`): KL-penalty strength against the reference.
* `label_smoothing` (float, default `0.0`, range `[0, 0.5]`): enables conservative DPO (cDPO).

**Metrics on `fb.metrics`**: the DPO pass reports `loss`, the mean DPO loss over the batch. A fresh policy starts near `ln 2` (about `0.693`) and decreases as the policy learns to prefer chosen responses.

## Next steps

* [`TrainingClient` reference](/reference/sdk/loops/training-client) for the full `forward_backward()` signature.
* [Helpers reference](/reference/sdk/loops/helpers) for `attach_reference_logprobs`.
* [Loops concepts](/loops/concepts) explains the paired trainer + sampler model that makes reference-policy snapshots possible.