Gradient-based methods are foundational in deep learning. The backpropagation algorithm, as popularized in the seminal work by Rumelhart et al. (1986), remains the dominant paradigm for training neural networks. Backpropagation leverages the chain rule of differentiation to compute the gradient of a global loss function with respect to all network parameters efficiently. These gradients are then typically used within variants of stochastic gradient descent (SGD) to iteratively update the parameters and minimize the loss.

To address issues such as slow convergence, navigation of complex loss landscapes, and sensitivity to learning rates, several improvements to vanilla SGD have been proposed. Momentum-based methods, including Polyak’s momentum (Polyak, 1964) and Nesterov accelerated gradient (Nesterov, 1983), incorporate a history of past updates to accelerate progress along consistent descent directions and dampen oscillations. Adaptive methods such as AdaGrad (Duchi et al., 2011), RMSProp (Tieleman, 2012), and Adam (Kingma and Ba, 2014) adjust learning rates on a per-parameter basis, adapting to the geometry of the loss landscape and helping networks train more robustly, especially with sparse gradients. For large models, memory-efficient variants of backpropagation have also been developed, e.g., reducing activation storage or recomputation overheads (Korthikanti et al., 2023; Xu and others, 2025).

Despite their success, gradient-based methods face several challenges. They can converge to suboptimal local minima or encounter prevalent high-error saddle points, particularly in high-dimensional non-convex optimization problems (Dauphin et al., 2014; Choromanska et al., 2015). Furthermore, they can suffer from vanishing or exploding gradients in deep networks (Hochreiter, 1991; Bengio et al., 1994), require the objective function and network components to be differentiable (or sub-differentiable), and can be sensitive to hyperparameter choices. The computational cost and sequential nature of computing and propagating gradients through deep networks can also be significant bottlenecks for parallelization across layers. These limitations motivate the exploration of alternative training schemes.

A direct alternative is to forgo gradient calculation entirely and use derivative-free optimization (DFO) or zeroth-order methods. These methods typically rely only on evaluations of the objective function (the global loss). Evolution Strategies (ES) (Rechenberg, 1978; Schwefel and Schwefel, 1977) are a prominent example. ES optimizes parameters by sampling points in the parameter space around the current estimate, evaluating the loss at these points, and moving the estimate in a direction informed by these evaluations, effectively approximating a gradient direction or directly seeking improvement (Salimans et al., 2017). While computationally intensive due to the need for multiple forward passes per update, ES has shown surprising effectiveness, particularly in reinforcement learning domains (Salimans et al., 2017). Other population-based methods like Genetic Algorithms (GAs) (Holland, 1992) and Particle Swarm Optimization (PSO) (Kennedy and Eberhart, 1995) maintain a population of candidate solutions (network parameters) and iteratively refine them using operators inspired by biological evolution (selection, crossover, mutation) or social behavior (particle movement). The objective is typically to maximize a fitness function directly related to the global loss. Evolutionary approaches have also been applied to meta-learning problems, such as optimizing learning rules or hyperparameters (Schmidhuber, 1987). The primary objective of these DFO methods is the minimization of the global loss function. However, they often exhibit high sample complexity compared to gradient-based methods, requiring many loss evaluations, which can be prohibitive for large networks and datasets. Scalability to extremely high-dimensional parameter spaces remains a challenge.

Drawing inspiration from neuroscience, another class of methods employs local learning rules that do not require global error backpropagation. Rooted in neuroscientific theory, Hebbian learning follows the principle “cells that fire together, wire together.” This principle, introduced by Donald Hebb in his 1949 book (Hebb, 1949), states that the synaptic strength between two neurons increases if they are co-activated. Early neural network research used simplified Hebbian rules for unsupervised representation learning. Modern variants, such as Oja’s rule (Oja, 1982) and the Bienenstock–Cooper–Munro (BCM) rule (Bienenstock et al., 1982), build on the classical Hebbian framework to ensure weight stabilization and learn principal components or other low-dimensional representations. Relatedly, attractor networks like the Hopfield network (Hopfield, 1982) also utilize Hebbian-style rules to store patterns as stable states (attractors) within an energy function, suitable for tasks like associative memory and pattern completion. Another important family of models leveraging local updates includes energy-based models like Boltzmann Machines (BMs) and particularly their simplified variant, Restricted Boltzmann Machines (RBMs) (Hinton, 2002). RBMs are stochastic networks trained to model a probability distribution over their inputs, typically via Maximum Likelihood. While exact gradient calculation is intractable for general BMs, RBMs can be effectively trained using algorithms like Contrastive Divergence (CD) (Hinton, 2002), which relies on Gibbs sampling and local computations between connected layers to approximate the gradient. This made RBMs instrumental as building blocks for Deep Belief Networks (Hinton et al., 2006), offering a biologically more plausible route to unsupervised feature learning and generative modeling. A common appealing property of these various local learning rules, whether Hebbian or energy-based, is their independence from global gradient signals. Instead, they rely on updates based on local neuronal activity (e.g., correlations, timing, sampling statistics) or principles like energy minimization. The objective is often related to maximizing correlation, capturing variance, achieving representational stability, minimizing a network energy function, or modeling data distributions, rather than directly minimizing a global supervised loss function for input-output mappings. This locality can be advantageous for large-scale or biologically plausible architectures where global signals are unavailable or costly. However, directly applying or adapting these principles to match the performance of gradient-based methods on complex supervised tasks remains challenging.

Inspired by the brain’s event-driven architecture, neuromorphic computing focuses on spiking neural networks (SNNs) that communicate via discrete spike events instead of continuous-valued activations. Conventional backpropagation becomes difficult in spiking systems due to the non-differentiable nature of spike functions. As a result, a variety of alternative learning rules, such as Spike-Timing-Dependent Plasticity (STDP) in spiking neural networks (Song et al., 2000; Gerstner and Kistler, 2002), have emerged. STDP aligns closely with Hebbian-like principles, adjusting synaptic strengths based on the relative timing of pre- and post-synaptic spikes. The objective here is local synaptic modification driven by temporal correlations. Neuromorphic chips like IBM’s TrueNorth (Merolla et al., 2014) and Intel’s Loihi (Davies et al., 2018) have been designed to implement such models efficiently, potentially offering significant energy savings. Because learning in SNNs often occurs locally within each synapse, neuromorphic approaches can circumvent the need for global gradient signals. However, training SNNs for complex, real-world tasks comparable to those solved by deep learning remains an active research area, and performance often lags behind conventional networks. Such methods also inherently handle non-differentiable spike events, a property shared with the projection-based methods discussed in this work.

Several approaches attempt to retain some benefits of gradient-based learning while avoiding full backpropagation, often motivated by biological plausibility (e.g., the weight transport problem) or computational considerations. Target Propagation (TP) (Bengio, 2014; Lee et al., 2015) aims to compute layer-specific targets instead of gradients. An inverse mapping or autoencoder associated with each layer generates target activations that, if achieved, would reduce the global loss. The layer then updates its weights locally to better map its input to these targets. The objective is local: minimize the mismatch between layer outputs and computed targets, serving as a proxy for the global loss. TP can avoid propagating precise gradients but requires learning or defining inverse mappings, which can be challenging and potentially unstable. Equilibrium Propagation (Scellier and Bengio, 2017) is another line of work in this direction: it computes updates from differences between free and weakly clamped network states in energy-based models, avoiding explicit reverse-mode backpropagation through layers. Feedback Alignment (FA) (Lillicrap et al., 2016) replaces the transposed weight matrices used in backpropagation’s backward pass with fixed, random feedback matrices. Surprisingly, the network can learn by aligning its forward weights to leverage these random pathways for error signaling. The objective remains the global loss, but the gradient information is approximated. FA solves the weight transport problem but typically learns slower and may achieve lower final performance compared to backpropagation. Direct Feedback Alignment (DFA) (Nøkland, 2016) further simplifies this by sending the error signal directly from the output layer to each hidden layer via fixed random matrices. Decoupled Neural Interfaces (synthetic gradients) (Jaderberg et al., 2017) similarly reduce strict sequential dependencies by learning gradient predictors for intermediate modules. Layer-wise training provides another way to decouple learning. Initially popular for pre-training Deep Belief Networks (Hinton et al., 2006), this involves training layers sequentially, often using unsupervised objectives like reconstruction error before potential end-to-end fine-tuning. The objective is local per layer/stage. While useful for initialization, purely greedy layer-wise training may not yield globally optimal solutions for the final task. More recently, the Forward-Forward algorithm (Hinton, 2022) was proposed as a potential alternative inspired by biology. It discards backpropagation entirely, using two forward passes—one with positive (real) data and one with negative data—and updating weights based on a local goodness metric specific to each layer. The objective is layer-local: maximizing goodness for positive samples and minimizing it for negative samples. This avoids backpropagation but requires generating negative data and doubles the computation per update compared to a single forward pass. Its scalability and performance across diverse tasks are still under investigation. Further exploring alternatives to full backpropagation, Radhakrishnan et al. (2024) recently introduced the Average Gradient Outer Product (AGOP) as a backpropagation-free mathematical mechanism to characterize and enable feature learning. Their work demonstrates that AGOP captures learned features across diverse architectures and can instill feature learning capabilities in models like kernel machines, notably through their Recursive Feature Machine (RFM) algorithm.

Beyond heuristic or bio-inspired approaches, alternative mathematical optimization frameworks have been applied to neural network training, often by reformulating the learning problem. The Alternating Direction Method of Multipliers (ADMM) (Boyd et al., 2011) is a general framework for constrained optimization that decomposes a large problem into smaller, potentially easier subproblems that are solved iteratively. It has been explored for training neural networks, sometimes by introducing auxiliary variables and constraints to decouple layers or enforce structure (Glowinski and Marroco, 1975; Taylor et al., 2016). The Method of Auxiliary Coordinates (MAC) represents another such decomposition strategy for deeply nested systems (Carreira-Perpiñán and Wang, 2014). The objective in these frameworks is typically the original global loss, subject to reformulations. ADMM can handle non-differentiable regularizers and constraints but requires careful problem formulation, and solving the subproblems efficiently can be challenging, often involving overhead from dual variable updates or the solution of complex subproblems like large matrix inversions (Taylor et al., 2016). Block Coordinate Descent (BCD) methods optimize the network parameters block by block (e.g., layer by layer) while keeping other parameters fixed (Bertsekas, 1997). Each subproblem optimizes the global loss with respect to a subset of variables. BCD can be simpler to implement and potentially more memory-efficient than SGD for certain structures, but its convergence can be slow, dependent on the block partitioning strategy, and it may struggle with highly correlated parameters.

Crucially, the perspective of training as a feasibility problem (finding parameters ${\bm{\theta}}$ such that ${\bm{f}}({\bm{x}}_{i},{\bm{\theta}})={\bm{y}}_{i}$ for all $i$) provides a distinct reformulation suitable for iterative projection methods like AP or DR. This “learning without loss” approach was conceptually explored by Elser (2021), who proposed using the Difference Map algorithm (related to DR) to find feasible points in the intersection of constraints derived from individual data samples. Elser demonstrated the concept with illustrative examples. The projection-based method detailed in our work builds directly on this feasibility perspective, leveraging a fine-grained decomposition based on the computation graph’s primitive functions, which allows for efficient, parallelizable projection steps.

A less common direction treats network training or design as a discrete or symbolic problem. By formulating aspects of a neural network’s parameters or structure (e.g., activation functions, connectivity) as a discrete optimization problem, one can sometimes leverage mature combinatorial solvers like SAT solvers or integer programming techniques. These approaches often arise in the context of “neuro-symbolic AI,” aiming to integrate neural learning with symbolic reasoning. The objective might be to find a network satisfying certain logical constraints or optimizing a discrete objective. However, these methods often suffer from severe scalability challenges as network dimensions grow and are typically applicable only to specific problem types or network architectures.

PJAX is built upon core components implemented using JAX primitives. Accessed only indirectly via the high-level API, these components form the foundational building blocks, implementing the essential primitive functions and projection operators required for the optimization process:

Users interact with PJAX through a high-level API designed to closely resemble the JAX / NumPy interface, facilitating adoption for those already familiar with these libraries.

The pjax.optim module contains the algorithms that solve the feasibility problem defined by the computation graph. Currently available optimizers include Alternating Projections (AP), Cyclic Projections (CP), Douglas–Rachford (DR), and the Difference Map algorithm (Elser, 2003). The optimizer’s update function accepts:

• •

  A user-defined Python function that computes the desired output using PJAX operations. This function implicitly defines the computation graph and constraints. It is analogous to the function one might pass to jax.grad, often representing the forward pass of a model, ${\bm{\theta}}\mapsto{\bm{f}}({\bm{x}},{\bm{\theta}})$, whose output is then constrained according to the graph (e.g., via Output Constraint nodes).

• •

  A dictionary (or other pytree structure) holding the current state of the Parameter objects (e.g., network weights).

The optimizer then applies the function to the parameters, computes and stores intermediate outputs, forms a bipartition of the graph, and performs a specified number of update steps using the chosen projection algorithm. Finally, it returns a dictionary containing the updated consensus values for the parameters.

Inspired by Flax, pjax.nn simplifies model definition and training via a compositional Module interface for reusable components (e.g., layers, blocks). It automates parameter handling (initialization, naming, sharing), enabling users to define complex architectures with familiar concepts, which are readily compatible with the pjax.optim module for projection-based training.

To illustrate how a neural network is defined using the pjax.nn API and how PJAX subsequently constructs a detailed computation graph, we consider an MLP with a single hidden layer and ReLU activation. Mathematically, this MLP is defined as

where ${\bm{x}}$ is the input, ${\bm{W}}^{\text{hidden}}$ and ${\bm{b}}^{\text{hidden}}$ are the weights and biases for the hidden layer, respectively, and ${\bm{W}}^{\text{out}}$ represents the weights for the output layer. The corresponding PJAX code for this MLP is shown in .

When this MLP model is instantiated and applied to an input batch (e.g., 32 MNIST vectors, each $784$-dimensional, with a hidden layer of $16$ features and $10$ output classes), PJAX traces the operations. The resulting computation graph, shown in Fig. 7, visualizes this trace.

In this graph, nodes represent PJAX components. Array nodes hold constant data like the input batch and target data for the cross_entropy loss. Learnable parameters, whose names are assigned by our pjax.nn API, are represented as Parameter nodes; for this MLP, these include hidden.weight (corresponding to ${\bm{W}}^{\text{hidden}}$), relu.bias (corresponding to ${\bm{b}}^{\text{hidden}}$ and applied at the sum_relu stage), and out.weight (corresponding to ${\bm{W}}^{\text{out}}$). The primitive functions shown are dot (scalar dot product) and sum_relu. Shape transformation nodes, specifically reshape and repeat, prepare input tensors for the batched scalar operations. For instance, to implement the batched matrix multiplication in the first linear layer (${\bm{W}}^{\text{hidden}}{\bm{x}}$), the input Array (32,784) and the hidden.weight (784,16) parameter are reshaped and repeated to match dimensions for the dot operation.

This graph details the concrete sequence of operations as traced by PJAX. In contrast, a conceptual illustration (as in Fig. 2) would represent each neuron’s compound operations (dot product and activation) as distinct explicit nodes with edges fanning out to all neurons in the subsequent layer, rather than utilizing PJAX’s explicit shape transformations for vectorized scalar primitives. Another minor technical difference is that the conceptual graph does not include a node for the target data, rather, it’s implicitly represented in the loss function.

Below, we detail and derive the projection operators onto the graphs of several primitive functions used within the PJAX framework. Recall that for a function $f:\mathbb{R}^{n}\to\mathbb{R}^{m}$, its graph is defined as $\operatorname{Graph}(f)=\{({\bm{x}},f({\bm{x}}))\mid{\bm{x}}\in\mathbb{R}^{n}\}$. The projection onto this graph, $\mathrm{P}_{\operatorname{Graph}(f)}({\bm{x}}_{0},{\bm{y}}_{0})$, finds the point on the graph closest to $({\bm{x}}_{0},{\bm{y}}_{0})$.