PJAX

Training without backpropagation

For forty years, neural networks have been trained by gradient descent. Pick a differentiable loss, compute its gradient with backpropagation, take a step in the negative gradient direction, and repeat. Every modern deep learning framework, from PyTorch to JAX, is built around this loop. It works very well, and it is also the only training paradigm most practitioners ever encounter.

This work asks a basic question: what does it look like to train a neural network without descending a loss? We do not mean a new optimizer or a clever surrogate gradient, but a different formulation of what training is.

We build an alternative on a classical problem in optimization, the feasibility problem: find a point that lies in several constraint sets at once. When each set is easy to project onto but their intersection is not, the classical approach is a projection method, which repeatedly projects onto the individual sets. A projection maps the current point to a nearest point of the set. Projection methods go back to von Neumann's method of alternating projections and are widely used in signal and image reconstruction.

We reformulate neural network training as a feasibility problem and solve it approximately with iterative projections. The updates need no backpropagation, and they run in parallel across the network.

Training as a feasibility problem

We write the network as a computation graph of simple operations and introduce one variable for every edge, the value it carries. Each node then imposes a condition on its edges. An input node fixes its edges to the data, and a parameter node requires all copies of its parameter to agree. An operation node requires its output to equal the operation applied to its inputs, as in the dot-product set \{(x, w, h) \mid \langle x, w\rangle = h\}. A loss node requires the prediction to minimize the loss for its example, which for the squared loss means matching the label. For a batch, the graph contains one copy of the network per example, and all copies share the parameter nodes.

A state that satisfies every condition gives parameters that minimize the training loss. Such a state exists only if one set of parameters fits every example, so in practice we look for a state that satisfies the conditions approximately. For losses such as cross-entropy, whose minimum is not attained at finite outputs, the loss node takes a proximal step instead of enforcing a hard condition.

Solving it with projections

Projection methods such as alternating projections and Douglas–Rachford project onto one set at a time and repeat. For closed convex sets with a common point, they converge to a point in the intersection. Some of our sets, such as those of dot products and ReLUs, are not convex, so here the methods are heuristics. The same holds in phase retrieval, another nonconvex problem, where they have long worked well in practice.

Each condition involves only the edges of a single node, so its projection is a small local problem. For an operation f, it moves the current input and output values to a nearest pair (x, f(x)). Most of these projections are cheap. They have a closed form or, for the dot product, need a few Newton steps on a scalar equation. The price of locality is that information moves only between neighboring nodes in each projection step, so deep networks need many iterations.

Nodes that share no edge can be projected at the same time. Neural network computation graphs are often bipartite, as in the figure above, and any graph can be made bipartite by inserting identity nodes. One iteration therefore projects all nodes of one partition at once and then all nodes of the other. Backpropagation, in contrast, is sequential across layers. On a small MLP, the step time grows more slowly with depth than for SGD.

A projection needs only the graph of f, which is well defined whether or not f is differentiable. Quantization, max-pooling and margin constraints all have exact projections, so they need no surrogate gradients or straight-through estimators.

For each batch, we initialize the edge values with a forward pass from the current parameters, run a few projection iterations, and read the parameters off the parameter edges. As with stochastic gradient descent, each batch is only partly solved, so the parameters move a little toward fitting each batch. No optimizer state is carried between batches.

theta = init_params()
while not converged:
    batch = sample_batch(data)
    G     = build_computation_graph(loss_fn, batch)   # nodes + edges
    z     = forward_pass(G, theta)                    # value on every edge, from current params

    for _ in range(K):                                # projection iterations
        z = project(G, z, partition="A")              # all nodes of A in parallel
        z = project(G, z, partition="B")              # then all nodes of B

    theta = read_parameter_edges(z)

PJAX: the projection analogue of autodiff

Much of backpropagation's practical success comes from autodiff frameworks such as PyTorch and JAX. With them, implementing a new architecture or training method only requires its forward computation, and the framework computes the gradients. A new training paradigm, such as training with projections, needs a similar framework. We therefore built PJAX on top of JAX. It traces the forward computation, builds the feasibility problem, and runs the projection method. Its API closely follows JAX and NumPy. On top of it, we implemented a higher-level library similar to Flax. Layers such as linear, convolutional and embedding layers are modules that hold their own parameters, and networks are built by stacking them.

Autodiff needs a derivative rule for each primitive operation, and PJAX needs a projection. A few primitives, namely dot products, sums, maxima and ReLU, are enough for the MLPs, CNNs and simple RNNs we train, because these networks are built from them. PJAX also provides projections for quantization and the losses. A matrix product is a batch of dot products. A convolution extracts the image patches under the kernel and takes their dot products with the kernel weights.

Patch extraction, reshapes, transposes, indexing and padding change where values are stored, but not the values themselves. They need no projection of their own. PJAX passes values through them in both directions, with the forward transform and its inverse. Where a transform copies a value, as overlapping patches do, the way back averages the copies. Autodiff treats them the same way: the derivative rule of a reshape just reshapes the gradient. Adding a new operation to PJAX requires only its forward evaluation and its projection.

Does it work?

We trained feedforward, convolutional and recurrent networks on image classification, particle physics and character-level language modeling, and compared against standard training with gradient methods.

The honest answer is yes, within limits. On shallow MLPs, Douglas–Rachford comes close to SGD on MNIST, with 95.8% test accuracy against 96.8%. On deeper networks and on CNNs, SGD performs notably better. No projection method matches Adam in any setting we tried.

One of the main limitations is memory. Every use of a parameter is a separate edge with its own value, so the state holds a copy of the parameters for every example in the batch, and in CNNs and RNNs also for every position or time step. Gradient methods keep the parameters in memory only once.

Concluding thoughts

Gradient descent dominates deep learning because it works, not because it is the only way to train a network. Part of the reason it works so well is that the field has converged on it and built everything else around it. Architectures and their components, such as normalization layers and skip connections, as well as optimizers, were developed together with gradient-based training. If we replace only the training algorithm in such a jointly tuned system, it is not surprising that the results fall behind. Architectures designed for projection-based training might be much more competitive. Finding them requires a framework in which new operations and models are easy to implement. This is what we want to provide with PJAX.