Randomized Matrix Sketching for Neural Network Training and Gradient Monitoring
This work addresses memory constraints for neural network practitioners, offering an incremental improvement through a novel application of existing sketching techniques.
The paper tackles the memory scalability challenge in neural network training by adapting control-theoretic matrix sketching to layer activations, enabling memory-efficient gradient reconstruction and demonstrating a controllable accuracy-memory tradeoff on datasets like MNIST and CIFAR-10.
Neural network training relies on gradient computation through backpropagation, yet memory requirements for storing layer activations present significant scalability challenges. We present the first adaptation of control-theoretic matrix sketching to neural network layer activations, enabling memory-efficient gradient reconstruction in backpropagation. This work builds on recent matrix sketching frameworks for dynamic optimization problems, where similar state trajectory storage challenges motivate sketching techniques. Our approach sketches layer activations using three complementary sketch matrices maintained through exponential moving averages (EMA) with adaptive rank adjustment, automatically balancing memory efficiency against approximation quality. Empirical evaluation on MNIST, CIFAR-10, and physics-informed neural networks demonstrates a controllable accuracy-memory tradeoff. We demonstrate a gradient monitoring application on MNIST showing how sketched activations enable real-time gradient norm tracking with minimal memory overhead. These results establish that sketched activation storage provides a viable path toward memory-efficient neural network training and analysis.