LayerGenius | Custom Training Steps in Keras 3: One Model, Three Backends

Sooner or later a project needs something model.fit() does not do out of the box: a loss that depends on two forward passes, gradient clipping by a custom rule, adversarial or contrastive training, a GAN, a teacher-student distillation step. In Keras 2 the answer was always the same — subclass Model, override train_step, open a tf.GradientTape. In Keras 3 that code still runs, but only on the TensorFlow backend. The gradient mechanics differ per backend, and if you want the same model to train under JAX or PyTorch, you have to know which parts are portable and which are not.

This post covers the three versions, what stays shared, and how to prove the custom step behaves like the built-in one.

First: do you actually need a custom step?

Most requests we see in review do not need one. Before rewriting the loop, check whether the requirement fits an existing hook:

Extra loss terms that depend on layer activations — self.add_loss(...) inside the layer, no loop change.
Per-sample weighting — the sample_weight argument to fit().
Custom metric bookkeeping — a keras.metrics.Metric subclass.
Schedules, freezing, early stopping, logging — a keras.callbacks.Callback.
Gradient clipping by norm or value — clipnorm / clipvalue / global_clipnorm on the optimizer.

Each of those keeps you inside the compiled fit() path, which gets you distribution strategies, jit_compile, progress bars, and callbacks for free. A custom train_step costs portability and some performance tuning, so spend it only on logic that genuinely cannot be expressed as loss, metric, or callback.

What is portable: the math

The forward pass and loss computation port cleanly across backends if you write them with keras.ops instead of tf.*, jnp.*, or torch.*. keras.ops mirrors the NumPy API plus a keras.ops.nn namespace, and dispatches to whichever backend is active:

import keras
from keras import ops

def distillation_loss(student_logits, teacher_logits, y_true, alpha=0.1, T=4.0):
    hard = keras.losses.categorical_crossentropy(
        y_true, student_logits, from_logits=True
    )
    soft = keras.losses.kl_divergence(
        ops.softmax(teacher_logits / T),
        ops.softmax(student_logits / T),
    ) * (T ** 2)
    return alpha * hard + (1.0 - alpha) * soft

That function runs unchanged under all three backends. Keep every piece of model logic in this style — it is the difference between a one-file backend swap and a rewrite.

What is not portable is how you get gradients and how you apply them. That is the part you write three times, or once for the backend you have committed to.

TensorFlow: the familiar tape

import tensorflow as tf

class Distiller(keras.Model):
    def train_step(self, data):
        x, y = data
        teacher_logits = self.teacher(x, training=False)

        with tf.GradientTape() as tape:
            student_logits = self(x, training=True)
            loss = ops.mean(
                distillation_loss(student_logits, teacher_logits, y)
            )

        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))

        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, student_logits)
        return {m.name: m.result() for m in self.metrics}

Nothing surprising here if you have written Keras 2 code. Note only that metrics are updated explicitly and the return value is a plain dict of scalars.

PyTorch: autograd and zeroing gradients

Under the torch backend, Keras variables are backed by torch tensors and loss.backward() populates .grad. Two details bite people: you must zero gradients yourself, and you should detach the loss before handing it to a metric.

class Distiller(keras.Model):
    def train_step(self, data):
        x, y = data
        teacher_logits = self.teacher(x, training=False)

        self.zero_grad()
        student_logits = self(x, training=True)
        loss = ops.mean(distillation_loss(student_logits, teacher_logits, y))
        loss.backward()

        trainable_weights = [v for v in self.trainable_weights]
        gradients = [v.value.grad for v in trainable_weights]

        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, student_logits)
        return {m.name: m.result() for m in self.metrics}

JAX: stateless, and that changes the signature

JAX is the one that forces a different mental model. JAX transformations require pure functions, so Keras 3 threads all state through explicitly: instead of train_step(self, data), you implement train_step(self, state, data) where state is a tuple of variable value lists, and you return updated values rather than mutating anything.

class Distiller(keras.Model):
    def compute_loss_and_updates(
        self, trainable_variables, non_trainable_variables, x, y, training=False
    ):
        student_logits, non_trainable_variables = self.stateless_call(
            trainable_variables, non_trainable_variables, x, training=training
        )
        teacher_logits = self.teacher(x, training=False)
        loss = ops.mean(
            distillation_loss(student_logits, teacher_logits, y)
        )
        return loss, (student_logits, non_trainable_variables)

    def train_step(self, state, data):
        (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        ) = state
        x, y = data

        grad_fn = jax.value_and_grad(self.compute_loss_and_updates, has_aux=True)
        (loss, (y_pred, non_trainable_variables)), grads = grad_fn(
            trainable_variables, non_trainable_variables, x, y, training=True
        )

        trainable_variables, optimizer_variables = self.optimizer.stateless_apply(
            optimizer_variables, grads, trainable_variables
        )
        # metrics are updated statelessly too; see keras.io for the full pattern
        state = (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        )
        return logs, state

The verbosity is the price of jit-ability, and it is real: a JAX custom step is roughly twice the code of the TF one. The payoff is that the whole step compiles as a single XLA program, which on some models is a meaningful throughput win. Two practical notes. self.stateless_call is how you run the forward pass without touching variables — including BatchNorm moving statistics, which come back in non_trainable_variables and must be threaded through or they silently stop updating. And anything non-JAX inside the step (a Python print, a NumPy call, host-side logging) will either trace once and freeze or break compilation outright.

Choosing: one backend or three

Be honest about the requirement. Most consulting engagements need exactly one backend in production, in which case write one train_step, name the backend in the README, and add a test that fails loudly if keras.backend.backend() is not what the code expects. Multi-backend portability is worth paying for when you are shipping a library others will install, when a hardware decision (TPU via JAX, say) is still open, or when a migration is in progress.

If you do want all three, the maintainable structure is:

One shared module with the model, the loss, and any math — pure keras.ops.
Three thin train_step implementations selected at import time on keras.backend.backend().
One test suite that runs the same tiny fixture through whichever backend is installed.

That keeps the duplicated surface to about thirty lines per backend instead of the whole training stack.

Verify against fit() before you trust it

A custom training step is easy to get subtly wrong — regularization losses dropped, metrics updated with logits where the metric expects probabilities, BatchNorm statistics frozen by accident. Three checks we run every time:

Equivalence test. Implement the ordinary loss as a custom step first, with the custom term disabled (alpha=1.0 in the distillation example). Train a small model for a few steps from a fixed seed with both the custom step and stock fit(). Losses should match to floating-point tolerance. If they diverge, the harness is wrong, not the idea.
Overfit a single batch. Ten samples, no augmentation, no regularization — loss should go to near zero within a couple hundred steps. If it plateaus, gradients are not reaching some variables.
Count trainable variables and check model.losses. Print len(model.trainable_variables) before and after the rewrite, and confirm that any add_loss terms and weight regularizers are still included in the total you differentiate. Dropping them is the single most common bug: training proceeds happily and quietly regularizes nothing.

Also override test_step when you override train_step. Validation that silently uses the default path while training uses the custom one produces two metrics that are not comparable, and the gap gets diagnosed as overfitting for a week.

The trade-off, stated plainly

A custom train_step buys expressiveness and costs portability, some speed on the unoptimized paths, and a maintenance surface that has to be retested at every Keras upgrade. It is the right call for distillation, GANs, contrastive objectives, and anything with multiple coupled optimizers. It is the wrong call for extra loss terms, custom metrics, and schedules, all of which have supported hooks. Reach for it second, and when you do, write the equivalence test before the feature.