polystep

Gradient-free neural network training via optimal transport.

Accepted at Transactions on Machine Learning Research (TMLR).

PolyStep evaluates finite-radius parameter perturbations and updates parameters through a weighted barycenter. I use it for spiking networks, quantized layers, hard routing, and other models without useful gradients. For differentiable models, backpropagation is usually faster.

Version 1.1.1. See the changelog for release changes.

Paper · Interactive visualization · API reference · Julia package

One particle	Four particles, two minima

Animations from Viet T. Nguyen's interactive visualization.

Installation

pip install polystep
pip install "polystep[examples]"  # dependencies for the repository's examples

Python 3.11+ and PyTorch 2.8+ are required. Choose a CUDA build compatible with your driver from the PyTorch installation guide. For development, run pip install -e ".[dev]" from a checkout.

Quickstart

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from polystep import HybridSubspace, ParamLayout, PolyStepOptimizer, TrainConfig, train

torch.set_num_threads(1)
model = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 3))
loader = DataLoader(
    TensorDataset(torch.randn(256, 16), torch.randint(0, 3, (256,))),
    batch_size=32,
)
subspace = HybridSubspace.from_layout(
    ParamLayout.from_module(model), rank=4, max_subspace_dim=128,
)
optimizer = PolyStepOptimizer(model, subspace=subspace, solver="softmax")
train(model, loader, nn.CrossEntropyLoss(), optimizer, TrainConfig(epochs=5))

max_subspace_dim=128 caps the search at 128 coordinates, so each step scores 144 candidates. Biases and normalization weights keep their full width. On a larger model, the cap must cover those parameters plus at least one coordinate per weight matrix. If the cap is too small, PolyStep warns and reports the minimum dimension it can use.

step_batch uses the same candidate evaluation paths as train():

loss_fn = nn.CrossEntropyLoss()
for x, y in loader:
    optimizer.step_batch(x, y, loss_fn=loss_fn)
optimizer.release_evaluator()

torch.set_num_threads(1) avoids thread overhead for a model this small. Larger models may benefit from more threads; python -m polystep.benchmarks.threads compares several counts on an MLP workload. See performance for candidate paths, thread counts, compilation, and memory settings.

Ask/tell

from polystep import PolyStepES

es = PolyStepES(dim=20, epsilon=0.1, step_radius=0.3)
for _ in range(300):
    candidates = es.ask()
    es.tell((candidates - 1).square().sum(dim=1))
print(es.best_fitness, (es.mean - 1).square().sum().item())

Synthetic objectives

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))

Method

Each step samples a rotated polytope, evaluates its vertices, computes transport weights, and moves to their barycenter. Subspaces reduce the number of coordinates searched. Softmax normalizes each particle's vertex weights; Sinkhorn also enforces the target marginal.

Finite-radius probes can cross flat regions of discontinuous losses. Performance depends on the radius, representation, and evaluation budget. PolyStep's optional quadratic models and radius controllers are implementation heuristics; they do not inherit classical direct-search or trust-region guarantees.

SNN training	CartPole policy

The SNN example trains a network with hard LIF spikes on four synthetic classes, without surrogate gradients.

The CartPole example uses real-valued weights and a hard left/right argmax. PolyStep optimizes episode return without gradients through the actions or simulator.

Baselines

polystep.baselines provides openai_es, spsa, mezo, random_search, eggroll, and cma_es, which needs pip install cma. One evaluation means one candidate scored, even when candidates are batched:

from polystep.baselines import Objective, openai_es

objective = Objective(lambda x: (x - 1).square().sum(1), dim=64, budget=10_000)
result = openai_es(objective, sigma=0.05, popsize=32)

The paper runners match search spaces, budgets and tuning as the reproducibility guide describes.

Benchmarks

Five-seed mean ± sample standard deviation, with configurations and checkpoints selected on validation data; the probSAT and RC2 columns are single runs. Each baseline gets a matched evaluation or step budget, as the protocol describes. The tables were recorded before 1.0; the reproducibility guide lists the rows that later releases change, and the experiment index has the full tables.

Non-differentiable tasks

Test accuracy %. A dash means the method was not run on that task.

Task	PolyStep	CMA-ES	OpenAI-ES	SPSA	Adam (surrogate)	Non-diff op
SNN/LIF (MNIST)	93.0 ± 0.2	77.1 ± 13.3	79.6 ± 5.2	53.9 ± 3.0	86.3 ± 10.6	threshold()
Int8 quantized	97.0 ± 0.1	91.6 ± 0.4	94.4 ± 0.2	82.5 ± 0.5	97.8 ± 0.1	round()
Argmax attention	86.6 ± 0.4	79.5 ± 0.5	80.0 ± 0.4	67.4 ± 1.3	88.6 ± 0.1	argmax()
Staircase activation	94.3 ± 0.1	89.2 ± 0.4	88.9 ± 0.5	45.2 ± 4.2	97.5 ± 0.1	floor()
Hard MoE routing	90.6 ± 0.2	77.8 ± 1.5	82.5 ± 0.8	25.4 ± 3.5	-	argmax()

MAX-SAT (% clauses satisfied)

Variables	PolyStep	CMA-ES	OpenAI-ES	probSAT	RC2
100	98.0 ± 0.7	99.0 ± 0.3	99.4 ± 0.1	99.8	99.8
5,000	98.1 ± 0.1	95.2 ± 0.2	92.9 ± 0.2	99.9	timeout
100,000	98.1 ± 0.0	90.7 ± 0.1	88.9 ± 0.0	99.6	timeout

Differentiable sanity checks

Task	PolyStep	Adam
MNIST (2-layer MLP)	96.8 ± 0.1	97.7 ± 0.1
ETTh1 (LSTM, MSE; lower is better)	0.253 ± 0.023	0.247 ± 0.023

The SNN surrogate-gradient baseline varies substantially across seeds. Adam leads on the differentiable tasks, and specialized solvers lead on MAX-SAT. See limitations for unsupported configurations and comparison limits.

Documentation

• API reference
• Performance
• Examples
• Experiments and reproduction
• Limitations
• Determinism
• Contributing and changelog

Citation

@article{le2026training,
  title={Training Non-Differentiable Networks via Optimal Transport},
  author={Le, An T},
  journal={Transactions on Machine Learning Research},
  year={2026},
  url={https://openreview.net/forum?id=8mlcqTTMuU}
}

Acknowledgments

Thanks to Viet T. Nguyen for the interactive visualization.

License

Apache-2.0.