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 · Julia package

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, 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, ParamLayout, PolyStepOptimizer, TrainConfig, train

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. max_subspace_dim=128 caps the search at 128 coordinates, so each step scores 144 candidates. 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

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.

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.square().sum(1), dim=64, budget=10_000)
result = openai_es(objective, sigma=0.05, lr=0.1, 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. In the architecture tables, OpenAI-ES, CMA-ES and EGGROLL get PolyStep's candidate-evaluation budget, and SPSA, MeZO and random search get its step count. On argmax attention, hard MoE and ETTh1, the population methods instead get 32 evaluations per PolyStep step, spent at their own population size, and Adam trains for PolyStep's 30 epochs, except 50 on ETTh1. On MAX-SAT, CMA-ES and OpenAI-ES get 50 evaluations per PolyStep step. OpenAI-ES scores 50 candidates at each of PolyStep's steps, while CMA-ES doubles its population at each restart and takes fewer generations. The tables were recorded before 1.0; the reproducibility guide lists the rows that 1.0 changes. The experiment index has the full tables.

Non-differentiable tasks

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

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={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.