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. It supports spiking networks, quantized layers, hard routing, and other models without useful gradients. For differentiable models, backpropagation is usually faster.

Paper · Interactive visualization · API reference

Installation

pip install polystep
pip install "polystep[examples]"  # optional demo dependencies

Python 3.11+ and PyTorch 2.8+ are required. Choose a CUDA build compatible with your driver from the PyTorch installation guide. For development, install 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, PolyStepOptimizer, TrainConfig, train
from polystep.transform import ParamLayout

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

Replace the generated data with your dataset. Limit max_subspace_dim to control the forward budget. In manual loops, call register_evaluator() before each batch; train() does this automatically. 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.square().sum(dim=1))
print(es.best_fitness, es.best_solution)

Synthetic objectives

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 treats blocks independently; Sinkhorn also constrains target mass.

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.

Baselines

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

from polystep.baselines import Objective, openai_es

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

The paper runners support --fair for matched search spaces and evaluation budgets. EGGROLL uses a factored representation to retain matrix perturbations. See the reproduction protocol for tuning and budget details.

Benchmarks

Five-seed mean ± sample standard deviation. These architecture tables compare validation-tuned methods at matched optimizer steps. Evaluation- and time-matched studies are reported separately in the experiment index.

Non-differentiable tasks

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

MAX-SAT (% clauses satisfied)

Differentiable sanity checks

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
• Determinism
• Contributing and changelog

Citation

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

Acknowledgments

Thanks to Viet T. Nguyen for the interactive visualization.

License

Apache-2.0.