BSFA: Leveraging the Subspace Dichotomy to Accelerate Neural Network Training

TL;DR

This work identifies a fundamental subspace dichotomy in neural network optimization, where updates in the dominant Hessian directions stabilize training while bulk-subspace updates primarily drive convergence. It introduces BSFA, a plug-in accelerator that differentially scales updates via a projector $\mathcal{P}_{\alpha,\gamma}(\theta)$, with estimators (PPE/LPE) to efficiently approximate dominant directions and a block-wise variant (BPPE) for scalability. The approach yields practical speedups (≈2× on large Transformers and ViT, up to 4× in controlled settings) and maintains or improves accuracy, aided by memory-conscious design and optional 4-bit quantization. These results suggest that explicit, subspace-aware update modulation can substantially reduce training time for modern, large-scale models, with clear directions for reducing memory overhead in future work.

Abstract

Recent studies \citep{gur2018gradient,song2024does, wen2024understanding} highlight a fundamental dichotomy in deep learning optimization: Although parameter updates along the top eigendirections of the loss Hessian (Dom-space) capture most of the update magnitude, they often contribute minimally to loss reduction. In contrast, updates in the orthogonal component (Bulk-space) have smaller magnitudes but drive most learning progress. In this work, we further advance the understanding of this phenomenon and introduce the \textbf{Bulk-Space-Filtration-Accelerator (BSFA)}, a novel plug-and-play framework. BSFA accelerates training by differentially scaling update components projected onto these distinct subspaces, simultaneously enhancing stability by moderating updates in the dominant subspace and boosting convergence speed by amplifying those in the bulk-space. To ensure BSFA is both practical and scalable for contemporary large models, we introduce two key innovations: an efficient estimator using Principal Component Analysis (PCA) on historical updates for fast subspace estimation, and a block-wise strategy that applies this estimation on a per-parameter-block basis. These designs make BSFA computationally tractable and highly effective. We demonstrate BSFA's acceleration across various tasks, notably achieving approximately 2$\times$ speedup when pre-training LLaMA-72M on WikiText-103 and LLaMA-134M on OpenWebText compared to vanilla AdamW.

Figures (9)

• Figure 1: SGD in dominant and bulk subspace ($k=2$).
• Figure 2: Subspace‐specific Update Scaling. We follow the setting of Figure \ref{['fig:minhak']}(SGD), and after 4000 steps, we introduce two scaling factors, $\alpha$ and $\gamma$, to independently modulate the update magnitudes in the dominant and bulk subspaces. Training is terminated once the model's training accuracy reaches 0.99.
• Figure 3: SGDM with BSFA ($\alpha=0.5, \gamma=4$) achieves a $4\times$ acceleration compared to tuned SGDM. In both experiments, we employ a cosine learning‐rate schedule with different total epoch counts. In each case, BSFA consistently achieves lower validation loss at identical epoch checkpoints and matches the baseline’s 200‐epoch test accuracy in just 50 epochs. We provide training results for three random seeds.
• Figure 4: (a) The Variance of updates in Dom-space is much higher than Bulk-space. (b) Alignment of Dom-space estimated by PPE and LPE is close to 1. (c) PPE and LPE exhibit comparable efficacy.
• Figure 5: AdamMini--BSFA outperform AdamW baseline in ViT and LLaMA training.

Theorems & Definitions (3)

• Proposition 1: Top eigenspace recovery via PCA
• Proposition 2: Gradient–PCA recovers the top eigenspace
• proof