polystep 0.11.0__tar.gz → 0.13.0__tar.gz

polystep-0.13.0/CHANGELOG.md

@@ -0,0 +1,103 @@
+# Changelog
+
+## 0.13.0 - 2026-09-22
+
+- Fix KL row mass, nonfinite inputs, evaluator hooks, and chunk handling.
+- Reuse shared MLP forward work and load only requested MNIST images.
+- Fix example checkpoint selection and evaluation budgets; trim demo dependencies.
+
+## 0.12.0 - 2026-09-17
+
+- Fix cross-entropy targets, stale caches, attention masks, nonfinite best values,
+  and mixed-dtype evaluation.
+- Correct simplex gradients and whole-model trust predictions; migrate checkpoints
+  to format 5.
+- Reuse MLP prefixes, use native CPU attention, and compile site-local evaluation.
+- Average tied baseline ranks; add controlled experiments and forward benchmarks.
+- Validate sparse dimensions and factored ranks; simplify documentation.
+
+## 0.11.0 - 2026-08-14
+
+- Remove `LowRankSubspace`, `FactorSpec`, and `LinearSubspace`; use `HybridSubspace`
+  or `FactoredSubspace`. Replace `absorb_every` with `absorb_mode`.
+- Remove `polystep.benchmarks.baselines`, `evotorch`, `SinkhornSolver.solve(init_eps)`,
+  and solver warm-start methods. `CMAAdaptiveSubspace.base` is removed.
+- Require `dim` in `probe_scale_of`; correct baseline probe scaling.
+- Fix nonfinite costs, zero-mass rows, quadratic evaluation counts, and resume state.
+- Add tight-frame quadratic fits and cache radius-jitter sampling.
+
+## 0.10.1 - 2026-08-01
+
+- Require `objective_token` for probe reuse and reject invalid optimizer settings.
+- Fix blockwise evaluation, checkpoint aliasing, dtype changes, and tied weights.
+- Add site-local evaluation to blockwise mode and reduce subspace allocations.
+
+## 0.10.0 - 2026-07-30
+
+- Default to simplex sampling; retune radii when upgrading.
+- Change seeded rotation streams and clear momentum on basis changes.
+- Enable forward compilation on the in-place path by default.
+- Remove `projection_mode`, `use_csa`, low-level block strategies, and unused exports.
+- Move NumPy to optional extras; validate schedules and constructor arguments.
+
+## 0.9.0 - 2026-07-27
+
+- Add sparse-delta evaluation, `FactoredSubspace`, greedy solvers, and evaluator
+  registration helpers.
+- Fix probe reuse, screening, subspace state, tied weights, and best-weight restoration.
+- Extend checkpoint state and reduce projection and evaluation allocations.
+
+## 0.8.0 - 2026-07-24
+
+- Add optimizer checkpointing and forward-compilation controls.
+- Fix mixed precision, scheduled epsilon, and invalid solver inputs.
+
+## 0.7.0 - 2026-07-20
+
+- Publish on PyPI and reuse blockwise candidate buffers.
+- Fix constant-offset cost stability and biased rotations.
+
+## 0.6.1 - 2026-07-15
+
+- Add the hard decision-tree example.
+- Fix mixed precision, barycentric normalization, and candidate memory estimates.
+- Cache hybrid projections between rotations.
+
+## 0.6.0 - 2026-07-12
+
+- Add direct F1 optimization and optimizer-variant experiments.
+- Fix CMA scaling, solver validation, and trust-region state.
+
+## 0.5.0 - 2026-07-07
+
+- Add `PolyStepES`, `minimize`, and the binary-network example.
+- Fix Anderson acceleration, Newton selection, screening, and step adaptation.
+
+## 0.4.0 - 2026-06-21
+
+- Share solver cost handling and remove low-rank Sinkhorn.
+- Fix cost scaling, dual recentering, and activation semantics.
+
+## 0.3.0 - 2026-05-27
+
+- Reduce GPU synchronization in solvers and blockwise updates.
+- Fix CMA rotation forwarding and API examples.
+
+## 0.2.2 - 2026-05-21
+
+- Use unit-radius vertex generators with keyword-only `radius`.
+- Fix CMA transport weighting and stale API examples.
+
+## 0.2.1 - 2026-05-14
+
+- Add the Loihi-style SNN adaptation example.
+
+## 0.2.0 - 2026-05-14
+
+- Require Python 3.11, PyTorch 2.8, and NumPy 2.0 for optional extras.
+- Reduce solver synchronization and use inference mode for evaluation.
+
+## 0.1.0 - 2026-05-03
+
+- Initial release: gradient-free training, OT solvers, subspaces, vmap-compatible
+  layers, examples, and experiment runners.

{polystep-0.11.0 → polystep-0.13.0}/CITATION.cff

@@ -6,8 +6,8 @@ authors:
   - family-names: Le
     given-names: An T.
 license: Apache-2.0
-version: 0.11.0
-date-released: "2026-08-14"
+version: 0.13.0
+date-released: "2026-09-22"
 repository-code: "https://github.com/anindex/polystep"
 keywords:
   - gradient-free optimization

polystep-0.13.0/CONTRIBUTING.md

@@ -0,0 +1,54 @@
+# Contributing
+
+Install the development dependencies and run the same checks as CI:
+
+```bash
+pip install -e ".[dev]"
+ruff check .
+ruff format --check .
+pytest tests/ -q
+```
+
+Create a branch, make the change, add regression coverage where needed, and open a
+pull request. Keep comments specific and public APIs typed.
+
+## Tests
+
+The default suite excludes `slow` and `gpu` tests. To select them:
+
+```bash
+pytest tests/ -m slow
+pytest tests/ -m gpu
+pytest tests/ -m ""
+```
+
+Tests use one CPU thread by default; `POLYSTEP_TEST_THREADS` overrides it. Slow tests
+may download MNIST. Tests importing `experiments/` must skip when that directory is
+absent, since it is excluded from distributions.
+
+Pushes and pull requests test Python 3.11 with PyTorch 2.8 and Python 3.14 with
+current PyTorch. Scheduled and manual runs add Python 3.12/3.13, coverage, real-data
+learning checks, and the shipped tests against the installed wheel. Every run
+builds both distributions and checks the wheel's imports and type marker.
+
+Before a release:
+
+```bash
+uv build
+uvx twine check --strict dist/*
+```
+
+## Releases
+
+Update `polystep.__version__`, the version and date in `CITATION.cff`, and
+`CHANGELOG.md`. `pyproject.toml` reads the package version automatically.
+
+Publishing a GitHub Release tagged `vX.Y.Z` triggers the PyPI workflow. It checks
+version consistency, lint, tests, and package metadata before publishing. The
+release also runs the shipped tests against the wheel. Only the separate publish
+job has OIDC permission; it uses the `pypi` environment and PyPI Trusted Publishing.
+
+## Issues
+
+Include a minimal reproduction, the traceback, Python and PyTorch versions, and
+hardware details relevant to the failure.

polystep-0.13.0/LIMITATIONS.md

@@ -0,0 +1,92 @@
+# Limitations
+
+## Cost and search space
+
+PolyStep needs many candidate forwards per update. Differentiable models usually
+train faster with backpropagation. Small subspaces reduce cost but may exclude useful
+directions; very small projection ratios can fail on large models.
+
+`HybridSubspace` caps each layer at its parameter count. Large dense projections may
+use sparse signed columns, which are not orthonormal. Claims requiring an orthogonal
+projection do not apply to those columns.
+
+`FactoredSubspace` restricts perturbations to the input directions in its fixed
+factor. Rotation changes those directions. Retune radii when switching subspaces.
+
+Momentum is cleared when the basis changes. `AdaptiveSubspace` rotates every step
+by default, so momentum requires a longer rotation interval to accumulate.
+Heavy-ball momentum can amplify the barycentric displacement; `step_radius` bounds
+the displacement before momentum.
+
+CMA adaptation keeps its basis fixed between absorbs and normalizes covariance to
+`trace(C) = n`. Its default `mu_eff=1` gives rank-one adaptation only. Cumulative
+step-size adaptation is not implemented.
+
+## Evaluation
+
+- `step()` expects `closure(batched_params) -> losses` with shape `(N,)`.
+- Manual loops must register the evaluator and current batch to enable fast paths.
+  The evaluator and closure must compute the same objective.
+- Mixed-dtype models require `HybridSubspace` or `FactoredSubspace`. Full-space and
+  global-subspace coordinates use one dtype and reject mixed parameter dtypes.
+- `mixed_precision=True` uses BF16 geometry and forwards; solver costs use FP32.
+  Models needing autocast internally must provide it.
+- `adaptive_probes` needs an unchanged whole configuration and `objective_token`.
+  It cannot reuse losses across changing minibatches.
+- Deferred `trust_region` comparisons require the same objective. `train()` rejects
+  this option; use manual steps on a stationary objective.
+- `trust_region` and `use_adaptive_radius` both scale the step. Choose one controller.
+- Multi-fidelity screening imputes dropped costs using an additive offset. This can
+  bias the full-fidelity cost estimate and is disabled with quadratic-model options.
+- Blockwise mode does not support every monolithic adaptation option; unsupported
+  options warn at construction.
+- The library does not distribute optimizer state across multiple GPUs.
+
+## Attention and recurrent layers
+
+[`VmapSafeMultiHeadAttention`](src/polystep/layers/attention.py) requires batch-first
+inputs and `kdim == vdim == embed_dim`. It does not support `add_bias_kv`,
+`add_zero_attn`, returned attention weights, or `is_causal=True`; supply an explicit
+causal mask. It returns an output tensor, not an `(output, weights)` tuple.
+
+Boolean masks use `True` for masked positions; float masks are additive. Attention
+masks accept `(query, key)`, `(batch * heads, query, key)`, and broadcastable 4D shapes.
+The wrapper also accepts `(batch, query, key)` masks.
+
+CPU uses native scaled-dot-product attention; CUDA uses explicit matmuls. CPU vmap
+may emit a batching-fallback warning. Evaluate dropout layers in `model.eval()` mode
+for deterministic candidate comparisons.
+
+[`VmapSafeLSTM`](src/polystep/layers/rnn.py) requires batch-first dense inputs. It does
+not support bidirectional recurrence, projection, or `PackedSequence`. Its explicit
+cell computations can be slower than fused cuDNN recurrence.
+
+## Solvers and mathematical scope
+
+`SoftmaxSolver` enforces source row sums only; it ignores the target marginal and
+warns for a nonuniform one. Use `SinkhornSolver` to enforce both marginals.
+Independent particle rotations mean Sinkhorn's shared column index does not identify
+a common direction across particles.
+
+Sinkhorn requires `omega` in `[0.5, 1.95]`. Anderson acceleration and adaptive omega
+are inactive in fixed-iteration mode (`threshold <= 0`). Both solver families promote
+FP16/BF16 costs to FP32.
+
+All nonfinite costs are invalid. Their finite replacement can affect global cost
+scaling, including `scale_cost='max_cost'`.
+
+Quadratic fits are finite-radius models. Simplex gradients can be biased on
+anisotropic quadratics, and block models omit cross-block curvature. A fit does not
+supply a classical derivative at a discontinuity. The deferred radius controller
+does not perform a trust-region acceptance test.
+
+## Benchmark comparisons
+
+The architecture tables use validation-selected checkpoints. `--allow-test-leakage`
+is an explicit opt-in for test-selected runs, which standard aggregation rejects.
+
+The SNN surrogate-gradient baseline has substantial variation across seeds. MAX-SAT
+search heuristics use flip budgets, which are not equivalent to neural candidate
+budgets. Head-only GPT-2 experiments do not establish full-model fine-tuning results.
+
+See the [experiment index](experiments/EXPERIMENT_INDEX.md) for protocols and results.

polystep-0.13.0/PKG-INFO

@@ -0,0 +1,237 @@
+Metadata-Version: 2.4
+Name: polystep
+Version: 0.13.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 :: 4 - Beta
+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: snntorch>=0.7; 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[box2d]>=0.29; extra == "rl"
+Requires-Dist: stable-baselines3>=2.3; extra == "rl"
+Requires-Dist: swig>=4.2; 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](docs/api_overview.md)
+
+<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]"  # 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](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, 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](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
+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:
+
+```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 support `--fair` for matched search spaces and evaluation budgets.
+EGGROLL uses a factored representation to retain matrix perturbations. See the
+[reproduction protocol](docs/reproducibility.md) 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](experiments/EXPERIMENT_INDEX.md).
+
+<!-- BENCH:START -->
+### Non-differentiable tasks
+
+Test accuracy %. A dash means the result is not in this release.
+
+| 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](LIMITATIONS.md) for unsupported configurations and comparison limits.
+
+## Documentation
+
+- [API reference](docs/api_overview.md)
+- [Performance](docs/performance.md)
+- [Examples](examples/README.md)
+- [Experiments](experiments/README.md) and [reproduction](docs/reproducibility.md)
+- [Determinism](docs/determinism.md)
+- [Contributing](CONTRIBUTING.md) and [changelog](CHANGELOG.md)
+
+## Citation
+
+```bibtex
+@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](https://vietngth.github.io/) for the
+[interactive visualization](https://vietngth.github.io/polystep-visualization/).
+
+## License
+
+[Apache-2.0](LICENSE).

polystep-0.13.0/README.md

@@ -0,0 +1,180 @@
+# 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](docs/api_overview.md)
+
+<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]"  # 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](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, 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](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
+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:
+
+```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 support `--fair` for matched search spaces and evaluation budgets.
+EGGROLL uses a factored representation to retain matrix perturbations. See the
+[reproduction protocol](docs/reproducibility.md) 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](experiments/EXPERIMENT_INDEX.md).
+
+<!-- BENCH:START -->
+### Non-differentiable tasks
+
+Test accuracy %. A dash means the result is not in this release.
+
+| 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](LIMITATIONS.md) for unsupported configurations and comparison limits.
+
+## Documentation
+
+- [API reference](docs/api_overview.md)
+- [Performance](docs/performance.md)
+- [Examples](examples/README.md)
+- [Experiments](experiments/README.md) and [reproduction](docs/reproducibility.md)
+- [Determinism](docs/determinism.md)
+- [Contributing](CONTRIBUTING.md) and [changelog](CHANGELOG.md)
+
+## Citation
+
+```bibtex
+@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](https://vietngth.github.io/) for the
+[interactive visualization](https://vietngth.github.io/polystep-visualization/).
+
+## License
+
+[Apache-2.0](LICENSE).