Metadata-Version: 2.4
Name: polystep
Version: 1.0.0
Summary: PyTorch PolyStep Optimizer - gradient-free neural network training via entropic optimal transport
Author: An T. Le
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/anindex/polystep
Project-URL: Repository, https://github.com/anindex/polystep
Project-URL: Documentation, https://github.com/anindex/polystep/blob/main/docs/api_overview.md
Project-URL: Changelog, https://github.com/anindex/polystep/blob/main/CHANGELOG.md
Project-URL: Issues, https://github.com/anindex/polystep/issues
Keywords: pytorch,optimal-transport,gradient-free,sinkhorn,polystep,optimization,neural-network
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.8
Provides-Extra: examples
Requires-Dist: numpy>=2.0; extra == "examples"
Requires-Dist: matplotlib>=3.8; extra == "examples"
Requires-Dist: Pillow>=10.0; extra == "examples"
Requires-Dist: gymnasium[classic-control]>=0.29; extra == "examples"
Provides-Extra: dev
Requires-Dist: numpy>=2.0; extra == "dev"
Requires-Dist: setuptools>=77; extra == "dev"
Requires-Dist: pytest>=8.0; extra == "dev"
Requires-Dist: pytest-timeout>=2.2; extra == "dev"
Requires-Dist: pytest-rerunfailures>=14.0; extra == "dev"
Requires-Dist: pytest-cov>=5.0; extra == "dev"
Requires-Dist: pytest-xdist>=3.5; extra == "dev"
Requires-Dist: ruff==0.15.13; extra == "dev"
Provides-Extra: experiments
Requires-Dist: numpy>=2.0; extra == "experiments"
Requires-Dist: pandas>=2.0; extra == "experiments"
Requires-Dist: python-sat; extra == "experiments"
Requires-Dist: torchvision>=0.25; extra == "experiments"
Requires-Dist: cma>=3.3; extra == "experiments"
Requires-Dist: datasets>=2.14; extra == "experiments"
Requires-Dist: transformers>=4.40; extra == "experiments"
Requires-Dist: gymnasium>=0.29; extra == "experiments"
Provides-Extra: rl
Requires-Dist: gymnasium>=0.29; extra == "rl"
Requires-Dist: stable-baselines3>=2.3; extra == "rl"
Dynamic: license-file

# polystep

[![PyPI](https://img.shields.io/pypi/v/polystep.svg)](https://pypi.org/project/polystep/)
[![arXiv](https://img.shields.io/badge/arXiv-2605.01928-b31b1b.svg)](https://arxiv.org/abs/2605.01928)
[![Python 3.11+](https://img.shields.io/badge/python-3.11%2B-blue.svg)](https://www.python.org/downloads/)
[![PyTorch 2.8+](https://img.shields.io/badge/PyTorch-2.8%2B-ee4c2c.svg)](https://pytorch.org/)
[![License](https://img.shields.io/badge/license-Apache%202.0-green.svg)](https://github.com/anindex/polystep/blob/main/LICENSE)

Gradient-free neural network training via optimal transport.

Accepted at [Transactions on Machine Learning Research (TMLR)](https://openreview.net/forum?id=8mlcqTTMuU).

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](https://arxiv.org/abs/2605.01928) ·
[Interactive visualization](https://vietngth.github.io/polystep-visualization/) ·
[API reference](https://github.com/anindex/polystep/blob/main/docs/api_overview.md) ·
[Julia package](https://github.com/anindex/PolyStep.jl)

<p align="center">
  <img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/method_diagram.png" width="840"
       alt="PolyStep: search coordinates, polytope probes, cost evaluation, soft assignment, and a barycentric update.">
</p>

## Installation

```bash
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](https://pytorch.org/get-started/locally/).
For development, install `pip install -e ".[dev]"` from a checkout.

## Quickstart

```python
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](https://github.com/anindex/polystep/blob/main/docs/performance.md)
for candidate paths, thread counts, compilation, and memory settings.

### Ask/tell

```python
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

```python
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:

```python
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](https://github.com/anindex/polystep/blob/main/docs/reproducibility.md#protocol) 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](https://github.com/anindex/polystep/blob/main/docs/reproducibility.md#changes-since-the-published-tables)
lists the rows that 1.0 changes. The
[experiment index](https://github.com/anindex/polystep/blob/main/experiments/EXPERIMENT_INDEX.md)
has the full tables.

<!-- BENCH:START -->
### 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 |
<!-- BENCH:END -->

The SNN surrogate-gradient baseline varies substantially across seeds. Adam leads
on the differentiable tasks, and specialized solvers lead on MAX-SAT. See
[limitations](https://github.com/anindex/polystep/blob/main/LIMITATIONS.md)
for unsupported configurations and comparison limits.

## Documentation

- [API reference](https://github.com/anindex/polystep/blob/main/docs/api_overview.md)
- [Performance](https://github.com/anindex/polystep/blob/main/docs/performance.md)
- [Examples](https://github.com/anindex/polystep/blob/main/examples/README.md)
- [Experiments](https://github.com/anindex/polystep/blob/main/experiments/README.md) and [reproduction](https://github.com/anindex/polystep/blob/main/docs/reproducibility.md)
- [Determinism](https://github.com/anindex/polystep/blob/main/docs/determinism.md)
- [Contributing](https://github.com/anindex/polystep/blob/main/CONTRIBUTING.md) and [changelog](https://github.com/anindex/polystep/blob/main/CHANGELOG.md)

## Citation

```bibtex
@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](https://vietngth.github.io/) for the
[interactive visualization](https://vietngth.github.io/polystep-visualization/).

## License

[Apache-2.0](https://github.com/anindex/polystep/blob/main/LICENSE).
