Build train steps from objective functions

[seanmor5]

Closes #595.

Axon.Loop.train_step/4 and Axon.Loop.trainer/4 only accepted a loss of the form loss(y_true, y_pred). Anything that needs more than one prediction and its targets, such as a contrastive objective that runs the model on an anchor and a positive, a penalty computed from the parameters, or any objective that wants to see the whole model state, had to re-implement the entire training step. notebooks/vision/metric-learning.livemd did exactly that: sixty lines of batch_step/init boilerplate to express a four-line objective with two forward passes.

Both functions now accept an arity-4 objective function anywhere they accept a loss:

objective = fn forward_fn, model_state, inputs, targets ->
  %{prediction: y_pred} = output = forward_fn.(model_state, inputs)
  penalty = model_state.data["dense_0"]["kernel"] |> Nx.pow(2) |> Nx.sum()
  loss = Axon.Losses.mean_squared_error(targets, y_pred, reduction: :mean)
  {Nx.add(loss, Nx.multiply(1.0e-3, penalty)), output}
end

model
|> Axon.Loop.trainer(objective, :sgd)
|> Axon.Loop.run(data)

The objective receives the model's training-mode forward function, so users never build the model a second time or remember mode: :train, the %Axon.ModelState{} being trained, and the batch. It returns the unscaled loss and the forward output to use as the prediction, which must be the %{prediction: _, state: _} map the forward returns in training mode. That keeps y_pred available to metrics and lets Axon.ModelState.update/3 merge the updated layer state (batch norm statistics, dropout keys) back as before. With several forward passes the user picks which output counts as the prediction. Dispatch is by arity: an arity-4 function is an objective, everything else goes through the existing build_loss_fn/1, which now wraps the loss in the default objective {loss_fn.(targets, prediction), forward_fn.(model_state, inputs)}. The step's return value is checked at trace time, so an objective that returns only the loss raises a clear ArgumentError before anything is compiled.

This also removes the "trainable parameters as grad" hack the issue calls out. The step now differentiates with respect to the whole %Axon.ModelState{} (it is an Nx.Container whose parameters/state/frozen_parameters metadata are kept fields) and selects the trainable subset of the resulting gradient struct with Axon.ModelState.trainable_parameters/1, so the optimizer sees exactly the parameters its state was initialised from. Gradients of frozen parameters and layer state are dead nodes and are eliminated by the compiler; Axon.ModelState.update/3 still passes frozen leaves through untouched, so buffer donation behaves as before. This is what lets an objective read model_state.data[...] directly and have it differentiated. The private tree_merge/3 in Axon.Loop only existed for the hack and is gone.

For loops built from an objective, trainer/4 reports loss from the running average the step keeps in its state instead of recomputing it from y_true/y_pred, which is what build_batch_fn/2 already did for the loss metric of a plain loss.

Limitations

• validate/4 re-attaches training metrics to the evaluator with the default [:y_true, :y_pred] transform. An objective's loss cannot be recomputed that way, so validate/4 now only carries over metrics given as an atom or an arity-2 function and a loop trained with an objective gets validation_<metric> for its other metrics but no validation_loss. Previously any other metric arity crashed with a BadArityError at the end of the first epoch, so this is not a regression. An evaluator/eval_step that accepts an objective is a natural follow-up.
• init_fn now traces the objective rather than just the forward pass to learn the prediction shape, since with a user objective that is the only way to know it. Only shapes are kept, nothing is lowered, but it does mean the targets handed to init_fn must be ones the loss can consume. Two existing tests passed dummy targets of the wrong shape to init_fn directly and were updated; Axon.Loop.run/4 always initialises from a real batch so this does not affect loops.

Docs

train_step/4 gains an "Objective functions" section with the contract and two examples, trainer/4 a "Custom objective" example and the validate/4 note, validate/4 documents which metrics carry over, the custom-loss guide gains a "Using custom objectives in training loops" section, and the metric-learning notebook now uses Axon.Loop.trainer(model, &objective_fn/4, :adam) with an objective that returns the similarity logits as the prediction instead of its hand-written step.

Tests

test/axon/loop_test.exs:

• train_step/3 with an objective that runs two forward passes plus a parameter penalty: step state shape and a decreasing loss over five steps.
• train_step/3 with an objective equivalent to :mean_squared_error produces the same parameters, loss and loss-scale state as the loss form, for both :identity and :dynamic loss scaling.
• The objective receives the training-mode forward: batch-norm running statistics update through an objective.
• An objective returning only the loss raises ArgumentError at init.
• Frozen parameters stay untouched while trainable ones move, exercising trainable_parameters/1 on the gradient struct.
• trainer/3 with an objective runs end to end with a loss and a :mean_absolute_error metric; an arity-3 function is rejected with a message pointing at the objective form.
• validate/4 on an objective loop reports validation_mean_absolute_error and no validation_loss.

All nine fail against main's lib/axon/loop.ex. mix test: 886 passed, 47 excluded. USE_EXLA=1 mix test test/axon/loop_test.exs: 70 passed, which covers the buffer-donation tests through the rewritten gradient path.

🤖 Generated with Claude Code