---
title: 'PJAX: Feasibility-Based Neural Training'
url: https://www.emergentmind.com/topics/pjax
type: topic
---

# PJAX: Feasibility-Based Neural Training

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 [2506.05878].

## 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 $\{z_{p_k\to u}\}$
- **Hidden function nodes** enforce $\bigl(z_{h\to w}\bigr)_{w} = f_h\bigl(z_{u\to h}\bigr)_u$
- **Target (loss) nodes** impose goal-oriented constraints, typically matching to target labels via $\arg\min_{z'} \ell(z', y)$ or a proximal step

For each constraint set $C_v \subseteq \Re^{d_v}$, the projection operator is defined as
\[
P_{C_v}(z) = \arg\min_{y \in C_v} \|y-z\|_2^2
\]
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 $\ell$.

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 $K$ 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 $G$ is bipartite ($V = A \cup B$ with edges only between $A$ and $B$), all projections $P_{C_v}$ for $v$ in partition $A$ (or $B$) 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 $A$-node and $B$-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):

```python
import jax
import pjax
from pjax import nn, optim

class MLP(nn.Module):
    def __init__(self, in_features, hidden, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden)
        self.relu = nn.ReLU(hidden)
        self.fc2 = nn.Linear(hidden, num_classes)

    def __call__(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        return self.fc2(x)

model = MLP(784, 256, 10)
params = model.init(jax.random.PRNGKey(0))
opt = optim.DouglasRachford(steps_per_update=50)

for (x_batch, y_batch) in dataloader:
    def loss_fn(p):
        logits = model.apply(p, x_batch)
        return nn.cross_entropy(logits, y_batch)
    params, loss = opt.update(loss_fn, params)
```

Crucially:
- `model.init(key)` initializes a tree of `pjax.Parameter`s
- `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 $10^{-3}$, 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 [2506.05878]

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.

Source: https://www.emergentmind.com/topics/pjax