Randomized matrix sketching for neural network training and gradient monitoring

View at publisher
Authors: Harbir Antil, Deepanshu Verma
Cited by: 1
References: 15

Abstract

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

Fields of Study & Subjects

Stochastic Gradient Optimization Techniques Neural Networks and Reservoir Computing Model Reduction and Neural Networks

Publication details

Journal / source
Mathematical Models and Methods in Applied Sciences
Publisher
World Scientific Pub Co Pte Ltd
Volume
36
Issue
08
Pages
1623–1648
Publication date
1 April 2026
Language
EN