polystep 0.12.0__tar.gz → 1.0.0__tar.gz

{polystep-0.12.0 → polystep-1.0.0}/CHANGELOG.md

@@ -1,5 +1,43 @@
 # Changelog
 
+## 1.0.0 - 2026-10-13
+
+### Upgrade notes
+
+- Use `PolyStepOptimizer` for neural networks and `PolyStep` for standalone objectives.
+- Replace diagonal CMA with `AdaptiveSubspace`, tempered softmax with `SoftmaxSolver`,
+  and `auto_epsilon` with `ProgressiveEpsilon`. Probe reuse has been removed.
+- Checkpoints use format 6. Rebuild the optimizer from saved model weights when
+  upgrading from 0.13; older optimizer checkpoints cannot be loaded.
+- Some low-level exports and experiment helpers moved or were removed. The
+  [migration guide](docs/api_overview.md#upgrading-from-013) lists replacements and
+  changes to evaluation, configuration, and callbacks.
+
+### Fixes
+
+- Preserve live model weights, buffers, momentum, and trust-region state across
+  evaluation, subspace updates, and checkpoint restoration.
+- Handle nonfinite costs, tiny transport rows, integer losses, and half-precision
+  reductions consistently. Sparse CUDA projections repeat in deterministic mode.
+- Accept device strings in rotation sampling and float32 masks in half-precision
+  CUDA attention.
+- Correct baseline budgets, tied ranks, rollout accounting, and download handling.
+
+### Performance and examples
+
+- Reuse unchanged forward work, update only affected weights, and size candidate
+  chunks from measured activations. Speed up blockwise, factored, and adaptive steps.
+- Reduce sparse projection storage and host synchronization; compile Sinkhorn loops.
+- Refresh all 11 examples, their evaluation paths, and their documented settings.
+  See [performance](docs/performance.md) for measurements and tradeoffs.
+- Consolidate redundant tests and keep release validation focused on the built wheel.
+
+## 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,

{polystep-0.12.0 → polystep-1.0.0}/CITATION.cff

@@ -6,8 +6,8 @@ authors:
   - family-names: Le
     given-names: An T.
 license: Apache-2.0
-version: 0.12.0
-date-released: "2026-09-17"
+version: 1.0.0
+date-released: "2026-10-13"
 repository-code: "https://github.com/anindex/polystep"
 keywords:
   - gradient-free optimization
@@ -22,9 +22,9 @@ preferred-citation:
   authors:
     - family-names: Le
       given-names: An T.
-  journal: "arXiv preprint arXiv:2605.01928"
+  journal: "Transactions on Machine Learning Research"
   year: 2026
-  url: "https://arxiv.org/abs/2605.01928"
+  url: "https://openreview.net/forum?id=8mlcqTTMuU"
 references:
   - type: conference-paper
     title: "Accelerating Motion Planning via Optimal Transport"

polystep-1.0.0/CONTRIBUTING.md

@@ -0,0 +1,65 @@
+# Contributing
+
+Install the development dependencies. Then run the CI checks:
+
+```bash
+pip install -e ".[dev]"
+ruff check .
+ruff format --check .
+pytest tests/ -q -n auto
+```
+
+Work on a branch. Add a regression test when the change needs one.
+Open a pull request after the checks pass. Test observable behavior.
+Keep one small regression test per bug.
+Extend an existing case when it already covers the same path; avoid separate tests
+for trivial defaults, getters, or repeated shapes. 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 ""
+```
+
+The fast suite uses one CPU thread; `POLYSTEP_TEST_THREADS` overrides it. Slow tests
+use `ncpu - 8` threads, clamped to 2-16, and may download datasets. Keep downloads
+and training-quality checks in the slow suite; use synthetic inputs for unit tests.
+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.
+
+## Releases
+
+Update `polystep.__version__`. Set the same version and the release date in
+`CITATION.cff`. Add a `CHANGELOG.md` entry with the heading `## X.Y.Z - YYYY-MM-DD`.
+`pyproject.toml` reads the package version automatically.
+If the project's maturity changes, update its `Development Status` classifier.
+
+Build the distributions. Then check their metadata:
+
+```bash
+uv build
+uvx twine check --strict dist/*
+```
+
+Run the full manual CI workflow on the release commit.
+This covers the Python matrix and repository-only experiment tests.
+After CI passes, publish a GitHub Release with the tag `vX.Y.Z`.
+This triggers the PyPI workflow. The workflow checks version consistency, lint, and
+package metadata. It tests the installed wheel before it publishes the package.
+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.12.0 → polystep-1.0.0}/LIMITATIONS.md

@@ -13,33 +13,46 @@ 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.
+Momentum, the `amortize_steps` coasting direction, and the `biased_rotation`
+direction are cleared when the basis changes. `AdaptiveSubspace` rotates every step
+by default. These options need a longer rotation interval to take effect.
+The optimizer warns when you combine them with rotation on every step.
+Heavy-ball momentum can amplify the barycentric displacement.
+`step_radius` bounds the displacement before momentum.
 
 ## Evaluation
 
 - `step()` expects `closure(batched_params) -> losses` with shape `(N,)`.
+  The candidate tensors may share storage that the next chunk overwrites, so clone
+  any you keep after the closure returns.
 - Manual loops must register the evaluator and current batch to enable fast paths.
   The evaluator and closure must compute the same objective.
+- `NNCostEvaluator` and `train()` score candidates in eval mode: dropout is off and
+  batch-norm running statistics never update. A custom closure that runs the model in
+  train mode compares candidates under different dropout masks.
 - 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.
+  Models needing autocast internally must provide it. Half-precision candidate outputs
+  reduce in FP32 only under the builtin cross-entropy, MSE and L1 losses at their
+  defaults apart from `reduction`. Other losses may round nearby candidate scores
+  to the same value. With FP16 or BF16 parameters, `FactoredSubspace` materializes
+  every candidate and runs slower. See
+  [precision](docs/performance.md#compilation-and-precision).
+- `train()` rejects `trust_region`. See
+  [radius control](docs/api_overview.md#quadratic-model-and-radius-control) for its
+  manual-step use.
+- `trust_region` and `use_adaptive_radius` both scale the step. Choose one; the
+  optimizer warns when both are on.
 - 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.
+  bias the full-fidelity cost estimate.
+- Full-space blockwise mode does not support every monolithic adaptation option;
+  unsupported options warn at construction.
+- Automatic chunk sizing counts outputs of layers with parameters. On a site-local
+  path, it counts all module outputs from the perturbed layer onward, provided the
+  forward neither draws random numbers nor reads the perturbed parameter before
+  calling its layer. Attention scores and other temporary tensors inside a module's
+  forward are not counted. Set `chunk_size` by hand if such a model runs out of memory.
 - The library does not distribute optimizer state across multiple GPUs.
 
 ## Attention and recurrent layers
@@ -53,9 +66,7 @@ Boolean masks use `True` for masked positions; float masks are additive. Attenti
 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.
+CPU vmap may emit a batching-fallback warning.
 
 [`VmapSafeLSTM`](src/polystep/layers/rnn.py) requires batch-first dense inputs. It does
 not support bidirectional recurrence, projection, or `PackedSequence`. Its explicit
@@ -70,10 +81,22 @@ 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'`.
+FP16/BF16 costs to FP32. Cost differences and reductions must also fit the solver's
+dtype: finite FP32 inputs near its limits can still overflow when recentered or
+averaged. Rescale the objective before evaluation or return FP64 costs in that case.
+
+All nonfinite costs are invalid. `min_cost_greedy` and `top_k_mean` skip a vertex
+with any nonfinite probe, and the solve leaves a particle with no feasible vertex
+where it is. The weighted solvers replace each nonfinite cost with twice the largest
+finite magnitude plus one, clamped to the dtype limit. An infeasible candidate can
+therefore retain weight, depending on epsilon. With `scale_cost='mean'` or
+`'max_cost'`, the solvers compute the penalty after scaling. Otherwise, that weight
+also depends on the cost offset.
+Balanced Sinkhorn sends every column its prescribed mass, even when all of that
+column's candidates are infeasible. The monolithic step records a vertex's cost as
+the mean of its probes with each nonfinite probe already penalized, so with several
+probes a vertex with only some nonfinite probes can still set the best-vertex
+progress that drives the adaptive radius.
 
 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
@@ -89,4 +112,5 @@ The SNN surrogate-gradient baseline has substantial variation across seeds. MAX-
 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.
+The [reproducibility guide](docs/reproducibility.md#protocol) gives the protocols, and the
+[experiment index](experiments/EXPERIMENT_INDEX.md) has the results.

{polystep-0.12.0 → polystep-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polystep
-Version: 0.12.0
+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
@@ -10,7 +10,7 @@ Project-URL: Documentation, https://github.com/anindex/polystep/blob/main/docs/a
 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: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Intended Audience :: Developers
 Classifier: Operating System :: OS Independent
@@ -27,13 +27,12 @@ License-File: LICENSE
 Requires-Dist: torch>=2.8
 Provides-Extra: examples
 Requires-Dist: numpy>=2.0; extra == "examples"
-Requires-Dist: torchvision>=0.25; 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"
-Requires-Dist: python-sat; 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"
@@ -46,14 +45,12 @@ 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: gymnasium>=0.29; extra == "rl"
 Requires-Dist: stable-baselines3>=2.3; extra == "rl"
-Requires-Dist: swig>=4.2; extra == "rl"
 Dynamic: license-file
 
 # polystep
@@ -66,6 +63,8 @@ Dynamic: license-file
 
 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,
@@ -73,20 +72,19 @@ 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)
+[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: subspace projection, polytope probes, loss evaluation, and a weighted parameter update.">
+       alt="PolyStep: search coordinates, polytope probes, cost evaluation, soft assignment, and a barycentric update.">
 </p>
 
-> Want to play around with parameters? [**Viet T. Nguyen**](https://vietngth.github.io/) built a gorgeous interactive walkthrough that animates every step of the method -> **[explore the PolyStep visualization](https://vietngth.github.io/polystep-visualization/)**.
-
 ## Installation
 
 ```bash
 pip install polystep
-pip install "polystep[examples]"  # optional demo dependencies
+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
@@ -99,8 +97,7 @@ For development, install `pip install -e ".[dev]"` from a checkout.
 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
+from polystep import HybridSubspace, ParamLayout, PolyStepOptimizer, TrainConfig, train
 
 model = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 3))
 loader = DataLoader(
@@ -114,10 +111,11 @@ 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.
+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
 
@@ -134,6 +132,7 @@ print(es.best_fitness, es.best_solution)
 ### Synthetic objectives
 
 ```python
+import torch
 from polystep import Ackley
 from polystep.solver import PolyStep
 
@@ -145,7 +144,8 @@ state = solver.run(torch.randn(100, 10))
 
 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.
+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
@@ -155,8 +155,8 @@ 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:
+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
@@ -165,20 +165,30 @@ 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.
+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. 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).
+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 cell is not in this release.
+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 |
 |---|---|---|---|---|---|---|
@@ -206,16 +216,17 @@ Test accuracy %. A dash means the cell is not in this release.
 
 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.
+[limitations](https://github.com/anindex/polystep/blob/main/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)
+- [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
 
@@ -223,15 +234,17 @@ on the differentiable tasks, and specialized solvers lead on MAX-SAT. See
 @article{le2026training,
   title={Training Non-Differentiable Networks via Optimal Transport},
   author={Le, An T},
-  journal={arXiv preprint arXiv:2605.01928},
-  year={2026}
+  journal={Transactions on Machine Learning Research},
+  year={2026},
+  url={https://openreview.net/forum?id=8mlcqTTMuU}
 }
 ```
 
 ## Acknowledgments
 
-A huge thank you to [**Viet**](https://vietngth.github.io/) for building a beautiful interactive [PolyStep visualization](https://vietngth.github.io/polystep-visualization/), it brings the method to life and makes every step click!
+Thanks to [Viet T. Nguyen](https://vietngth.github.io/) for the
+[interactive visualization](https://vietngth.github.io/polystep-visualization/).
 
 ## License
 
-[Apache-2.0](LICENSE).
+[Apache-2.0](https://github.com/anindex/polystep/blob/main/LICENSE).

{polystep-0.12.0 → polystep-1.0.0}/README.md

@@ -8,6 +8,8 @@
 
 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,
@@ -15,20 +17,19 @@ 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)
+[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: subspace projection, polytope probes, loss evaluation, and a weighted parameter update.">
+       alt="PolyStep: search coordinates, polytope probes, cost evaluation, soft assignment, and a barycentric update.">
 </p>
 
-> Want to play around with parameters? [**Viet T. Nguyen**](https://vietngth.github.io/) built a gorgeous interactive walkthrough that animates every step of the method -> **[explore the PolyStep visualization](https://vietngth.github.io/polystep-visualization/)**.
-
 ## Installation
 
 ```bash
 pip install polystep
-pip install "polystep[examples]"  # optional demo dependencies
+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
@@ -41,8 +42,7 @@ For development, install `pip install -e ".[dev]"` from a checkout.
 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
+from polystep import HybridSubspace, ParamLayout, PolyStepOptimizer, TrainConfig, train
 
 model = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 3))
 loader = DataLoader(
@@ -56,10 +56,11 @@ 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.
+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
 
@@ -76,6 +77,7 @@ print(es.best_fitness, es.best_solution)
 ### Synthetic objectives
 
 ```python
+import torch
 from polystep import Ackley
 from polystep.solver import PolyStep
 
@@ -87,7 +89,8 @@ state = solver.run(torch.randn(100, 10))
 
 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.
+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
@@ -97,8 +100,8 @@ 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:
+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
@@ -107,20 +110,30 @@ 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.
+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. 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).
+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 cell is not in this release.
+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 |
 |---|---|---|---|---|---|---|
@@ -148,16 +161,17 @@ Test accuracy %. A dash means the cell is not in this release.
 
 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.
+[limitations](https://github.com/anindex/polystep/blob/main/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)
+- [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
 
@@ -165,15 +179,17 @@ on the differentiable tasks, and specialized solvers lead on MAX-SAT. See
 @article{le2026training,
   title={Training Non-Differentiable Networks via Optimal Transport},
   author={Le, An T},
-  journal={arXiv preprint arXiv:2605.01928},
-  year={2026}
+  journal={Transactions on Machine Learning Research},
+  year={2026},
+  url={https://openreview.net/forum?id=8mlcqTTMuU}
 }
 ```
 
 ## Acknowledgments
 
-A huge thank you to [**Viet**](https://vietngth.github.io/) for building a beautiful interactive [PolyStep visualization](https://vietngth.github.io/polystep-visualization/), it brings the method to life and makes every step click!
+Thanks to [Viet T. Nguyen](https://vietngth.github.io/) for the
+[interactive visualization](https://vietngth.github.io/polystep-visualization/).
 
 ## License
 
-[Apache-2.0](LICENSE).
+[Apache-2.0](https://github.com/anindex/polystep/blob/main/LICENSE).