Custom Loss Functions
For use cases outside of the built-in losses, Tinker provides the more flexible (but slower) method forward_backward_custom to compute a more general class of loss functions.
Usage
A custom loss function receives the data and the model's logprobs, and returns a scalar loss plus optional metrics:
def logprob_squared_loss(data: list[Datum], logprobs: list[torch.Tensor]) -> tuple[torch.Tensor, dict[str, float]]:
loss = sum((lp ** 2).sum() for lp in logprobs)
return loss, {"logprob_squared_loss": loss.item()}
Call it with forward_backward_custom:
future = training_client.forward_backward_custom(data, logprob_squared_loss)
result = await future
print(result.metrics)
You can also define loss functions which operate on multiple sequences at a time. For example, a loss function that computes the variance across the sequences (although practically useless) can be implemented as:
def variance_loss(data: list[Datum], logprobs: list[torch.Tensor]) -> tuple[torch.Tensor, dict[str, float]]:
flat_logprobs = torch.cat(logprobs)
variance = torch.var(flat_logprobs)
return variance, {"variance_loss": variance.item()}
A more practical use case would be to compute a Bradley-Terry loss on pairwise comparison data -- a classic approach in RL from human feedback, as introduced and popularized by Learning to Summarize. Similarly, we can also implement Direct Preference Optimization, which also computes a loss involving pairs of sequences; see the DPO & Preferences tutorial for more details.
Multi-target loss
You can also use forward_backward_custom with multi-target target_tokens to define losses over a small candidate set at specific sequence positions. For example, the following masked cross-entropy loss trains a model on a multiple-choice question by renormalizing over four answer-letter tokens instead of the full vocabulary. Because this renormalization is nonlinear in the extracted target-token logprobs, it's somewhat cumbersome to express in terms of the cross_entropy primitive:
import torch
import tinker
messages = [{"role": "user", "content": (
"What is the capital of France?\n"
"A) London\n"
"B) Paris\n"
"C) Berlin\n"
"D) Madrid\n"
"Answer Letter:"
)}]
prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
# Extract token IDs for the four answer letters.
a_token = tokenizer.encode(" A", add_special_tokens=False)[0]
b_token = tokenizer.encode(" B", add_special_tokens=False)[0]
c_token = tokenizer.encode(" C", add_special_tokens=False)[0]
d_token = tokenizer.encode(" D", add_special_tokens=False)[0]
target_tokens = torch.zeros(len(prompt_tokens), 4, dtype=torch.long)
target_tokens[-1] = torch.tensor([a_token, b_token, c_token, d_token])
datum = tinker.Datum(
model_input=tinker.ModelInput.from_ints(prompt_tokens),
loss_fn_inputs={"target_tokens": target_tokens},
)
CORRECT_ANSWER_IDX = 1 # B) Paris
def masked_ce_loss(data, logprobs_list):
# logprobs_list[i] has shape [N, K=4] for each datum
# Renormalize over just the four answer tokens instead of the full vocabulary.
answer_logprobs = logprobs_list[0][-1] # shape [K=4]
answer_logprobs = answer_logprobs - torch.logsumexp(answer_logprobs, dim=-1)
loss = -answer_logprobs[CORRECT_ANSWER_IDX]
return loss, {"masked_ce": loss.item()}
future = training_client.forward_backward_custom(
[datum],
masked_ce_loss,
)
result = await future
If you're using a custom loss function that you think is generally useful, please let us know, and we'll add it to the list of built-in loss functions.
forward_backward_custom returns an APIFuture immediately, so you can submit an optimizer step before awaiting it.
Recover with a new client after a submitted custom-loss failure
If a submitted client-side custom loss operation fails or is cancelled, its training client rejects queued and subsequent training requests. Correct the callback and resume from a saved checkpoint with a new client.
How forward_backward_custom works
The following describes passing a Python callback, as in the examples above. Tinker does NOT pickle your function or send it to the server. Instead, Tinker decomposes the gradient computation into a forward call followed by a forward_backward call on an appropriately designed weighted cross-entropy loss, which lets it compute exactly the right gradient.
Mathematically, this works as follows. First, consider the full nonlinear loss function:
We construct a loss function that is linear in the logprobs, but has the same gradient with respect to params as the full nonlinear loss:
logprobs = compute_target_logprobs(params)
surrogate_loss = (logprobs * logprob_grads).sum()
# where logprob_grads = dLoss/dLogprobs
Here's what happens under the hood — the client and server collaborate in two passes:
This client-side callback path does an extra forward pass, so it requires 1.5x as many FLOPs as a single forward_backward. It'll also take up to 3x as long (wall time), due to implementation details of how forward_backward operations are scheduled.