Loss Functions - Tinker Documentation

Loss Functions

Tinker provides built-in loss functions for supervised learning and reinforcement learning. You select a loss by passing a string to forward_backward:

future = await training_client.forward_backward_async(data, loss_fn="cross_entropy")
result = await future.result_async()

How it works

Each training example is a Datum — a single sequence with all the information needed to compute the loss. A Datum contains:

The key design principle: everything the loss needs is in the Datum. This means each Datum is self-contained — no batch-level state, no external lookups. The simplest example is cross-entropy for SFT:

import tinker
from tinker import types

# A single training example: predict target_tokens from input_tokens
datum = types.Datum(
    model_input=types.ModelInput.from_ints(input_tokens),  # shape: (N,)
    loss_fn_inputs={
        "target_tokens": target_tokens,  # shape: (N,) — what to predict at each position
        "weights": weights,              # shape: (N,) — 0 for prompt, 1 for completion
    }
)

# forward_backward computes the loss and gradients in one call
future = await training_client.forward_backward_async([datum], loss_fn="cross_entropy")
result = await future.result_async()
print(f"Loss: {result.loss}")

For RL losses, the Datum also includes the sampling log-probabilities and advantages:

rl_datum = types.Datum(
    model_input=types.ModelInput.from_ints(tokens),    # shape: (N,)
    loss_fn_inputs={
        "target_tokens": target_tokens,                 # shape: (N,)
        "weights": weights,                             # shape: (N,)
        "logprobs": sampling_logprobs,                  # shape: (N,) — from the rollout policy
        "advantages": advantages,                       # shape: (N,) — reward signal per token
    }
)

future = await training_client.forward_backward_async([rl_datum], loss_fn="importance_sampling")

forward_backward returns a ForwardBackwardOutput with output tensors in result.loss_fn_outputs (e.g., the model's logprobs for each token).

At a glance

Loss Use case Key idea
cross_entropy Supervised learning Maximize log-probability of target tokens
importance_sampling RL (policy gradient) Correct for off-policy sampling with p/qp/qp/q ratio
ppo RL (clipped) Clip the p/qp/qp/q ratio to prevent large updates
cispo RL (clipped grad) Clip the ratio but use it as a gradient coefficient
dro RL (off-policy) Quadratic penalty on policy divergence
forward_backward_custom Any Write arbitrary loss over logprobs

Notation

We denote the training model as pθp_{\theta}pθ​, the sampling distribution as qqq, and advantages as AAA. For notation simplicity we omit the query and denote the full model completion sequence of tokens as xxx.

All losses are applied at the token level. Unless noted otherwise, tensors have shape (N,) where N is model_input.length. They can be provided as numpy.ndarray or torch.Tensor, and the return values will use the same tensor type.

Additional notes on RL losses: