polystep 0.13.0__tar.gz → 1.1.0__tar.gz

polystep-1.1.0/CHANGELOG.md

@@ -0,0 +1,186 @@
+# 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
+
+- 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) covers the main API
+  replacements and checkpoint requirements.
+
+### 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,
+  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.13.0 → polystep-1.1.0}/CITATION.cff

@@ -6,8 +6,8 @@ authors:
   - family-names: Le
     given-names: An T.
 license: Apache-2.0
-version: 0.13.0
-date-released: "2026-09-22"
+version: 1.1.0
+date-released: "2026-10-24"
 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.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.1.0/LIMITATIONS.md

@@ -0,0 +1,162 @@
+# 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 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.
+
+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 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; 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).
+- `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.
+- 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.
+- 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.
+  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`, 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.
+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.
+
+## 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`). 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.
+
+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
+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.
+
+The [reproducibility guide](docs/reproducibility.md#protocol) gives the protocols, and the
+[experiment index](experiments/EXPERIMENT_INDEX.md) has the results.