STAC-R: Subspace-Tracked Activation Compression with Residuals
Abstract
Training-memory reduction requires identifying which saved tensors drive peak usage and ensuring high-fidelity backward contractions from compressed states. We formalize saved-activation compression as a target-surface prediction problem: a pre-training framework to estimate the replaceable memory surface, computational overhead, and parameter-gradient error. We introduce STAC-R (Subspace Tracked Activation Compression with Residuals), a saved-tensor wrapper storing activations as persistent low-rank coefficients paired with a sparse quantized residual. Weight gradients are contracted directly from these factors, bypassing full decompression. To guide deployment, we propose the Saved-Surface Predictor (), a statistic that upper-bounds peak memory reduction. We also derive a gradient bridge linking reconstruction error to parameter-gradient cosine similarity, determining if a low-precision state preserves update fidelity.Evaluations across GPT-2 (Medium/Large) and Pythia-1.4B pre-training show STAC-R reduces peak memory by 7.2-12.4% while keeping long-horizon validation perplexity within 1.5% of baseline (Pythia-1.4B 10k steps: +1.14%) with mean gradient cosine 0.98. On BERT fine-tuning (SST-2), STAC-R achieves 12.45% memory reduction with a 0.15pp accuracy decrease. Mechanistic benchmarks yield a mean gradient cosine of 0.9813, outperforming iso-byte SVD (0.9236) and untracked-subspace (0.6949) controls. On the empirical rate-distortion frontier, STAC-R is Pareto-dominant over controls across stored-byte budgets. Composed with GaLore, a "Turbo-fast" variant matches full STAC-R's memory footprint within 0.005 GiB while reaching 0.67𝑥 wall time. Finally, our analysis identifies memory-increasing regimes in SwiGLU architectures (Qwen-2.5/Llama-3.2), establishing STAC-R as a compressor with predictable memory and fidelity conditions checkable before training.
Related papers
See more related submissions when the market opens.