PJAX: Feasibility-Based Neural Training

• PJAX is a neural training framework that reformulates optimization as a feasibility problem using local projections rather than conventional gradient descent.
• It integrates a JAX-native architecture that traces computation graphs and applies projection algorithms (AP, CP, DR) to enforce constraints in parallel.
• Empirical evaluations demonstrate lower per-update latency and support for non-differentiable operations, though it sometimes lags behind Adam in test accuracy.

PJAX is a JAX-based software framework for neural network training that implements a feasibility-seeking paradigm based on projections, distinct from conventional gradient-based optimization. Rather than minimizing a global loss via backpropagation, PJAX formulates training as a feasibility problem defined by the intersection of local constraints derived from a network's computation graph. Training thus consists of iterative local projection operations, which can be massively parallelized. The framework natively supports GPU/TPU acceleration and is extensible to diverse network architectures and non-differentiable primitives (Bergmeister et al., 6 Jun 2025).

1. Feasibility-Seeking Formulation

Conventional neural network training solves the empirical risk minimization problem

$$\min_{\theta}\;\frac1N\sum_{i=1}^N\ell\bigl(f_\theta(x_i),\,y_i\bigr)$$

for parameter vector $\theta$ and loss function $\ell$. PJAX replaces this with a feasibility formulation, where the configuration is encoded by edge variables $z=(z_{uv})_{(u\to v)\in E}$ associated with the computation graph $G=(V,E)$ of the model: $$\text{Find } z\quad\text{s.t.}\quad z \in \bigcap_{v\in V} C_v$$ Each node $v \in V$ is associated with a constraint set $C_v$:

• Input nodes enforce $z_{c_j\to w} = x_j$ for all children $w$
• Parameter nodes require consensus among all outgoing edges $\theta$0
• Hidden function nodes enforce $\theta$1
• Target (loss) nodes impose goal-oriented constraints, typically matching to target labels via $\theta$2 or a proximal step

For each constraint set $\theta$3, the projection operator is defined as

$\theta$4

Key node types have specific projection forms, for example, parameter consensus uses edge-averaging, hidden function nodes decompose via Theorem 2.1 into averaging, solving a projection onto the function graph, and replication, while loss nodes use the proximal operator with respect to $\theta$5.

PJAX implements classical iterative projection algorithms:

• Alternating Projections (AP): sequentially project onto bipartitions of the constraints
• Cyclic Projections (CP): cycle through all constraints
• Douglas–Rachford (DR): applies reflection-based relaxed updates between partitions

2. Software Architecture

PJAX is architected as a projection-analog of autodiff libraries, with its core backend implemented atop JAX. The system traces user code involving pjax.Array and pjax.Parameter types to generate a computation graph at the scalar primitive level. All projection kernels are compiled via JAX's JIT/XLA mechanisms to run efficiently on CPU, GPU, or TPU.

Core primitives encompass both a forward computation (e.g., linear, activation, pooling) and a closed-form or iterative projection operator onto the associated computational graph. Shape-manipulating operations (such as reshape, transpose, concat, conv_patch) are implemented as extensible, invertible no-ops that retain data routing semantics without generating new constraints.

The high-level API includes:

• pjax.nn.Module: compositional layers (e.g., Linear, ReLU, Conv) similar to Flax
• pjax.optim: projection-based optimizers (AlternatingProjections, CyclicProjections, DouglasRachford) that:
  1. Accept a loss_fn(params) user function
  2. Build the constraint graph via code tracing
  3. Run $\theta$6 projection steps
  4. Return consensus parameters

• pjax.vmap: provides automatic vectorization/batching across data samples akin to jax.vmap, parallelizing projection code for minibatches

3. Parallelization and Algorithmic Strategy

PJAX leverages the bipartite structure of the computation graph for parallelization. If $\theta$7 is bipartite ($\theta$8 with edges only between $\theta$9 and $\ell$0), all projections $\ell$1 for $\ell$2 in partition $\ell$3 (or $\ell$4) can be computed simultaneously since they act on disjoint edge-variable coordinates, as formalized by Theorem 2.2.

Block-iterative variants exploit this parallelism:

• AP or DR alternate between applying all $\ell$5-node and $\ell$6-node projections in parallel
• CP can organize projections by layer, also enabling intra-layer parallelism

The DR algorithm uses reflection and relaxation steps, which empirically improve convergence rates. This design enables PJAX to scale efficiently on accelerated hardware by maximizing concurrency.

4. Illustrative Usage Patterns

PJAX is applied analogously to standard autodiff libraries, with user code specifying models and objectives in a declarative fashion. Consider a minimal multi-layer perceptron (MLP) trained on MNIST (784-dimensional input, 256 hidden units, 10 outputs):

$\ell$8

Crucially:

• model.init(key) initializes a tree of pjax.Parameters
• opt.update(loss_fn, params) traces the loss function, constructs the associated feasibility graph, executes 50 projection steps on the edge-variable state, and produces updated consensus parameters

The same pattern is extensible to CNNs and RNNs. A plausible implication is that any model assemblable from supported primitives (including non-differentiable ones) can be trained in this paradigm without fundamental algorithmic changes.

5. Empirical Evaluation

PJAX was empirically validated on MLPs, CNNs, and RNNs over MNIST, CIFAR-10, HIGGS, and Shakespeare next-character prediction. Standard optimizers (SGD, Adam with learning rate $\ell$7, batch size 256) serve as baselines; projection-based methods use 50 projection steps per batch.

Empirical findings include:

• Step time: Projection methods (AP/DR) exhibit substantially lower per-update latency (≈0.5 ms on a 1-hidden-layer MLP) compared to SGD/Adam (≈3 ms)
• Convergence (steps): DR typically converges in fewer steps than AP but more than Adam
• Test accuracy: Adam achieves the highest accuracy (98.1% on MNIST-MLP), DR attains 95–96%, and AP reaches 91–93%
• Network depth: Pure projection methods struggle in deep MLPs without skip/readout connections, but the use of skip links restores competitive performance
• RNNs: Projection-based methods obviate explicit backpropagation-through-time and converge in significantly fewer steps than SGD (albeit with more expensive steps)
• Non-differentiable operations: PJAX directly supports primitives such as quantization and max-pooling, unlike gradient-based methods, which usually require surrogates

6. Advantages and Limitations

PJAX introduces several intrinsic advantages:

• Gradient-free training: Only projections onto local constraints (primitive graphs) required, facilitating optimization of non-differentiable or discrete operations
• Local, parallel updates: Projections affect only edge-variables adjacent to the involved nodes; maximal parallelism is realizable on suitable hardware
• Biological plausibility: Updates propagate via local consensus and constraint matching, avoiding weight transport mechanisms inherent in backpropagation
• Framework extensibility: New primitive operations with defined projections can be seamlessly incorporated

However, there are notable limitations:

• Memory overhead: Storing one edge-variable per directed edge (per sample) increases memory demands relative to backpropagation, especially marked with parameter-shared architectures (e.g., CNNs, RNNs)
• Accuracy gap: Adaptive gradient optimizers (Adam) surpass projection methods in observed test accuracy across tasks
• Depth challenges: For very deep networks, vanilla projection algorithms require skip connections or specialized initialization to avoid collapse of network signal
• Open convergence theory: Lacking global guarantees for nonconvex feasibility problems, PJAX typically finds “best-fit” solutions; theoretical convergence in highly nonconvex settings remains an active area of investigation (Bergmeister et al., 6 Jun 2025)

PJAX thus constitutes a compositional, JAX-native approach for training neural networks via iterative projections, particularly compelling for applications involving non-differentiable modules or when exploiting extreme parallelism is essential.