polystep

Gradient-free neural network training via optimal transport.

PolyStep optimizes neural networks without backpropagation. At each step, it samples polytope vertices around current parameters, evaluates losses via forward passes only, and computes softmax-weighted projections to find descent directions. That trains models with non-differentiable components (spiking networks, quantized layers, blackbox modules) where gradients are unavailable or undefined.

Based on the Sinkhorn Step algorithm (Le et al., NeurIPS 2023), extended with subspace compression, a softmax solver, and convergence analysis for piecewise-smooth losses.

How it works

1. Sample polytope vertices around the current parameters in a compressed subspace.
2. Evaluate the loss at each vertex (forward pass only, no gradients).
3. Compute softmax weights over the cost matrix.
4. Update the parameters by barycentric projection from the weighted vertices.

Want to play around with parameters? Viet T. Nguyen built a gorgeous interactive walkthrough that animates every step of the method -> explore the PolyStep visualization.

Installation

pip install polystep                  # from PyPI (core: torch only)
uv add polystep                       # or with uv

From source:

pip install -e .                      # core library (torch only)
pip install -e ".[examples]"          # + numpy, torchvision, matplotlib, Pillow, gymnasium, python-sat
pip install -e ".[dev]"               # + numpy, pytest (+ plugins), ruff
pip install -e ".[experiments]"       # + numpy, pandas, python-sat, torchvision, cma, snntorch, datasets, transformers, gymnasium
pip install -e ".[rl]"                # + gymnasium[box2d], stable-baselines3, swig (Box2D envs)

GPU: pip install torch --index-url https://download.pytorch.org/whl/cu130, or pick the CUDA build matching your driver from the PyTorch install page.

Quickstart

Synthetic optimization

import torch
from polystep import Ackley
from polystep.solver import PolyStep

solver = PolyStep.create(Ackley(dim=10), epsilon=0.5, max_iterations=50)
state = solver.run(torch.randn(100, 10))
print(f"Best cost: {min(state.costs):.4f}")

Neural network training

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from polystep import PolyStepOptimizer, train, TrainConfig
from polystep.epsilon import CosineEpsilon
from polystep.hybrid_subspace import HybridSubspace
from polystep.transform import ParamLayout

model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))

# Replace this with your real DataLoader.
train_loader = DataLoader(
    TensorDataset(torch.randn(1024, 784), torch.randint(0, 10, (1024,))),
    batch_size=64,
)

# HybridSubspace compresses the parameter space per layer. Cost per step scales with
# subspace_dim, and rank alone does not bound it: rank=4 on this 101k-parameter model
# gives subspace_dim=4338 and 600 ms/step on one CPU core at batch 64. Cap it
# directly instead; max_subspace_dim=256 gives 26 ms/step.
layout = ParamLayout.from_module(model)
subspace = HybridSubspace.from_layout(layout, rank=4, max_subspace_dim=256)

# Cosine schedules: broad exploration -> fine exploitation.
optimizer = PolyStepOptimizer(
    model, subspace=subspace, solver="softmax",
    epsilon=CosineEpsilon(init=10.0, target=0.5, decay=0.02),
    step_radius=CosineEpsilon(init=5.0, target=1.0, decay=0.008),
    probe_radius=CosineEpsilon(init=10.0, target=2.0, decay=0.016),
)

train(model, train_loader, nn.CrossEntropyLoss(), optimizer, TrainConfig(epochs=5))

Expect this to cost far more per step than backprop. Use a GPU build.

Two things dominate wall-clock:

• Pin the CPU thread count below nproc. PolyStep's per-step ops are small, so torch's default core-count intra-op pool oversubscribes and OpenMP spin-wait takes over: on a 24-core box examples/03 runs 2.1 s at torch.set_num_threads(1) against 427 s at the default. One thread suits most models; raise it only if your own objective issues wide ops. See docs/performance.md.
• Register the evaluator if you write your own loop. train() does it for you. A bare optimizer.step(closure) leaves closure() as the only route to the objective, so the fused in-place, factored and sparse-delta evaluators never run. Call optimizer.register_evaluator(evaluator, inputs, targets) before each step.

Drop-in gradient-free optimizer (ask/tell)

PolyStep also exposes the ask/tell interface used by evolution-strategy libraries, so it drops into ES benchmark harnesses (evosax, NeuroEvoBench) and any black-box loop:

import torch
from polystep import PolyStepES

es = PolyStepES(dim=20, epsilon=0.1, step_radius=0.3, x0=torch.full((20,), 2.0))
for _ in range(300):
    candidates = es.ask()           # (popsize, dim) points to evaluate
    es.tell(objective(candidates))  # lower fitness is better
print(es.best_fitness, es.mean)

experiments/bench_ask_tell.py compares it head-to-head with a Gaussian ES on the standard synthetic suite under a matched evaluation budget.

See examples/ for eleven runnable demos.

Spiking net (hard LIF thresholds) trained with forward passes only

CartPole policy search: no value function, no gradients

When to use PolyStep

PolyStep is designed for models where gradients are unavailable or unreliable:

• Spiking neural networks: hard LIF thresholds, discrete spike events
• Quantized layers: int8 weights, binary/ternary networks
• Blackbox modules: external simulators, API-based models, hardware-in-the-loop
• Hard routing: argmax gating, hard mixture-of-experts
• Combinatorial optimization: MAX-SAT, discrete assignment problems

If your model is fully differentiable, Adam/SGD will be faster and more accurate.

How PolyStep relates to other gradient-free methods

Zeroth-order optimization splits into two camps. Memory-efficient fine-tuning (MeZO and successors) skips the backprop tape but still assumes a useful local gradient. Evolution strategies and SPSA (CMA-ES, OpenAI-ES) estimate one from small perturbations.

PolyStep is a randomized direct search: it probes a finite-radius polytope around the current parameters and moves toward the lowest-cost vertices through an optimal-transport barycenter. Where the loss is piecewise-constant, the local gradient is zero almost everywhere, so finite differences carry no signal and ES/SPSA stall while a finite radius steps across the flat regions. Example 09 shows the separation on a hard decision tree.

Classical direct search (generalized pattern search, MADS) motivates that finite radius, but PolyStep does not inherit its guarantees: it evaluates no incumbent, moves to the barycenter every step, and drives its radius from the epsilon schedule.

Gradient-free baselines

polystep.baselines runs six gradient-free methods against the same protocol PolyStep uses, so a comparison is not confounded by the search space, the probe radius, the minibatch stream or the evaluation budget.

from polystep.baselines import Objective, openai_es, random_search

# fn: (N, dim) -> (N,) losses, lower is better. One call per generation, so a
# stochastic fn draws one minibatch per call and every candidate sees the same data.
obj = Objective(fn, dim=64, budget=10_000)
result = openai_es(obj, sigma=0.05, lr=0.1, popsize=32)
result.best_loss, result.evals   # evals is never above budget

To search a subspace instead, only the objective changes:

obj = Objective.from_subspace(hybrid, base_sd, loss_batch, budget=10_000)
result = random_search(obj, sigma=0.1)   # isolates the subspace from the update rule

One evaluation means one candidate scored. Objective counts rows, so a method that vmaps 32 candidates behind one call spends 32, exactly like one that makes 32 calls. It refuses a batch that would exceed the budget, so no method can overspend.

Note on EGGROLL: its rank-r perturbations A B^T are defined on weight matrices, and Objective.shapes says where those are. Objective.from_layout preserves the per-tensor shapes; Objective.from_subspace cannot, because HybridSubspace coordinates carry no matrix structure, so EGGROLL there degenerates to dense Gaussian ES on the coordinates. FactoredSubspace is the subspace that already implements the A B^T parameterization.

Fairness mode

The paper runners take --fair, which hands every gradient-free method in a table the same subspace (same class, rank and seed), the same candidate budget derived from what PolyStep spends over its configured epochs, the same minibatch stream and the same probe radius. Each result JSON records subspace_class, subspace_rank, eval_budget and evals_used, plus a per-generation trajectory against cumulative candidate evaluations for the accuracy-vs-evaluations figure. EGGROLL is the one recorded exception: it gets FactoredSubspace, for the reason above.

python experiments/runners/run_mnist.py --fair --methods polystep openai_es spsa mezo random_search eggroll cma_es
# tuning cost, in the same units, for both sides of the table
python experiments/runners/variant_sweep.py --envs mnist_mlp --stage screen   # writes tuning_cost.json
python experiments/runners/variant_sweep.py --envs mnist_mlp --baseline spsa  # appends to it

--theory-mode runs the unaccelerated reference configuration instead of the tuned one: jitter 0.05 on the probe and step radii with the smooth mollifier density, independently sampled rotations, flat epsilon, step radius r_0 (t+1)^-(1/2+0.1) (polystep.PowerDecay), orthoplex, HybridSubspace, and no momentum, amortization or Anderson acceleration. Jitter and the amortization heuristics are mutually exclusive and PolyStepOptimizer enforces that itself: probe_radius_jitter > 0 turns off adaptive_probes and amortize_steps, because jitter makes the per-step cost a noisy estimate that those heuristics read as progress.

Benchmarks

5-seed mean ± std, every method tuned on validation at an equal budget and matched on optimizer steps. Protocol and per-experiment notes: experiments/EXPERIMENT_INDEX.md. Tables are generated from the result JSONs, not hand-maintained.

Non-differentiable tasks

Test accuracy %. A dash means the cell is not in this release.

MAX-SAT (% clauses satisfied)

Differentiable sanity checks

SNN memory scaling (forward-only vs. BPTT)

PolyStep leads every gradient-free baseline on all five non-differentiable tasks, and beats the gradient surrogate only where the non-differentiability is hard. It does not beat Adam on differentiable problems, or domain solvers on MAX-SAT.

Features

• OT solvers: entropic Sinkhorn (full-space default), softmax (subspace default), KL-softmax interpolation, and greedy selection.
• Subspace compression: HybridSubspace (recommended), FactoredSubspace, AdaptiveSubspace, and sparse projection for very large models.
• Candidate evaluation without materializing weights: for nn.Sequential MLPs the step scores every candidate through a shared base forward plus a low-rank correction, 4.4x on a 203K MLP. Automatic. See docs/performance.md.
• Sub-linear memory: forward-only evaluation, no BPTT activation tape (~30x savings at long SNN horizons).
• Block-wise OT for per-layer decomposition.
• Vmap-safe layers: drop-in attention and LSTM that work under torch.vmap.
• torch.compile opt-in on the OT kernels (compile=True) and on the candidate forward (compile_evaluator; compile_forward auto-enables on the in-place path and takes an explicit False to opt out).
• Ask/tell API: PolyStepES drops into evosax / NeuroEvoBench-style ES harnesses and black-box loops.

Limitations

• Compute cost. Roughly tens of millions of forward passes (on the SNN benchmark, around 30M) vs. tens of thousands of Adam gradient steps for the same MNIST accuracy. This is inherent to zeroth-order methods.
• High-dimensional NLP. All-parameter fine-tuning of GPT-2 124M through a 128-dim projection collapses to random predictions; the projection ratio is far below the Johnson-Lindenstrauss floor.
• Adam baseline. The stronger surrogate-gradient / BPTT baseline for SNNs is not bundled with this release; see the arXiv preprint.

Full discussion in LIMITATIONS.md.

Documentation

Citation

Citation:

@article{le2026training,
  title={Training Non-Differentiable Networks via Optimal Transport},
  author={Le, An T},
  journal={arXiv preprint arXiv:2605.01928},
  year={2026}
}

Acknowledgments

A huge thank you to Viet for building a beautiful interactive PolyStep visualization, it brings the method to life and makes every step click!

License

Apache License 2.0. See LICENSE.