OASIS: Online Activation Subspace Learning for Memory-Efficient Training

Abstract: Training LLMs is constrained by memory requirements, with activations accounting for a substantial fraction of the total footprint. Existing approaches reduce memory using low-rank weight parameterizations or low-rank gradient subspaces for optimizer states, while activation memory is addressed through architectural modifications or compression schemes based on periodically updated projections. We propose OASIS, an online activation subspace learning algorithm for memory-efficient training that tracks and continuously updates a low-dimensional activation subspace during training. Intermediate activations are projected onto this evolving subspace, reducing memory without modifying forward-pass computations. The evolving activation subspace induces low-rank gradient representations, enabling both gradients and optimizer states to be maintained directly in this subspace, while a projection-aware optimizer consistently transports optimizer states across subspace updates for stable training. Across various finetuning and pretraining tasks, OASIS achieves up to $2\times$ lower peak memory than full fine-tuning while matching its performance and outperforming prior low-rank methods.

• The paper introduces an online activation subspace learning method that compresses activations during LLM training to significantly reduce memory usage.
• It leverages continuous updates via Oja’s rule to adaptively maintain a low-rank representation, achieving up to 2× memory savings compared to conventional methods.
• Empirical results across various Llama-family models show that OASIS outperforms techniques like LoRA, GaLore, and LDAdam while ensuring robust downstream performance.

OASIS: Online Activation Subspace Learning for Memory-Efficient Training

Motivation and Problem Formulation

LLM training is fundamentally constrained by memory demands, with activations dominating the memory footprint, especially under large batch sizes, long sequences, or increasing model parameterization. Traditional techniques, such as low-rank parameterizations (LoRA, QLoRA, DoRA) or gradient subspace tracking (GaLore, LDAdam), primarily reduce parameter or optimizer state memory but typically leave the substantial activation memory requirements unmitigated. Architecturally invasive or compression-based techniques for activations tend to degrade task performance or require modifications that limit applicability or generalizability.

OASIS Framework

OASIS (“Online Activation Subspace Learning for Memory-Efficient Training”) proposes a unified memory-reduction strategy by maintaining a low-dimensional, online-tracked activation subspace. The methodology leverages two key theoretical underpinnings: (1) the gradients of a linear layer’s parameters reside within the span of the input activations, and (2) principal directions of activations evolve dynamically throughout training due to continuous changes in network parameters.

The framework operates as follows:

• Subspace Projection: At each training iteration, input activations are projected onto a learned, low-dimensional subspace $U_t$, yielding $\widetilde{X}_t = X_t U_t$.
• Low-Rank Gradient and Optimizer States: Gradients and optimizer states are computed and maintained in this subspace, yielding $O(N r)$ and $O(rm)$ activation and optimizer memory costs (vs. $O(Nd)$ and $O(dm)$ full precision).
• Online Subspace Update: Crucially, OASIS updates $U_t$ iteratively using Oja’s rule, a streaming PCA algorithm. Each step increments the basis to track dominant activation variance directions, followed by explicit re-orthonormalization and adaptive learning rate normalization, efficiently capturing activation dynamics.

Figure 1: Memory breakdown during training, showing the predominance of activation memory in the training footprint.

Figure 2: OASIS reduces the memory requirements across activations, gradients, and optimizer states, yielding the lowest overall memory usage relative to baselines.

The optimizer must be projection-aware: as $U_t$ evolves, previous optimizer states are aligned to the new subspace using subspace transition matrices, ensuring consistent and stable optimization (cf. the coordinate transport procedure in LDAdam).

Experimental Evaluation

The authors conduct comprehensive evaluations on both finetuning and pretraining settings using various Llama-family models (7B, 1B, 350M, 130M), spanning tasks such as GSM8K (math word problems), HumanEval (code generation), and C4 (language modeling).

Key empirical findings include:

• Peak memory reduction: OASIS regularly achieves up to $2\times$ lower peak memory usage compared to full fine-tuning and over $30\%$ reduction relative to gradient subspace methods (GaLore, LDAdam), matching or exceeding their downstream accuracy, even at aggressive low ranks (e.g., $\widetilde{X}_t = X_t U_t$0 or $\widetilde{X}_t = X_t U_t$1).
• Performance-Compression Tradeoffs: OASIS consistently outperforms LoRA, GaLore, and LDAdam at corresponding ranks with stronger accuracy and lower memory, both in finetuning and pretraining.

OASIS ablates update frequency, subspace learning rate, and subspace rank, revealing that:

• Continuous online subspace updates outperform periodic PCA/grid-tuned updates in both convergence speed and eventual performance.
• Performance saturates beyond moderate ranks, demonstrating OASIS’s robustness to subspace dimensionality under practical choices.

Figure 3: OASIS consistently outperforms periodic PCA at all tested subspace ranks on GSM8K finetuning.

Analysis of Subspace Dynamics

Subspace drift analyses confirm that the optimal activation subspace is inherently non-stationary during training, both for finetuning and pretraining. Drift metrics (Frobenius distance) indicate that activations continuously evolve, which invalidates static or infrequently updated projections as used by prior methods.

OASIS’s online subspace adaptation closely tracks these dynamics, preserving the representation quality and ensuring optimizer convergence.

Figure 4: The effect of PCA update interval demonstrates that periodic subspace recomputation is inferior to OASIS’s continual adaptation.

Practical and Theoretical Implications

The results demonstrate that memory-efficient training is not strictly bounded by parameter- or optimizer-focused techniques. By moving activation storage into a highly expressive, adaptive low-rank subspace, OASIS fundamentally re-allocates the memory budget, enabling either larger batch sizes, longer input sequences, or deployment on hardware with more restrictive device memory.

On the theoretical front, OASIS substantiates that low-rank approximation of activations is tenable even under the non-stationary conditions of realistic LLM training, provided the basis evolves online in concordance with model updates.

Practically, OASIS is immediately compatible with standard training pipelines, requiring neither architectural modifications nor constraining the update expressivity, positioning it as a general drop-in component for memory scaling.

Future Perspectives

Several directions warrant further exploration:

• Learning Rate Adaptation: The subspace learning rate is a significant hyperparameter; dynamic or data-driven adjustment strategies may yield further robustness.
• Extended Applicability: Extensions to non-linear or structured activation manifolds may further generalize the framework.
• Integration with Additional Compression Techniques: Investigating synergistic effects with quantization, pruning, or on-device continual learning optimizations is promising.

Conclusion

OASIS presents an online activation subspace learning strategy that unifies activation, gradient, and optimizer state compression in a single, theoretically sound and empirically validated algorithm. It outperforms baseline low-rank and activation compression techniques in both memory efficiency and downstream performance across a range of LLM training regimes. Its adaptability to subspace non-stationarity and integration with existing optimizers position it as a strong candidate for scalable training of increasingly large neural models.