polystep 1.0.0__tar.gz → 1.1.0__tar.gz

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

@@ -1,5 +1,56 @@
 # Changelog
 
+## 1.1.0 - 2026-10-24
+
+### Features and performance
+
+- Add `step_batch(inputs, targets, loss_fn=...)` for supervised training. It shares
+  the fast evaluation paths with `train()`; both accept tuple and dict inputs.
+  `register_evaluator` now returns a closure for the registered batch.
+- Add `is_causal` support to `VmapSafeMultiHeadAttention`.
+- Extend fast MLP evaluation to functional losses and traceable custom forwards.
+  Reuse unchanged forward work, reduce CPU cross-entropy overhead, and replay
+  site-local CUDA forwards with CUDA graphs.
+- Measured full-space MLP steps are 1.9–3.3x faster on one CPU thread and 2.6x on a
+  shared RTX 5090 than 1.0.0. Tiny tied-weight models can be slower on CUDA. See
+  [performance](docs/performance.md) for the evaluation paths and tradeoffs.
+- Refresh the examples and their metrics, clarify simulated Loihi adaptation, and
+  add SNN training and CartPole GIFs. The MAX-SAT example reaches about 98% SAT with
+  600 default steps instead of 1500.
+
+### Fixes
+
+- Restore forward-written buffers between scoring calls and vmap chunks, even on
+  failure, and preserve scores that share those buffers. Protect frozen embedding
+  weights changed by `max_norm`, reject evaluators for another model, and follow
+  model dtype and device changes.
+- Preserve finite-difference gradients and curvature at small probe radii. Correct
+  KL zero-support handling and marginal diagnostics, including tiny target masses.
+- Fix CUDA rotation for large adaptive subspaces, carry FP16/BF16 absorb rounding
+  error forward, and correct biased rotation, screening, amortized schedules,
+  epsilon restarts, half-precision attention, and loss precision in ES and baselines.
+- Fall back to sequential scoring when vmap cannot batch a model, including stock
+  `nn.LSTM`.
+
+### Upgrade notes
+
+- Pass the model to `PolyStepOptimizer`, not `model.parameters()`. Rebuild the
+  optimizer if its trainable parameter set changes. Separate parameters or buffers
+  that share trainable storage now raise `ValueError`; tie weights by assigning
+  the same Parameter, or register a clone.
+- `FactoredSubspace` redraws `B` every 5 OT steps by default; set
+  `rotation_interval=0` for the old behavior. Hybrid and factored seeded bases change.
+- `SinkhornSolver` and `KLSoftmaxSolver` default to `threshold=3e-5`, using relative
+  marginal error and the dual step divided by epsilon, respectively. Stalled solves
+  above the threshold report `converged=False`.
+- `train()` uses `step_batch`; `restore_best` compares weights on the first training
+  batch. Seeded results can differ from 1.0.
+- For compiled hybrid/factored runs, rebuild the optimizer from saved model weights:
+  checkpoint parameter names no longer contain the `_orig_mod.` prefix. The
+  [upgrade guide](docs/api_overview.md#upgrading-from-10) covers all compatibility
+  changes; [reproducibility](docs/reproducibility.md#changes-since-the-published-tables)
+  lists the affected published results.
+
 ## 1.0.0 - 2026-10-13
 
 ### Upgrade notes
@@ -10,8 +61,8 @@
 - 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.
+  [migration guide](docs/api_overview.md#upgrading-from-013) covers the main API
+  replacements and checkpoint requirements.
 
 ### Fixes
 

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

@@ -6,8 +6,8 @@ authors:
   - family-names: Le
     given-names: An T.
 license: Apache-2.0
-version: 1.0.0
-date-released: "2026-10-13"
+version: 1.1.0
+date-released: "2026-10-24"
 repository-code: "https://github.com/anindex/polystep"
 keywords:
   - gradient-free optimization

polystep-1.1.0/CONTRIBUTING.md

@@ -0,0 +1,65 @@
+# Contributing
+
+From the repository root, install the development dependencies and run the CI checks:
+
+```bash
+pip install -e ".[dev]"
+ruff check .
+ruff format --check .
+pytest tests/ -q -n auto --maxprocesses=4
+```
+
+A bug fix should have a small regression test of observable behavior. Extend an
+existing case when it already covers the same path; avoid separate tests for trivial
+defaults, getters, or repeated shapes. Keep public APIs typed. Comments should
+explain constraints, numerical choices, or non-obvious behavior.
+
+## Tests
+
+The default suite excludes `slow` and `gpu` tests. To select them:
+
+```bash
+pytest tests/ -m slow
+pytest tests/ -m gpu
+pytest tests/ -m ""
+```
+
+Capping parallel test workers at four limits startup and memory costs. The fast
+suite uses one CPU thread per worker; `POLYSTEP_TEST_THREADS` overrides it. Slow
+tests use `ncpu - 8` threads, clamped to 2-16 and split across xdist workers, 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. Install `pip install -e ".[dev,experiments]"` to run
+tests that need optional experiment dependencies. Missing optional dependencies and
+unavailable hardware are valid reasons to skip a test; unsupported parameter pairs
+should not be collected. Tests run without automatic retries.
+
+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, optional
+experiment dependencies, 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.
+
+```bash
+uv build
+uvx twine check --strict dist/*
+```
+
+After the full manual CI workflow passes on the release commit, publish a GitHub
+Release tagged `vX.Y.Z`. The release workflow checks versions, lint, metadata, and
+the installed wheel before publishing. Only the 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-1.0.0 → polystep-1.1.0}/LIMITATIONS.md

@@ -10,41 +10,66 @@ directions; very small projection ratios can fail on large models.
 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.
+`FactoredSubspace` restricts each step to the input directions spanned by its factor
+`B`. It redraws `B` every 5 OT steps by default; with `rotation_interval=0` the whole
+run stays in one set of directions. Retune radii when switching subspaces.
 
-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.
+A basis change clears momentum, the `amortize_steps` coasting direction and the
+`biased_rotation` direction, so these options need a rotation interval longer than one
+step. `AdaptiveSubspace` rotates every step by default, and the optimizer warns when
+any subspace combines these options with `rotation_interval=1`. `FactoredSubspace`
+clears them every 5 OT steps without a warning. Under `amortize_steps`,
+`rotation_interval` and `absorb_interval` count OT steps only.
+
+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.
+  The candidate tensors may be expanded views or share storage that the next chunk
+  overwrites, so treat them as read-only and clone any you keep after the closure
+  returns.
+- `step_batch()` and `train()` register each batch. A manual `step(closure)` loop must
+  register the evaluator and current batch to use the fast paths. The evaluator's
+  loss and the closure must compute the same objective.
+- Multi-input batches use vmap and the site-local path; the batched-matmul, delta and
+  factored paths and multifidelity screening need one input tensor.
+- Tied weights must be the same Parameter object, and buffers may not share memory
+  with trainable parameters. Register a clone or compute the tensor in forward instead.
+- `NNCostEvaluator`, `step_batch()` and `train()` score candidates in eval mode:
+  standard dropout is off and batch-norm layers do not update their running statistics.
+  A custom closure that runs the model in train mode compares candidates under
+  different dropout masks.
+- The evaluator restores buffers written by the forward, such as a call counter,
+  after each batched chunk and each candidate in the sequential and in-place loops.
+  Embedding weights changed by `max_norm` are restored even when frozen. Update
+  persistent state outside the closure.
+- Writes that leave the tensor version counter unchanged, such as
+  `F.batch_norm(training=True)` running statistics or writes through `.data`, are found
+  by comparing values after the first forward and the first candidate of each
+  sequential, in-place or chunked vmap sweep. A write first triggered after those
+  checks can go undetected and persist into later evaluations.
+- Float inputs are cast to the dtype of the model's first floating weight only when
+  the forward raises on them as given. The evaluator checks this with one extra
+  forward per input dtype and puts back the buffers it changed. The delta and
+  factored paths, which accept plain Linear stacks only, always cast.
 - 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. 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
+  reduce in FP32 only under the builtin cross-entropy, MSE and L1 losses; other losses
+  may round nearby candidate scores to the same value. With FP16 or BF16 parameters,
+  `FactoredSubspace` materializes every candidate, and the fast paths can differ from
+  the model's own Linear output by about one unit in the last place (ulp), enough to
+  flip a hard threshold. `use_inplace=True` scores with the model's forward. See
   [precision](docs/performance.md#compilation-and-precision).
-- `train()` rejects `trust_region`. See
+- `step_batch()` and `train()` reject `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
+- Multifidelity screening imputes dropped costs using an additive offset. This can
   bias the full-fidelity cost estimate.
 - Full-space blockwise mode does not support every monolithic adaptation option;
   unsupported options warn at construction.
@@ -53,14 +78,21 @@ Heavy-ball momentum can amplify the barycentric displacement.
   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.
+  Under vmap, CUDA attention materializes the scores of every candidate in the chunk.
+  `PolyStepOptimizer(chunk_size=...)` caps the candidates scored per step chunk, and
+  `NNCostEvaluator(chunk_size=...)` caps how many of them each vmap call batches.
+- On CUDA the site-local path replays CUDA graphs, which freeze Python-side state the
+  forward reads, such as a Python counter or a flag. Pass `compile_forward=False` for
+  such a forward.
 - 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.
+`add_zero_attn`, or returned attention weights. `is_causal` is a hint that needs
+`attn_mask`, as in `nn.MultiheadAttention`. 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.
@@ -68,6 +100,19 @@ The wrapper also accepts `(batch, query, key)` masks.
 
 CPU vmap may emit a batching-fallback warning.
 
+Stock `nn.MultiheadAttention` and `nn.TransformerEncoderLayer` work, but vmap has no
+batching rule for their fused kernels, so it runs them once per candidate and PyTorch
+prints a performance-drop warning. On CUDA this makes a step several times slower, and
+the evaluator warns and points to `VmapSafeMultiHeadAttention`. On CPU stock
+`nn.MultiheadAttention` measured faster than `VmapSafeMultiHeadAttention`, so the
+evaluator stays silent there.
+
+Stock `nn.LSTM`, `nn.GRU` and `nn.RNN` have no vmap batching rule, so the evaluator
+scores them one candidate at a time in a sequential loop, with a warning. On a small
+CPU model, `NNCostEvaluator(..., use_inplace=True)` ran 1.1 to 1.3 times faster than
+the loop, and `VmapSafeLSTM` 3 to 5 times faster. `step_batch()` and `train()` build
+their own evaluator, so they take the loop for these layers.
+
 [`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.
@@ -80,7 +125,7 @@ Independent particle rotations mean Sinkhorn's shared column index does not iden
 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
+are inactive in fixed-iteration mode (`threshold <= 0`). Sinkhorn and KLSoftmax promote
 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.
@@ -92,6 +137,7 @@ finite magnitude plus one, clamped to the dtype limit. An infeasible candidate c
 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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polystep
-Version: 1.0.0
+Version: 1.1.0
 Summary: PyTorch PolyStep Optimizer - gradient-free neural network training via entropic optimal transport
 Author: An T. Le
 License-Expression: Apache-2.0
@@ -35,7 +35,6 @@ 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"
@@ -75,10 +74,11 @@ backpropagation is usually faster.
 [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>
+| One particle | Four particles, two minima |
+|:---:|:---:|
+| [<img src="docs/figures/polystep_particle.gif" width="440" alt="One particle probes an objective and moves toward its minimum through weighted barycentric steps.">](https://vietngth.github.io/polystep-visualization/) | [<img src="docs/figures/polystep_particles.gif" width="440" alt="Four particles probe a two-basin objective and move toward its minima through weighted barycentric steps.">](https://vietngth.github.io/polystep-visualization/) |
+
+Animations from [Viet T. Nguyen's interactive visualization](https://vietngth.github.io/polystep-visualization/).
 
 ## Installation
 
@@ -89,7 +89,7 @@ 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.
+For development, run `pip install -e ".[dev]"` from a checkout.
 
 ## Quickstart
 
@@ -99,6 +99,7 @@ from torch import nn
 from torch.utils.data import DataLoader, TensorDataset
 from polystep import HybridSubspace, ParamLayout, PolyStepOptimizer, TrainConfig, train
 
+torch.set_num_threads(1)
 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,))),
@@ -111,9 +112,24 @@ 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
+`max_subspace_dim=128` caps the search at 128 coordinates, so each step scores
+144 candidates. Biases and normalization weights keep their full width. On a larger
+model, the cap must cover those parameters plus at least one coordinate per weight
+matrix. If the cap is too small, PolyStep
+warns and reports the minimum dimension it can use.
+
+`step_batch` uses the same candidate evaluation paths as `train()`:
+
+```python
+loss_fn = nn.CrossEntropyLoss()
+for x, y in loader:
+    optimizer.step_batch(x, y, loss_fn=loss_fn)
+optimizer.release_evaluator()
+```
+
+`torch.set_num_threads(1)` avoids thread overhead for a model this small. Larger
+models may benefit from more threads; `python -m polystep.benchmarks.threads` compares
+several counts on an MLP workload. See
 [performance](https://github.com/anindex/polystep/blob/main/docs/performance.md)
 for candidate paths, thread counts, compilation, and memory settings.
 
@@ -125,8 +141,8 @@ 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)
+    es.tell((candidates - 1).square().sum(dim=1))
+print(es.best_fitness, (es.mean - 1).square().sum().item())
 ```
 
 ### Synthetic objectives
@@ -152,6 +168,17 @@ depends on the radius, representation, and evaluation budget. PolyStep's optiona
 quadratic models and radius controllers are implementation heuristics; they do not
 inherit classical direct-search or trust-region guarantees.
 
+| SNN training | CartPole policy |
+|:---:|:---:|
+| <img src="docs/figures/polystep_snn_progress.gif" width="440" alt="Training loss and batch accuracy over 60 PolyStep updates of a spiking network with hard LIF thresholds."> | <img src="docs/figures/rl_cartpole_policy.gif" width="440" alt="A trained CartPole policy uses hard left/right actions to keep the pole upright."> |
+
+The [SNN example](examples/02_snn_starter.py) trains a network with hard LIF spikes
+on four synthetic classes, without surrogate gradients.
+
+The [CartPole example](examples/03_rl_cartpole.py) uses real-valued weights and a
+hard left/right `argmax`. PolyStep optimizes episode return without gradients
+through the actions or simulator.
+
 ## Baselines
 
 `polystep.baselines` provides `openai_es`, `spsa`, `mezo`, `random_search`, `eggroll`,
@@ -161,8 +188,8 @@ 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)
+objective = Objective(lambda x: (x - 1).square().sum(1), dim=64, budget=10_000)
+result = openai_es(objective, sigma=0.05, popsize=32)
 ```
 
 The paper runners match search spaces, budgets and tuning as the
@@ -171,17 +198,12 @@ The paper runners match search spaces, budgets and tuning as the
 ## 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
+selected on validation data; the probSAT and RC2 columns are single runs. Each
+baseline gets a matched evaluation or step budget, as the
+[protocol](https://github.com/anindex/polystep/blob/main/docs/reproducibility.md#protocol)
+describes. 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
+lists the rows that later releases change, and the
 [experiment index](https://github.com/anindex/polystep/blob/main/experiments/EXPERIMENT_INDEX.md)
 has the full tables.
 
@@ -225,6 +247,7 @@ for unsupported configurations and comparison limits.
 - [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)
+- [Limitations](https://github.com/anindex/polystep/blob/main/LIMITATIONS.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)
 

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

@@ -20,10 +20,11 @@ backpropagation is usually faster.
 [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>
+| One particle | Four particles, two minima |
+|:---:|:---:|
+| [<img src="docs/figures/polystep_particle.gif" width="440" alt="One particle probes an objective and moves toward its minimum through weighted barycentric steps.">](https://vietngth.github.io/polystep-visualization/) | [<img src="docs/figures/polystep_particles.gif" width="440" alt="Four particles probe a two-basin objective and move toward its minima through weighted barycentric steps.">](https://vietngth.github.io/polystep-visualization/) |
+
+Animations from [Viet T. Nguyen's interactive visualization](https://vietngth.github.io/polystep-visualization/).
 
 ## Installation
 
@@ -34,7 +35,7 @@ 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.
+For development, run `pip install -e ".[dev]"` from a checkout.
 
 ## Quickstart
 
@@ -44,6 +45,7 @@ from torch import nn
 from torch.utils.data import DataLoader, TensorDataset
 from polystep import HybridSubspace, ParamLayout, PolyStepOptimizer, TrainConfig, train
 
+torch.set_num_threads(1)
 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,))),
@@ -56,9 +58,24 @@ 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
+`max_subspace_dim=128` caps the search at 128 coordinates, so each step scores
+144 candidates. Biases and normalization weights keep their full width. On a larger
+model, the cap must cover those parameters plus at least one coordinate per weight
+matrix. If the cap is too small, PolyStep
+warns and reports the minimum dimension it can use.
+
+`step_batch` uses the same candidate evaluation paths as `train()`:
+
+```python
+loss_fn = nn.CrossEntropyLoss()
+for x, y in loader:
+    optimizer.step_batch(x, y, loss_fn=loss_fn)
+optimizer.release_evaluator()
+```
+
+`torch.set_num_threads(1)` avoids thread overhead for a model this small. Larger
+models may benefit from more threads; `python -m polystep.benchmarks.threads` compares
+several counts on an MLP workload. See
 [performance](https://github.com/anindex/polystep/blob/main/docs/performance.md)
 for candidate paths, thread counts, compilation, and memory settings.
 
@@ -70,8 +87,8 @@ 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)
+    es.tell((candidates - 1).square().sum(dim=1))
+print(es.best_fitness, (es.mean - 1).square().sum().item())
 ```
 
 ### Synthetic objectives
@@ -97,6 +114,17 @@ depends on the radius, representation, and evaluation budget. PolyStep's optiona
 quadratic models and radius controllers are implementation heuristics; they do not
 inherit classical direct-search or trust-region guarantees.
 
+| SNN training | CartPole policy |
+|:---:|:---:|
+| <img src="docs/figures/polystep_snn_progress.gif" width="440" alt="Training loss and batch accuracy over 60 PolyStep updates of a spiking network with hard LIF thresholds."> | <img src="docs/figures/rl_cartpole_policy.gif" width="440" alt="A trained CartPole policy uses hard left/right actions to keep the pole upright."> |
+
+The [SNN example](examples/02_snn_starter.py) trains a network with hard LIF spikes
+on four synthetic classes, without surrogate gradients.
+
+The [CartPole example](examples/03_rl_cartpole.py) uses real-valued weights and a
+hard left/right `argmax`. PolyStep optimizes episode return without gradients
+through the actions or simulator.
+
 ## Baselines
 
 `polystep.baselines` provides `openai_es`, `spsa`, `mezo`, `random_search`, `eggroll`,
@@ -106,8 +134,8 @@ 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)
+objective = Objective(lambda x: (x - 1).square().sum(1), dim=64, budget=10_000)
+result = openai_es(objective, sigma=0.05, popsize=32)
 ```
 
 The paper runners match search spaces, budgets and tuning as the
@@ -116,17 +144,12 @@ The paper runners match search spaces, budgets and tuning as the
 ## 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
+selected on validation data; the probSAT and RC2 columns are single runs. Each
+baseline gets a matched evaluation or step budget, as the
+[protocol](https://github.com/anindex/polystep/blob/main/docs/reproducibility.md#protocol)
+describes. 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
+lists the rows that later releases change, and the
 [experiment index](https://github.com/anindex/polystep/blob/main/experiments/EXPERIMENT_INDEX.md)
 has the full tables.
 
@@ -170,6 +193,7 @@ for unsupported configurations and comparison limits.
 - [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)
+- [Limitations](https://github.com/anindex/polystep/blob/main/LIMITATIONS.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)