Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation

Killian Bakong, Laurent Massoulié, Edouard Oyallon, Kevin Scaman

TL;DR

This work addresses the computational and memory bottlenecks of backpropagation by replacing exact vector-Jacobian products with unbiased randomized estimators. It develops a comprehensive sketching framework for backpropagation, proving optimality results for low-rank and diagonal sketches and analyzing how variance propagates through the DAG during the reverse pass. A suite of practical methods—uniform masks, rank-constrained sketches, and data-dependent diagonal sketches—are proposed and evaluated on MLPs, BagNet, and ViT, showing meaningful cost reductions with limited impact on accuracy. The findings offer a scalable path toward more bandwidth- and compute-efficient backpropagation in pipeline-parallel and large-model settings, with future work targeting coordinated variance control and adaptive hyperparameters.

Abstract

In this work we introduce methods to reduce the computational and memory costs of training deep neural networks. Our approach consists in replacing exact vector-jacobian products by randomized, unbiased approximations thereof during backpropagation. We provide a theoretical analysis of the trade-off between the number of epochs needed to achieve a target precision and the cost reduction for each epoch. We then identify specific unbiased estimates of vector-jacobian products for which we establish desirable optimality properties of minimal variance under sparsity constraints. Finally we provide in-depth experiments on multi-layer perceptrons, BagNets and Visual Transfomers architectures. These validate our theoretical results, and confirm the potential of our proposed unbiased randomized backpropagation approach for reducing the cost of deep learning.

Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation

TL;DR

This work addresses the computational and memory bottlenecks of backpropagation by replacing exact vector-Jacobian products with unbiased randomized estimators. It develops a comprehensive sketching framework for backpropagation, proving optimality results for low-rank and diagonal sketches and analyzing how variance propagates through the DAG during the reverse pass. A suite of practical methods—uniform masks, rank-constrained sketches, and data-dependent diagonal sketches—are proposed and evaluated on MLPs, BagNet, and ViT, showing meaningful cost reductions with limited impact on accuracy. The findings offer a scalable path toward more bandwidth- and compute-efficient backpropagation in pipeline-parallel and large-model settings, with future work targeting coordinated variance control and adaptive hyperparameters.

Abstract

In this work we introduce methods to reduce the computational and memory costs of training deep neural networks. Our approach consists in replacing exact vector-jacobian products by randomized, unbiased approximations thereof during backpropagation. We provide a theoretical analysis of the trade-off between the number of epochs needed to achieve a target precision and the cost reduction for each epoch. We then identify specific unbiased estimates of vector-jacobian products for which we establish desirable optimality properties of minimal variance under sparsity constraints. Finally we provide in-depth experiments on multi-layer perceptrons, BagNets and Visual Transfomers architectures. These validate our theoretical results, and confirm the potential of our proposed unbiased randomized backpropagation approach for reducing the cost of deep learning.
Paper Structure (52 sections, 10 theorems, 86 equations, 4 figures, 6 algorithms)

This paper contains 52 sections, 10 theorems, 86 equations, 4 figures, 6 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Sketching Reverse-Mode Automatic Differentiation
  4. Variance in Stochastic Gradient Descent
  5. Stochastic Gradient Surrogate
  6. Gradient Computation on a DAG
  7. Standard Gradient Computations through a Computational Graph.
  8. Gradient Estimation on a DAG
  9. Randomized Vector-Jacobian Products in Linear Settings
  10. Randomized Sketching Framework.
  11. Optimal Unbiased Low-Rank Matrix Approximation.
  12. Application to VJPs.
  13. Diagonal Sketches.
  14. Jacobian Approximations
  15. First Strategy: Applying Uniform Masks
  16. ...and 37 more sections

Key Result

Proposition 2.2

Assume the seed at the output node is exact (i.e. $\widehat{g}_{\text{out}} = g_{\text{out}}$) and that as:local_vjp_unbiased holds. Then, for every node $i$ of the DAG,

Figures (4)

  • Figure 1: Comparison of sampling strategies and scoring methods.
  • Figure 2: Comparison of types of weighted methods.
  • Figure 3: Sketching on larger architectures.
  • Figure 4: Impact of VJP Approximation Location in MLPs.

Theorems & Definitions (16)

  • Proposition 2.2
  • Lemma 3.1: Optimal unbiased random sketch under rank constraint
  • Lemma 3.2: Distortion in Linear Nodes
  • proof
  • Proposition 3.3: Minimal Distortion rank $r$ Unbiased Sketch
  • proof : Sketch of proof
  • Lemma 3.4: Diagonal mask with expected size at most $r$
  • Proposition 1.1: Restatement of \ref{['prop:variance_backprop']}
  • proof
  • Lemma 1.2: Restatement of \ref{['thm:svd_sketch']}
  • ...and 6 more