SymJAX: Projection-Based Neural Training

• SymJAX is a projection-based framework built on JAX that reformulates neural network training as a large-scale feasibility problem using local constraint projection operators.
• It employs iterative algorithms like Alternating Projections, Cyclic Projections, and Douglas–Rachford Splitting to enforce network node constraints and optimize performance.
• The framework leverages high parallelism and vectorization, enabling efficient training of MLPs, CNNs, and RNNs on diverse hardware such as CPUs, GPUs, and TPUs.

PJAX is a JAX-based software framework for neural network training that implements a projection-based, gradient-free optimization paradigm. Distinct from conventional gradient-based minimization, PJAX reformulates optimization as a large-scale feasibility problem, using local projection operators and iterative projection algorithms. The framework is designed for high parallelism, compositionality, and the explicit accommodation of non-differentiable or discrete operations (Bergmeister et al., 6 Jun 2025).

1. Feasibility-Seeking Formulation

PJAX replaces traditional loss minimization

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

with the problem of finding edge-variables

$z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$

subject to the intersection of local constraint sets: $z\;\in\;\bigcap_{v\in V}\;C_v$ Here, the computational graph $G=(V,E)$ is constructed such that each node $v$ (input, parameter, hidden function, or target/loss node) imposes a local constraint $C_v$.

Local constraint classes are as follows:

• Input nodes $c_j$: Enforce $z_{c_j\to w}=x_j$.
• Parameter nodes $p_k$: Require consensus among all outgoing edge variables $\{z_{p_k\to u}\}$.
• Hidden function nodes $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$0: Enforce

$z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$1

• Target nodes $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$2: Impose loss minimality, either via $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$3 or proximal update $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$4.

Node constraints are enforced by orthogonal projection: $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$5

PJAX supports the following iterative projection algorithms:

• Alternating Projections (AP)
• Cyclic Projections (CP)
• Douglas–Rachford (DR) Splitting, leveraging reflections and relaxation for empirically improved convergence.

2. Software Architecture

PJAX operates as a "projection-analog" of autodiff libraries, leveraging JAX as its backend. The system traces user code operating on pjax.Array and pjax.Parameter objects to generate a computation graph comprising primitive scalar operations. JAX’s JIT and XLA compilation enable projection kernels to execute on CPU, GPU, or TPU.

• Core primitives include both forward and projection (closed-form or iterative) implementations for operations such as identity, sum, dot, nonlinearities, max, quantization, etc. Shape transformations are handled as invertible “no-op” nodes.
• High-level API: The framework provides pjax.nn.Module (analogous to Flax) for network definition, and pjax.optim for optimizers implementing projection schemes. Typical workflow:
  1. Define a Python loss_fn(params) using PJAX operations.
  2. Trace and build the constraint graph.
  3. Run $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$6 projection steps per update.
  4. Extract updated consensus parameters.

• Vectorization: pjax.vmap wraps user functions to apply projection code independently to each sample in a batch.

3. Parallelization and Block-Iterative Schemes

PJAX exploits the bipartite structure of the computation graph $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$7, partitioning nodes into disjoint sets $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$8 and $z=(z_{uv})_{(u\to v)\in E}\in\Re^{|E|}$9 such that all projections $z\;\in\;\bigcap_{v\in V}\;C_v$0 for $z\;\in\;\bigcap_{v\in V}\;C_v$1 (or $z\;\in\;\bigcap_{v\in V}\;C_v$2) affect disjoint coordinates and are thus fully parallelizable. The main strategies are:

• Bipartite Alternation: Projection algorithms (AP or DR) alternate between synchronous updates across all nodes in $z\;\in\;\bigcap_{v\in V}\;C_v$3 and all nodes in $z\;\in\;\bigcap_{v\in V}\;C_v$4.
• Layerwise Parallelism: CP can organize projections by network layers, paralleling projections within each layer.
• Douglas–Rachford Synchronization: Combines reflection and relaxation for improved convergence over naive alternating projections.

This formulation enables maximal parallelism on multi-core and accelerator hardware.

4. Practical Usage and Example Workflows

PJAX integrates into research code called in a manner similar to other autodiff-oriented machine learning frameworks. For example, the minimal training loop for an MLP on MNIST involves:

$z\;\in\;\bigcap_{v\in V}\;C_v$9 Key elements include parameter initialization, tracing loss functions to build constraint graphs, and stepping the optimizer to repeatedly project onto the network’s constraint sets. Similar idioms apply to CNNs and RNNs, including those with skip or “readout” connections.

5. Empirical Evaluation

PJAX’s empirical performance has been evaluated across diverse models—MLPs, CNNs, and RNNs (with and without skip connections)—on MNIST, CIFAR-10, HIGGS, and Shakespeare next-character benchmarks.

Key findings include:

• Step time: DR/AP projection updates on a 1-hidden-layer MLP are approximately $z\;\in\;\bigcap_{v\in V}\;C_v$5 ms per step, compared to $z\;\in\;\bigcap_{v\in V}\;C_v$6 ms for SGD/Adam, providing %%%%22$\{z_{p_k\to u}\}$23%%%% faster per-update speed.
• Convergence: DR typically requires fewer steps than AP, though more than Adam.
• Test accuracy: Adam achieves the best final accuracy (98.1% on MNIST-MLP); DR approaches 95–96%, AP around 91–93%.
• Deep MLPs: Without skip connections, projection-based schemes underperform; adding skip connections restores accuracy performance.
• RNNs: Projection approaches sidestep backpropagation-through-time, converging in far fewer steps than SGD (though steps are individually more expensive).
• Non-differentiable primitives: PJAX directly optimizes models with operations such as quantization or max-pooling, unlike gradient-based approaches that require surrogate gradients.

6. Advantages and Limitations

PJAX’s projection-based approach offers a unique set of benefits and trade-offs:

Strengths	Limitations
Gradient-free operation	Higher memory overhead (per-edge variable storage)
Handles non-differentiable	Final accuracy gap vs. Adam/SGD
Discrete operations	Deep networks require skip-connections
Highly parallel local updates	Lack of global convergence guarantees
Biologically plausible: no global weight-transport
Extensible to new primitives

PJAX does not require gradients and is thus inherently capable of training over non-differentiable or discrete primitives. Its local update rules—focused on modifying edge variables immediately adjacent to a node—allow massive, structured parallelism, especially under bipartite partitioning. The compositionality enables rapid prototyping and extension. However, memory requirements are elevated relative to gradient-based methods; this is particularly acute in architectures with significant parameter sharing (e.g., CNNs, RNNs). While training is robust, adaptive optimizers such as Adam currently offer higher final accuracy. Vanilla projection approaches struggle with deep architectures unless enhanced by skip connections or tailored initialization. Theoretical guarantees for convergence in nonconvex settings remain an open question, with empirical results demonstrating “best-fit” rather than globally optimal solutions (Bergmeister et al., 6 Jun 2025).