Linear Gradient Prediction with Control Variates

TL;DR

This paper introduces a novel training method for neural networks that uses approximate gradient predictions with control variates, reducing training costs while maintaining unbiasedness, inspired by Neural Tangent Kernel theory.

Contribution

It presents a new gradient prediction technique based on control variates and NTK theory, improving training efficiency without sacrificing accuracy.

Findings

Effective on vision transformer classification tasks

Reduces training cost while maintaining unbiased gradient estimates

Demonstrates practical benefits of the proposed method

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.