Linear Gradient Prediction with Control Variates

2025-11-07

AI Summary

To address the high computational cost of gradient computation in neural network training, this paper proposes an unbiased and efficient approximate gradient estimation method. Our approach introduces three key innovations: (1) a lightweight gradient predictor inspired by the Neural Tangent Kernel theory, enabling rapid gradient approximation in forward mode; (2) control variates to rigorously guarantee estimator unbiasedness and controllable variance; and (3) integration with a linear approximation model, validated end-to-end on Vision Transformers. Experiments on benchmarks including ImageNet demonstrate that our method reduces training FLOPs by up to 38% while maintaining convergence speed and final accuracy comparable to full-gradient training. The proposed framework offers a novel paradigm for efficient large-model training without sacrificing optimization fidelity or generalization performance.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Deep Neural Architectures and Foundation Models
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

We propose a new way of training neural networks, with the goal of reducing training cost. Our method uses approximate predicted gradients instead of the full gradients that require an expensive backward pass. We derive a control-variate-based technique that ensures our updates are unbiased estimates of the true gradient. Moreover, we propose a novel way to derive a predictor for the gradient inspired by the theory of the Neural Tangent Kernel. We empirically show the efficacy of the technique on a vision transformer classification task.

Problem

• Reducing neural network training cost using approximate predicted gradients
• Developing unbiased gradient estimates through control-variate techniques
• Creating gradient predictors inspired by Neural Tangent Kernel theory

Innovation

• Uses approximate predicted gradients for training
• Employs control variates for unbiased gradient estimates
• Derives gradient predictor using Neural Tangent Kernel theory