🤖 machine learning

A projection-based framework for gradient-free and parallel learning

This paper introduces PJAX, a JAX-based framework that reformulates neural network training as a parallelizable, gradient-free feasibility problem using iterative projection operators, offering a compelling alternative to conventional gradient-based optimization with advantages in handling non-differentiable operations and enabling massive parallelism.

Original authors: Andreas Bergmeister, Manish Krishan Lal, Stefanie Jegelka, Suvrit Sra

Published 2026-04-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Andreas Bergmeister, Manish Krishan Lal, Stefanie Jegelka, Suvrit Sra

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive, complex puzzle, like a giant 3D jigsaw or a Rubik's Cube, but you don't know the final picture.

The Old Way (Gradient-Based Learning)
Currently, most AI models learn using a method called "backpropagation." Think of this like a hiker trying to find the bottom of a foggy valley (the best solution). The hiker feels the slope under their feet (the gradient) and takes a step downhill. They keep doing this, step by step, until they reach a low point.

The Problem: Sometimes the hiker gets stuck in a small dip (a local minimum) that isn't the true bottom. Sometimes the path is so steep or flat that the hiker gets lost or moves too slowly. Also, to know which way is "down," the hiker has to send a signal all the way back from the bottom of the valley to the top, which is slow and requires a very specific, symmetrical path.

The New Way (Projection-Based Learning)
The authors of this paper propose a completely different strategy. Instead of trying to find the bottom of a valley, they treat training as a feasibility problem.

Imagine you have a room full of walls, each with a specific rule.

Wall A says: "The red block must be next to the blue block."
Wall B says: "The green block must be above the red block."
Wall C says: "The total weight must equal 50kg."

Your goal isn't to slide down a hill; it's to find one single arrangement of blocks where every single wall's rule is satisfied at the same time.

How It Works: The "Projection" Metaphor
The authors call their method "Projection-Based." Here is how they do it:

Break it Down: They break the giant puzzle (the neural network) into tiny, simple pieces called "primitive functions" (like simple math operations: adding numbers, multiplying them, or deciding if a number is positive).
The Local Fix: Instead of looking at the whole puzzle, they look at just one wall (one rule). If the blocks don't fit that wall's rule, they "project" the blocks onto the wall. Imagine shining a light on the blocks; the shadow they cast on the wall is the "correct" position for that specific rule.
Parallel Power: This is the magic part. Because each wall only cares about its own immediate neighbors, you can fix Wall A, Wall B, and Wall C all at the same time. You don't have to wait for Wall A to finish before starting Wall B. This is like having a team of 100 people fixing different parts of a house simultaneously, rather than one person fixing the roof, then the kitchen, then the bathroom, one by one.
Repeat: They do this over and over. Each time, they nudge the blocks to fit the local rules better. Eventually, the blocks settle into a position where they satisfy all the rules simultaneously. That is your trained AI.

Why This Is Cool (According to the Paper)

No "Slope" Needed: You don't need to calculate a "slope" (gradient). This means you can use rules that are "bumpy" or broken (non-differentiable), like a switch that is either ON or OFF. The old method struggles with these; this new method handles them easily.
Biological Plausibility: In the brain, neurons don't send a global "error signal" all the way back from the end of a thought to the beginning. They just adjust based on what their immediate neighbors are doing. This new method mimics that local, neighbor-to-neighbor adjustment.
Speed: Because everyone works in parallel, it can be much faster on modern computer chips (GPUs/TPUs) that are designed to do many things at once.

The Trade-Off: The "Memory" Cost
The paper admits there is a catch. To do this, the computer has to remember the position of every single "edge" in the puzzle at every step.

Analogy: In the old method, you just remember the current location of the hiker. In this new method, you have to remember the position of every single block in the room, and every single connection between them, for every single person in your team.
Result: This uses a lot more computer memory (RAM). The authors had to shrink some of their test models to fit them into their computer's memory, whereas the old method could handle larger models more easily.

The Results
The authors built a software tool called PJAX (Projection JAX) to test this. They tried it on different types of puzzles:

Simple patterns (MLPs)
Image recognition (CNNs)
Language prediction (RNNs)

They found that while the "old way" (using Adam or SGD optimizers) is still the champion for raw speed and final accuracy in many cases, this new "Projection" way works surprisingly well. It is a viable alternative that:

Learns without needing gradients.
Handles "bumpy" rules that confuse other methods.
Learns very efficiently on parallel hardware, especially for tasks like language modeling where the old method struggles with "vanishing gradients" (forgetting the beginning of a sentence).

In Summary
The paper says: "Stop trying to slide down a hill to find the answer. Instead, treat the problem as a set of local rules. Fix each rule locally and simultaneously, and eventually, the whole system will fall into place." It's a new way to train AI that is more parallel, more flexible with different types of math, but currently requires more memory.

1. Problem Statement

The paper addresses fundamental limitations in the dominant paradigm of deep learning training: gradient-based optimization via backpropagation. While highly effective, this approach faces several critical challenges:

Differentiability Requirement: It requires all network components and loss functions to be (sub)differentiable, excluding non-differentiable operations (e.g., quantization, discrete logic).
Sequential Bottleneck: The backward pass is inherently sequential, limiting parallelism and increasing update latency, especially in deep architectures.
Biological Implausibility: The mechanism of symmetric weight transport (sending error gradients back through the exact same weights used for forward propagation) is considered biologically unrealistic.
Optimization Pathologies: Gradient methods often struggle with vanishing/exploding gradients, local minima, and saddle points.

The authors propose a paradigm shift: reformulating neural network training not as a loss minimization problem, but as a feasibility problem. Instead of navigating a loss landscape, the goal is to find network parameters and intermediate states that simultaneously satisfy a collection of local constraints derived from the network's structure and input-output mappings.

2. Methodology

The core methodology relies on iterative projection algorithms applied to a fine-grained computation graph.

A. Feasibility Formulation

Training is defined as finding a state vector $z$ (representing all edge values in the computation graph) that lies in the intersection of multiple constraint sets $C_v$ :
$\text{Find } z \in \bigcap_{v \in V} C_v$
Where $V$ represents nodes in the computation graph. The constraints are defined locally for each node type:

Constant Input Nodes: Enforce that outgoing edges carry the input value.
Parameter Nodes: Enforce consensus among all outgoing edges from a parameter (ensuring a single weight value is used everywhere).
Hidden Function Nodes: Enforce that the output of an edge equals the result of applying the primitive function (e.g., ReLU, dot product) to its inputs.
Target/Loss Nodes: Enforce that the network output minimizes the loss (or satisfies a margin constraint) relative to the target.

B. Projection Operators

The solution is found using iterative algorithms like Alternating Projections (AP), Cyclic Projections (CP), and Douglas-Rachford (DR).

Local Projections: For each node, the algorithm projects the current state onto the constraint set defined by that node.
Primitive Functions: The paper derives closed-form or efficient numerical solutions for projecting onto the graphs of elementary operations (e.g., dot products, ReLU of sums, max pooling, quantization).
Non-Differentiability: Since projections rely on distance minimization rather than derivatives, the framework naturally handles non-differentiable functions (e.g., quantization, step functions).

C. Parallelization via Bipartition

A key theoretical contribution is Theorem 2, which leverages the bipartite nature of computation graphs.

The graph nodes can be split into two disjoint sets, $A$ and $B$ , such that no two nodes within the same set are adjacent.
Projections for all nodes in set $A$ can be computed independently and in parallel because they operate on disjoint sets of edge variables. The same applies to set $B$ .
This allows the entire update step to be fully parallelized across the network, eliminating the sequential dependency of backpropagation.

D. The PJAX Framework

The authors introduce PJAX (Projection JAX), a software framework built on top of JAX.

Compositional: Similar to how autodiff composes derivatives, PJAX composes projection operators.
API: Provides a NumPy/JAX-like API for defining models.
Hardware: Inherits JAX's JIT compilation and support for GPUs/TPUs.
Automation: Automatically constructs the computation graph and orchestrates the iterative projection steps.

3. Key Contributions

Theoretical Reformulation: A rigorous formulation of neural network training as a large-scale feasibility problem using edge variables and local constraints, proving that a feasible solution corresponds to an empirical risk minimizer.
Parallelizable Algorithms: A proof that projection updates can be fully parallelized by exploiting the bipartite structure of computation graphs, offering a distinct advantage over sequential backpropagation.
Projection Operators: Mathematical derivations for projection operators on key primitive functions (dot product, ReLU, max, quantization) and loss functions (cross-entropy, margin loss).
PJAX Framework: The first general-purpose, gradient-free training framework that supports arbitrary architectures (MLPs, CNNs, RNNs) and hardware acceleration.
Empirical Validation: Comprehensive experiments demonstrating the viability of the approach across diverse tasks and architectures.

4. Experimental Results

The authors evaluated PJAX on standard benchmarks (MNIST, CIFAR-10, HIGGS, Shakespeare) using MLPs, CNNs, and RNNs, comparing against SGD, Adam, Feedback Alignment (FA), and Forward-Forward (FF).

Performance:
- MLPs: Projection methods (especially Douglas-Rachford) achieved competitive accuracy to SGD, often with significantly faster step times (up to 10x faster per step in shallow networks).
- RNNs: Projection methods showed a structural advantage by avoiding Backpropagation Through Time (BPTT), sidestepping vanishing/exploding gradient issues and converging faster in terms of steps than SGD.
- CNNs: Performance gap with gradient-based methods was larger, and step times were slower due to high memory overhead from shared parameters (see below).
Depth and Skip Connections: Deep networks without skip connections struggled with information propagation in projection-based training. However, introducing skip connections (shortening paths from early layers to the output) significantly improved performance, validating the need for architectural adjustments in this paradigm.
Efficiency:
- Time per Step: Projection methods (AP/CP) generally had lower time-per-step than SGD due to massive parallelism, though the total wall-clock time to convergence varied.
- Memory: The method requires storing state variables for every edge in the graph. This leads to $O(|E|)$ memory complexity, which is higher than the $O(|V|)$ of backpropagation, especially for architectures with heavy parameter sharing (CNNs, RNNs).

5. Significance and Future Directions

Biological Plausibility: The method aligns better with biological learning principles by using local updates and avoiding global error backpropagation.
Handling Non-Differentiability: It opens the door to training networks with discrete components (e.g., quantized neural networks, logical gates) without requiring gradient approximations.
Scalability: The parallel nature of the updates makes it highly suitable for modern hardware (GPUs/TPUs), potentially overcoming the sequential bottlenecks of backpropagation in very deep or wide networks.
Limitations: The primary limitation is memory usage, which scales with batch size, sequence length, and parameter sharing.
Future Work: The authors suggest exploring adaptive damping, hybrid schemes (mixing projections with gradients), memory-saving techniques (low-rank edge states), and tailored architectures to reduce parameter-data interaction.

In summary, this paper presents a compelling alternative to gradient-based training, demonstrating that feasibility-seeking via parallel projections is a viable, general, and efficient paradigm for deep learning, particularly for non-differentiable and recurrent architectures.