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.
step()expectsclosure(batched_params) -> losseswith 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()andtrain()register each batch. A manualstep(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()andtrain()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_normare restored even when frozen. Update persistent state outside the closure. - Some writes leave the tensor version counter unchanged. These include
F.batch_norm(training=True)running statistics and writes through.data. The evaluator compares values after the first forward to detect these writes. It repeats this check after 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
HybridSubspaceorFactoredSubspace. Full-space and global-subspace coordinates use one dtype and reject mixed parameter dtypes. mixed_precision=Trueuses 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,FactoredSubspacematerializes every candidate. Fast paths can differ from the model's own Linear output by about one unit in the last place (ulp). This can flip a hard threshold.use_inplace=Truescores with the model's forward. See precision.step_batch()andtrain()rejecttrust_region. See radius control for its manual-step use.trust_regionanduse_adaptive_radiusboth 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_sizeby 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, andNNCostEvaluator(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=Falsefor such a forward. - The library does not distribute optimizer state across multiple GPUs.
VmapSafeMultiHeadAttention 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 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.
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. With scale_cost='mean' or 'max_cost', solvers rescale
before subtracting the minimum, so extreme finite FP32 costs remain usable. Without
these modes, cost differences must fit the solver's dtype. Rescale the objective
before evaluation when the forward itself can overflow; solver scaling cannot
recover a finite value from an infinite input.
Balanced Sinkhorn needs equal positive total mass in both marginals. Zero entries carry exactly zero transport, and positive entries are logged without a mass floor. Masses and plan entries must still be representable in the solver's dtype. A result can underflow when individual entries are much smaller than their row sums.
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, after penalizing each nonfinite probe. With several probes, a partly infeasible vertex 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.
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.
Example 06 is a PyTorch SNN simulation. Its one-state neuron uses subtractive reset; it does not implement Loihi fixed-point arithmetic or the current and voltage states of Lava LIF. No hardware latency, energy, or adaptation claim follows from its results.
The reproducibility guide gives the protocols, and the experiment index has the results.