NSA-Flow: Nonnegative Stiefel Flow

• NSA-Flow is a matrix optimization framework that produces interpretable, nonnegative embeddings by balancing reconstruction fidelity with a soft orthogonality penalty.
• It integrates sparse matrix factorization, soft orthogonalization, and constrained manifold learning to enforce sparsity and mutual column decorrelation.
• The tunable parameter allows smooth interpolation between dense PCA approximations and sparse, structured representations, making it a versatile drop-in for dimensionality reduction pipelines.

Non-negative Stiefel Approximating Flow (NSA-Flow) is a matrix optimization framework designed to produce interpretable low-dimensional embeddings from high-dimensional data, particularly where interpretability, sparsity, and mutual column orthogonality are simultaneously desired. NSA-Flow operates by smoothly interpolating between data fidelity and column-wise decorrelation under a non-negativity constraint, leveraging a single tunable parameter to traverse this trade-off. The approach integrates concepts of sparse matrix factorization, soft orthogonalization, and constrained manifold learning, and is applicable as a drop-in module for pipelines such as PCA, Sparse PCA, and other structure-seeking dimensionality reduction methods.

1. Mathematical Formulation

NSA-Flow seeks a nonnegative matrix $Y \in \mathbb{R}^{p \times k}$ ($Y \geq 0$) that closely approximates a fixed target $X_0 \in \mathbb{R}^{p \times k}$, balancing reconstruction fidelity against the soft constraint of column-wise orthogonality. The objective function is

$$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$

where:

• $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$ is the reconstruction error.
• $$L_{\text{orth}}(Y) = \frac{1}{2}\|Y^T Y - I_k\|_F^2$$ is the soft orthogonality penalty.
• $w \in [0, 1]$ is a tunable parameter controlling the balance between fidelity and decorrelation.

An alternative scale-invariant orthogonality is given by: $$L_{\text{orth,inv}}(Y) = \frac{\|Y^T Y - \mathrm{diag}(\mathrm{diag}(Y^T Y))\|_F^2}{\|Y\|_F^4}$$ The full constrained optimization problem is thus: $$Y^* = \arg\min_{Y \ge 0} (1-w)\cdot \frac{1}{2}\|Y - X_0\|_F^2 + w\cdot \frac{1}{2}\|Y^T Y - I_k\|_F^2$$

The Euclidean gradient is: $\nabla_Y E(Y) = (1-w)(Y - X_0) + wY(Y^T Y - I_k)$ Non-negativity is enforced by proximal projection.

2. Flow Dynamics and Iterative Updates

NSA-Flow performs updates that blend standard gradient descent with retraction onto (or near) the Stiefel manifold, followed by interpolation and proximal projection. The update at iteration $t$ is generated by:

1. Euclidean Gradient Step:

$$\widetilde{Y} = Y^{(t)} - \eta [(1-w)(Y^{(t)}-X_0) + wY^{(t)}(Y^{(t)T}Y^{(t)}-I)]$$

2. Polar Retraction (Stiefel Projection):

$$Q = \widetilde{Y} (\widetilde{Y}^T \widetilde{Y})^{-1/2}$$

3. Soft Interpolation:

$Y_{\mathrm{int}} = (1-w)\widetilde{Y} + w Q$

4. Proximal Non-negativity:

$Y^{(t+1)} = \max(0, Y_{\mathrm{int}})$

In the continuous-time limit, the flow can be formulated as: $$\dot{Y}(t) = -(1-w)(Y - X_0) - wY(Y^T Y - I_k) - \rho \,\partial\iota_{\{Y\geq 0\}}(Y)$$ where $\partial\iota_{\{Y\geq 0\}}$ denotes the normal cone to the non-negativity constraint.

3. Algorithmic Implementation

A prototypical NSA-Flow implementation comprises the following steps:

Input: X0 ∈ ℝ^{p×k} # target matrix Y0 ∈ ℝ^{p×k}, Y0 ≥ 0 # initial iterate w ∈ [0,1] # orthogonality weight η0 > 0 # initial step size max_iter # e.g., 1000 tol # convergence tolerance optimizer_type # 'ASGD', 'LARS', or 'Adam' Initialize: Y ← Y0 η ← η0 best_E ← +∞; best_Y ← Y for t = 0 to max_iter-1: # Step 1: Gradient G ← (1-w)*(Y - X0) + w*Y*(Y.T @ Y - I_k) # Step 2: Optimizer update Y_gd ← optimizer_step(Y, G, η) # Step 3: Polar retraction V, D = eig(Y_gd.T @ Y_gd) T_inv_sqrt = V @ diag(D^{-1/2}) @ V.T Q = Y_gd @ T_inv_sqrt # Step 4: Soft interpolation Y_int = (1-w)*Y_gd + w*Q # Step 5: Proximal non-negativity Y_new = np.maximum(0, Y_int) # Step 6: Objective calculation E_new = (1-w)*0.5*norm(Y_new-X0, 'fro')**2 + w*0.5*norm(Y_new.T @ Y_new - I_k, 'fro')**2 if E_new < best_E: best_E = E_new; best_Y = Y_new if (abs(E_new - E_old)/E_old < tol) or (norm(Y_new - Y, 'fro') < tol): break # (Optional) Line search or learning rate scheduling Y = Y_new; E_old = E_new return best_Y, best_E

Hyperparameters:

• $w$ (orthogonality strength): 0.5 (balanced), 0.75–0.95 for increased sparsity
• $\eta_0$: 0.01–0.1, tune adaptively with scheduler or line search
• optimizer: "asgd" or "lars" recommended for speed–stability
• max_iter: 500–1000; tol: $1e{-6}$–$1e{-8}$

The computational complexity per iteration is $O(pk^2)$; thus, NSA-Flow scales well for $p \gg k$.

4. Geometric and Structural Insights

The NSA-Flow dynamics result in representation matrices whose columns have disjoint support when $Y \geq 0$ and the orthogonality weight $w$ is high. This is a consequence of the mutual orthogonality (decorrelation) pressure within the non-negative orthant, generating structured sparsity without explicit $\ell_1$ regularization. As $w \to 1$, the method approaches strict Stiefel manifold projections, maximally decorrelating columns and increasing sparsity. Conversely, $w \to 0$ recovers dense, purely Euclidean approximations.

The mechanism thus enables smooth interpolation between purely data-driven dense representations and maximally interpretable, orthogonal, sparse factor matrices. This approach differs fundamentally from classical regularization schemes, offering direct geometric manipulation of latent structure.

5. Applications and Integration

NSA-Flow can be integrated into established dimensionality reduction and representation learning workflows:

• PCA refinement: Setting $X_0 = Y_0$ as classical PCA loadings and running NSA-Flow yields interpretable, sparse, and nonnegative loadings with preserved explained variance.
• Sparse PCA (SPCA) inner loop: NSA-Flow can act as a drop-in replacement for the soft-threshold $\ell_1$ step in SPCA by launching an NSA-Flow cycle on the gradient-updated input.
• Hyperparameter tuning: The trade-off parameter $w$ is selected via cross-validation (using downstream classification, regression, or explained variance). The proportion of zeros and the fidelity–orthogonality trade-off are diagnostic: $w \in [0.05,0.25]$ for moderate orthogonality, $w \in [0.5,0.75]$ for balanced sparsity–fidelity, and $w \in [0.9,0.99]$ for high sparsity.

6. Empirical Performance and Benchmarks

NSA-Flow has been benchmarked on both canonical and real-world high-dimensional datasets:

Golub Leukemia Data ($n=72$, $p \approx 7000$, $k=2$):

Method	Explained Variance	Sparsity	Orth Defect	CV Accuracy
PCA	0.290	0.00	0.00	0.819
SPCA ($\ell_1$)	0.158	0.80	0.006	0.864
SPCA (NSA-Flow)	0.172	0.704	$\approx$0	0.883

ADNI Cortical Thickness ($N \times p=76$, $k=5$ networks):

• NSA-Flow versus PCA, AUC (random-forest subject scores):
  • CN vs MCI: NSA=0.675, PCA=0.595
  • CN vs AD: NSA=0.844, PCA=0.843
  • MCI vs AD: NSA=0.733, PCA=0.715
  • Multiclass: NSA=0.765, PCA=0.719 ($t=16.48$, $p<1e-4$)
• Regression on nine cognitive outcomes: NSA-Flow yielded lower $-\log p$ (better fit) on 5 of 9 measures.

In both settings, NSA-Flow maintained or improved downstream predictive performance and interpretability relative to both classical and sparse PCA.

7. Practical Recommendations

• Initialization: SVD/PCA on $X_0$, or random nonnegative start, is recommended to avoid poor local minima.
• Scaling: Pre-normalize $X_0$ (e.g., unit Frobenius norm) to regularize penalty balance.
• Optimizer selection: ASGD or LARS with moderate momentum; monitor for gradient divergence.
• Step size management: Use initial data-driven $\eta_0$ or Armijo backtracking. Reduce by half on plateau (patience $\approx$ 10).
• Monitoring and diagnostics: Plot fidelity and orthogonality defect over iterations; assess sparsity versus $w$.
• Computational scaling: $O(pk^2)$ per iteration; highly scalable for $p \gg k$.

NSA-Flow enables interpretable, structured representations for exploratory and predictive analytics across domains, notably in genomics and neuroimaging, with minimal modification to existing matrix factorization pipelines (Avants et al., 9 Nov 2025).