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 $Y \geq 0$0 is generated by:

1. Euclidean Gradient Step:

$Y \geq 0$1

2. Polar Retraction (Stiefel Projection):

$Y \geq 0$2

3. Soft Interpolation:

$Y \geq 0$3

4. Proximal Non-negativity:

$Y \geq 0$4

In the continuous-time limit, the flow can be formulated as: $Y \geq 0$5 where $Y \geq 0$6 denotes the normal cone to the non-negativity constraint.

3. Algorithmic Implementation

A prototypical NSA-Flow implementation comprises the following steps:

$$L_{\text{orth}}(Y) = \frac{1}{2}\|Y^T Y - I_k\|_F^2$$1

Hyperparameters:

• $Y \geq 0$7 (orthogonality strength): 0.5 (balanced), 0.75–0.95 for increased sparsity
• $Y \geq 0$8: 0.01–0.1, tune adaptively with scheduler or line search
• optimizer: "asgd" or "lars" recommended for speed–stability
• max_iter: 500–1000; tol: $Y \geq 0$9–$X_0 \in \mathbb{R}^{p \times k}$0

The computational complexity per iteration is $X_0 \in \mathbb{R}^{p \times k}$1; thus, NSA-Flow scales well for $X_0 \in \mathbb{R}^{p \times k}$2.

4. Geometric and Structural Insights

The NSA-Flow dynamics result in representation matrices whose columns have disjoint support when $X_0 \in \mathbb{R}^{p \times k}$3 and the orthogonality weight $X_0 \in \mathbb{R}^{p \times k}$4 is high. This is a consequence of the mutual orthogonality (decorrelation) pressure within the non-negative orthant, generating structured sparsity without explicit $X_0 \in \mathbb{R}^{p \times k}$5 regularization. As $X_0 \in \mathbb{R}^{p \times k}$6, the method approaches strict Stiefel manifold projections, maximally decorrelating columns and increasing sparsity. Conversely, $X_0 \in \mathbb{R}^{p \times k}$7 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 \in \mathbb{R}^{p \times k}$8 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 $X_0 \in \mathbb{R}^{p \times k}$9 step in SPCA by launching an NSA-Flow cycle on the gradient-updated input.
• Hyperparameter tuning: The trade-off parameter $$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$0 is selected via cross-validation (using downstream classification, regression, or explained variance). The proportion of zeros and the fidelity–orthogonality trade-off are diagnostic: $$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$1 for moderate orthogonality, $$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$2 for balanced sparsity–fidelity, and $$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$3 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 ($$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$4, $$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$5, $$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$6):

Method	Explained Variance	Sparsity	Orth Defect	CV Accuracy
PCA	0.290	0.00	0.00	0.819
SPCA ($$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$7)	0.158	0.80	0.006	0.864
SPCA (NSA-Flow)	0.172	0.704	$$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$80	0.883

ADNI Cortical Thickness ($$E(Y) = (1-w)\,L_{\text{fid}}(Y, X_0) + w\,L_{\text{orth}}(Y)$$9, $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$0 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 ($$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$1, $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$2)
• Regression on nine cognitive outcomes: NSA-Flow yielded lower $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$3 (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 $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$4, or random nonnegative start, is recommended to avoid poor local minima.
• Scaling: Pre-normalize $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$5 (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 $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$6 or Armijo backtracking. Reduce by half on plateau (patience $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$7 10).
• Monitoring and diagnostics: Plot fidelity and orthogonality defect over iterations; assess sparsity versus $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$8.
• Computational scaling: $$L_{\text{fid}}(Y, X_0) = \frac{1}{2}\|Y - X_0\|_F^2$$9 per iteration; highly scalable for $$L_{\text{orth}}(Y) = \frac{1}{2}\|Y^T Y - I_k\|_F^2$$0.

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