NSA-Flow: Non-negative Stiefel-Approximating Flow

A calibration-free relaxation between reconstruction and feature clustering.

import torch
from nsa_flow import nsa_flow

X0 = torch.rand(120, 8)          # a loading matrix to refine, e.g. |PCA loadings|
r  = nsa_flow(X0, w=0.7)         # w is a genuine convex weight in [0, 1]

print(r)
# NSAResult(w=0.7, iters=61, energy=..., fidelity=..., defect=...,
#           eff_rank=7.41, stop=grad_map, |Gmap|=4.2e-13)
print(r.Y.shape)                 # (120, 8), non-negative, near-orthogonal columns

What it computes

Minimise, over Y >= 0,

E_w(Y) = (1 - w) * ||Y - X0||_F^2 / ||X0||_F^2   +   w * D(Y) / (1 - 1/k)

where, with G = Y'Y / tr(Y'Y) the trace-normalised Gram matrix,

D(Y) = || G - I/k ||_F^2 = ||Y'Y||_F^2 / ||Y||_F^4 - 1/k

Both terms are dimensionless and O(1), so w needs no calibration: no normalising constants estimated from the data, no warm-up phase, no dependence on p, k or the scale of X0. A reported w means the same thing on every problem.

Why this D

D is the squared distance from the normalised Gram matrix to isotropy. It

• vanishes exactly on R_{>0} * St(p, k) — orthogonal columns of equal norm;
• equals k * Var(eigenvalues of G) = 1/EffectiveRank - 1/k, so minimising it drives the effective rank to k;
• is invariant under Y -> cY, Y -> UY and Y -> YV for orthogonal U, V — the same group that preserves the constraint being relaxed;
• charges at least 1/r - 1/k for rank r < k, so rank collapse is penalised;
• is bounded: 0 <= D <= 1 - 1/k;
• satisfies <grad D, Y> = 0, so it cannot change ||Y|| — scale is pinned by the fidelity term alone, and no renormalisation step is needed.

It splits as

D  =  sum_{i != j} G_ij^2   +   sum_i (G_ii - 1/k)^2
      \_______________/         \___________________/
        decorrelation              norm balance

The first term alone is the "invariant orthogonality defect" used in v1 and elsewhere. Dropping the second is what made that functional blind to conditioning, basis-dependent, and minimised by rank-deficient matrices.

What w does

w = 0 returns max(0, X0). As w -> 1, D -> 0, and since Y >= 0 with Y'Y diagonal forces pairwise disjoint column supports, the limit is a hard clustering of the p features into k groups — the feasible set of orthogonal NMF, equivalent to k-means. In between, the overlap is bounded:

max_{i != j} <y_i, y_j>  <=  sqrt(D) * ||Y||_F^2

so "approximately disjoint factors" is a claim with a number attached.

Solver

Pure-PyTorch L-BFGS-B with optional native C++ kernel acceleration (see below); every accepted step is certified by one shared gradient-mapping norm, result.grad_map, and result.stop_reason / result.certificate say exactly what was and was not established — never a silent claim of convergence.

The inner loop forms one Gram product and two [p,k] x [k,k] products: O(p k^2), with no SVD, eigendecomposition or QR. Typical convergence is 15–300 deterministic iterations.

Native C++ Kernel Acceleration (v3.2.6)

When compiled, nsa_flow._native._lbfgsb_cpu executes fused iterations in C++ on CPU, eliminating Python interpreter and ATen dispatch overhead while achieving bit-exact numerical parity (< 1e-12 energy difference) with pure PyTorch.

• Speedup: 5–12× speedup on standard problem shapes:
  • Tall small data (500 x 50, k=5): 11.8× speedup (7.3 ms vs 86.6 ms; 49.8 µs/iter vs 601.7 µs/iter).
  • ADNI cortical signed (64 x 10, k=10): 11.1× speedup (9.1 ms vs 101.1 ms; 61.4 µs/iter vs 683.1 µs/iter).
  • Tall medium data (1000 x 100, k=10): 5.0× speedup (20.9 ms vs 104.9 ms; 142.4 µs/iter vs 723.7 µs/iter).
  • Wide genomics data (50 x 2000, k=5): 1.8× speedup (185.5 ms vs 334.3 ms; 1.25 ms/iter vs 2.26 ms/iter).
• Direct CBLAS & Zero Inner Allocations: Cauchy point, subspace minimization, compact quasi-Newton $M$ inversion, and line search operate on preallocated workspace buffers and direct CBLAS/LAPACK calls with stack solvers for small-to-moderate dimensions.
• v3.2.6 Optimizations:
  • CBLAS SYRK: Rank-k symmetric updates replace general matrix products for $A = X_V^T X_V$, $B = V^T V$, and $BW = W^T W$, cutting FLOPs by 50% for quadratic terms.
  • Zero-Copy Free Subspaces: Direct contiguous pointer indexing for $W_{\text{free}}$ when all coordinates are free ($n_{\text{free}} == n$).
  • Lazy Line-Search Copies: Defers workspace buffer copying until trial 1 fails Wolfe conditions, eliminating overhead on >85% of iterations.
  • Noise-Floor Stall Gating: Anchors stall classification to max(stall_slack * tol, sqrt(eps)) to certified plateaus at the float precision floor without false certification on far points.
• Pure-Python fallback: To force pure-Python execution, set NSA_FLOW_DISABLE_NATIVE=1. All reference instances match to $< 10^{-12}$.

Empirically E_w has a unique optimum for w < 1 — 24 random restarts agree to machine precision on every problem family tested — so there are no restarts, schedules or step-size heuristics to tune.

Evaluation results (v3.0.0, certified solver)

All numbers below were produced by the pure-PyTorch L-BFGS-B default with the shared certificate, on bases fitted inside every training fold or split. Earlier releases fitted the Golub basis on all 72 samples and reported the ADNI result from one 80/20 split of an unconverged solve; both are corrected here and the differences are stated.

Golub leukemia, 3 classes (B-ALL / T-ALL / AML; p=2000, k=3; 5-fold stratified CV, balanced accuracy; experiments/rapid_primary_benchmarks.py)

At w=0 the signed lifting is PCA and now returns it exactly (the previous solver reported defect_D = 0.10 there). The old transductive numbers were 0.06–0.07 higher across the board; that was leakage.

ADNI cortical thickness → CDRSB (n≈300, p=66, k=5; 20 repeated 80/20 splits, ΔR² against PCA on the same splits, 95% paired-t interval; experiments/adni_cdrsb_repeated.py)

Split-to-split SD of R² is 0.24–0.31, so the single-split value in earlier releases (+0.18 forest, from a solve at defect_D = 0.006) was not evidence either way. The effect survives with a converged solver: smaller, and now with an interval.

Public data (experiments/benchmark_new_public_data.py, 5-fold; defect column is defect_D for every method, comparable to PCA)

Signed w=0 equals PCA exactly on both classification sets; the previous solver reported defect_D of 0.04 (Sonar) and 0.50 (Prostate) at that setting. Fits are 10–20× faster than in 2.15 (Prostate: 12.6 s → 1.3 s per fit).

Speed against PCA (top_k_eigenvectors, exact, on-device; float64 CPU, planted structure, w=0.5, default tolerance; v3.1.0, measured on an otherwise idle machine)

v3.1.0 is ~3× faster than v3.0.0 on every row (ADNI-like 262 → 93 ms, Golub-like signed 2080 → 681 ms). Two changes did it, neither of them torch.compile: the L-BFGS-B subspace step is branchless (no per-iteration host syncs), and the signed lifting now uses the matrix-free X'(XV) route when p > n instead of forming the p×p Gram — the data mode already did. torch.compile on the pure-tensor pieces was measured at +5–9% and is available via nsa_flow.lbfgsb.set_compile(True) but off by default.

NSA-Flow is an iterative constrained method; PCA is one factorisation. A fit on the paper's imaging shape is a tenth of a second. The objective is ~10% of each iteration; the rest is L-BFGS-B's active-set bookkeeping in Python (~0.5–2 ms/iteration of tensor dispatch), which is also why the same code on MPS runs at ~3 ms/gradient with zero host↔device copies. It is not within an order of magnitude of PCA on small problems and this README does not claim it is. A compiled (C++/Metal) iteration is the remaining lever.

Torch layers

Two routes, both sound:

from nsa_flow import NSAFlowLinear, NSAFlowConv2d

# (preferred) penalty: a standard layer plus a regulariser
layer = NSAFlowLinear(256, 32, w=0.0)
loss = task_loss + 0.3 * layer.defect()        # defect() is O(p k^2), exact gradient

# parameterisation: the effective weight is driven toward the feasible set
layer = NSAFlowLinear(256, 32, w=0.5)                     # Stiefel blend, exact at w=1
layer = NSAFlowLinear(256, 32, w=0.5, nonneg="hard")      # non-negative AND near-orthogonal
layer.project_()                                          # exact anchored prox, in place, no autograd

Unified continuous blend and certified solver semantics (v3.3.2). Across all layers (NSAFlowLinear, NSAFlowConv2d, NSAFlowLayer), operator="auto" is the default:

• Signed (nonneg=None): uses the fast differentiable blend $W_{\text{eff}} = (1 - w) W + w,\text{project_scaled_stiefel}(W)$.
• Non-negative (nonneg="hard" / "softplus", default operator="auto" $\to$ "solver"): Forward pass computes the exact anchored prox via nsa_flow(mode="anchored", fidelity="anchor") — identical to layer.project_() and linear SiMLR. With nonneg="softplus", $W_{\text{base}} = \text{softplus}(W)$ is pre-applied and used in the Straight-Through Estimator (STE: $W_{\text{base}} + (W_{\text{eff}} - W_{\text{base}}).\text{detach}()$) so backward gradients scale smoothly by $\sigma(W)$. Inference is cached in eval().
• Explicit unrolled blend (operator="blend"): uses $W_{\text{eff}} = (1 - w) W_0 + w P(W_0)$ where $P(W_0)$ is the in-orthant defect flow run for n_steps iterations (default 8) independently of $w$, compiling a static graph without Dynamo recompilation overhead across $w$ sweeps.

NSAFlowLayer applies the same map per sample to a [B, p, k] batch.

Cost (66×5, one CPU thread; from experiments-style timing):

nn.Linear fwd 10 us
nonneg=None     w=0.5: train fwd   61 us (was 58) | eval fwd 9 us
nonneg=hard     w=0.5: train fwd  377 us (was 539) | eval fwd 9 us
nonneg=hard     w=1.0: train fwd  731 us (was 719) | eval fwd 9 us
nonneg=softplus w=1.0: train fwd  649 us (was 842) | eval fwd 9 us
operator=solver w=0.5: train fwd ~500 us          | eval fwd 9 us
MPS blend w=0.5 831 us (was 1227) | MPS hard w=1 5383 us (was 6372)

In eval() the effective weight is cached until the parameter changes, so inference costs the same as nn.Linear (9 µs). compile=True on any layer compiles the defect flow as one graph (fullgraph=True): measured 732 → 154 µs for the hard, w=1 training forward, 393 → 66 µs for the flow alone.

Documentation (ANTsPy Style)

Standard Sphinx documentation is built using sphinx_rtd_theme:

make docs        # builds HTML docs in docs/_build/html with 0 warnings
make docs-clean  # removes build cache

The documentation includes thematic guides for layers, solvers, energy functionals, data-anchored reconstruction, scikit-learn estimators, and diagnostic certificates, plus recursive API references.

API

High-Level Unified Interface

nsa_flow serves as the unified wrapper across all problem modes:

from nsa_flow import nsa_flow, NSAFlow

# 1. Non-negative data -> fits non-negative basis V >= 0 (auto data mode)
r_data = nsa_flow(X_positive, k=5, w=0.5)

# 2. Signed or centered data -> signed contrast lifting V = V+ - V- (auto signed mode)
r_signed = nsa_flow(X_centered, k=5, w=0.5, consolidate=True)

# 3. Target loadings matrix -> anchored flow refinement
r_anc = nsa_flow(PCA_loadings, w=0.5)

# 4. Scikit-learn Pipeline Integration
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LogisticRegression

pipe = Pipeline([
    ("nsa", NSAFlow(n_components=5, w=0.5, consolidate=True)),
    ("clf", LogisticRegression())
])
pipe.fit(X_train, y_train)

Optimizers & Performance

All solvers default to optimizer="lbfgsb": L-BFGS-B (Byrd, Lu, Nocedal & Zhu 1995). On CPU, a native C++ extension (nsa_flow._native._lbfgsb_cpu, built via Accelerate/BLAS with zero heap allocation in the inner loop) accelerates the solver by 6–11× (ADNI cortical thickness: 47.7 µs/iter vs 501.3 µs/iter in pure PyTorch; 2–4× on larger shapes) while maintaining bit-exact energy parity (< 1e-12). Setting NSA_FLOW_DISABLE_NATIVE=1 cleanly switches back to the pure PyTorch reference solver. The implementation runs unmodified on CPU, CUDA and MPS in float32 or float64 with no host↔device transfers (the top-k eigenvector initialisation is computed on-device too; see nsa_flow/linalg.py). It was chosen by measurement (experiments/optimizer_study.py): it matches SciPy's Fortran L-BFGS-B energy to twelve figures on every problem tried, finds a strictly lower minimum on the hardest, and is the only pure-torch method that certified convergence on every configuration. fista, spg and pqn remain available; torch_lbfgs is deprecated (it never certified convergence and froze the support at its initialisation).

Every result carries the same diagnostics under fixed definitions — grad_map is one scale-invariant certificate for every optimiser, defect_D/defect_Cg/defect_C are all evaluated on the returned basis, and converged is set only when a certificate is earned (result.certificate is "stationary" or "numerical_floor"). max_iter caps gradient evaluations. The default tol is 1e-6 (float64) / 1e-4 (float32), which returns the same support as 1e-9 for 30–50% less work. Stall classification reliably terminates at the float floor across both float32 and float64 without spinning in noise.

nsa_flow_signed writes each component as a contrast of two non-negative parts, which restores a signed basis's representational capacity: at w = 0 it reproduces signed PCA's reconstruction to the digit. The relaxation alone leaves each part concentrated but not sparse, so pass consolidate=True to round to exactly disjoint supports and re-solve with the support fixed. On ADNI cortical thickness at w = 0.5 that takes the largest part from 39 of 66 features to 18, leaves V⁺ and V⁻ each about 10% dense with no component lost, and improves held-out prediction — the small tail was noise.

Leave w at its default. Across nine ADNI cognitive outcomes the lifting is the best of the four variants at w = 0.5 (mean ΔR² over PCA +0.011 linear, +0.044 forest) and worse than PCA on 0 of 9 at w = 0.75. Since a linear model sees only the span, that collapse is a loss of span quality: pushing w up rounds the contrasts toward a partition that no longer spans what the data needs. w = 0.5 is also the best setting for sparsity, so there is nothing to trade off.

For k > p, orthonormal columns are impossible and inf D = 1/p - 1/k > 0; this is reported rather than hidden behind a silently row-orthonormal answer.

Install, test, reproduce

make install       # editable install with experiment + test extras
make test          # 403 assertions
make theory        # just the property battery (79 assertions)
make experiments   # regenerate paper/results/ and paper/figs/
make paper         # build paper/nsaflow.pdf

tests/test_theory.py states every proposition in the paper executably — the bounds, the zero set, the invariances, the spectral identity, the decomposition, the collapse floor, the gradient identities, the disjoint-support equivalence, the overlap bound, and term calibration for every option pair. If a claim in the paper is weakened, one of those fails.

Migrating from 1.x

The 1.x API is gone; nsa_flow(target, w=...) replaces nsa_flow_orth(Y0, X0=..., ...) and the retraction, optimiser and learning-rate-strategy modules are removed. The appendix of the paper lists the substantive changes and why each was made; the 1.x code and tests are preserved under attic/ for reference.

Citation

Avants, B. NSA-Flow: Non-negative Stiefel-Approximating Flow --- a
calibration-free relaxation between reconstruction and feature clustering.

MIT licensed.