Skip to content

Latest commit

 

History

History
175 lines (148 loc) · 10.5 KB

LIMITATIONS.md

File metadata and controls

175 lines (148 loc) · 10.5 KB

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.
  • 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 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. 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=True scores with the model's forward. See precision.
  • step_batch() and train() reject trust_region. See 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 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.

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. 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.

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.

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.