# Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation

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

<figure><img src="image_1.png" /></figure>

To cite this version:

Killian Bakong, Laurent Massoulié, Edouard Oyallon, Kevin Scaman. Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation. 2026. hal-05512493

HAL Id: hal-05512493

https://hal.science/hal-05512493v1

Preprint submitted on 16 Feb 2026

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
# Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation

Killian Bakong<fnref n="1" /> Laurent Massoulié<fnref n="1" /> Edouard Oyallon² Kevin Scaman<fnref n="1" />

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

## 1. Introduction

Backpropagation is the workhorse of neural network training: it applies the chain rule over a computational graph to compute gradients efficiently (Rumelhart et al., 1986; LeCun et al., 2002; Baydin et al., 2018). Formally, it is reverse-mode automatic differentiation (AD) applied to a directed acyclic graph (DAG) of primitive operations (Griewank & Walther, 2008). For scalar losses, reverse-mode AD composes a sequence of vector-Jacobian products (VJPs) and is operation- and memory-optimal (Linnainmaa, 1976; Baur & Strassen, 1983). However, despite this optimality, backpropagation remains a major speed bottleneck in modern training, which motivates numerous engineering techniques (e.g., activation checkpointing (Herrmann et al., 2019; Griewank & Walther, 2000; Rojas et al., 2020) and 4D parallelism (Li et al., 2021)) to scale

up those methods to more workers or larger models.

This paper studies *randomized approximations of VJPs* inside reverse-mode AD. The motivation of our method is threefold and aligns with the constraints of large-scale training, especially under pipeline parallelism (Huang et al., 2019) and its hybrid variants (Narayanan et al., 2019). **(i)** Communication: in pipeline parallelism, inter-layer activations often dominate cross-device traffic (Narayanan et al., 2019; Huang et al., 2019). Compressing these signals while preserving gradient unbiasedness can substantially reduce bandwidth and latency. **(ii)** Computation and memory: some randomized linear mapping (like random features sampling (Rahimi & Recht, 2007)) can accelerate intermediate computations (e.g., approximate VJPs (Coleman & Moré, 1983)) and lower memory footprint, enabling better pipeline balance and reducing stalls by trading controlled variance for compute. **(iii)** Deep learning specific: our framework subsumes standard gradient-compression techniques based on masking (Sun et al., 2017; Stich et al., 2018) as a special case, but operates at the level of local VJPs within arbitrary architectures and allows us to apply the approximation in a cascade of layers of a deep neural network.

Concretely, at a node $y = f(x)$ in the computational DAG with Jacobian $J_f(x)$, during backpropagation we replace the exact VJP $J_f(x)^\top g(x)$ by $\hat{J}_f(x)^\top g(x)$ where $\hat{J}_f(x)$ is a *randomized* operator such that $\mathbb{E}[\hat{J}_f(x)|g(x)] = J_f(x)$. We analyze how the variance due to such local substitutions propagates along the reverse pass and identify randomized operators which yield favorable accuracy-efficiency trade-offs. Specifically we propose to approximate operators $J_f$ by random rescaled low-rank projections –either along data-dependent directions or along fixed directions– the latter leading to diagonal mask-and-rescale schemes that keep a few coordinates at random. Both preserve unbiasedness and let us target high-impact directions while meeting a budget on sketch size.

**Contributions.** (1) We propose replacing exact VJPs in a DAG by unbiased randomized estimators at selected nodes to reduce the cost of backpropagation at the expense of limited additional (stochastic gradient) variance. (2) We char-

<footnote n="1">¹Inria Paris & DI ENS, ENS, PSL University ²Sorbonne University, CNRS, ISIR, Paris. Correspondence to: Killian Bakong <firstname.lastname@inria.fr>.</footnote>
acterize the additional variance caused by such random-
ized VJPs along the backpropagation and identify regimes
where errors dampen or amplify as a function of local Ja-
cobian spectra. (3) We identify minimal variance unbi-
ased VJPs for (i) general rank constraints and (ii) rank con-
straints for diagonal mask-and-rescale. (4) We compare the
training accuracy versus cost reduction trade-offs obtained
for a selection of candidate unbiased randomized VJPs via
experiments<fnref n="1" /> with BagNets (Brendel & Bethge, 2019), Vi-
sual Transformers (Dosovitskiy et al., 2021) and standard
MLPs architectures. This leads us to identify promising
candidates, in particular the “ℓ₁”-score approach which
combines high cost reduction with marginal impact on test
accuracy for fixed number of training steps.

A discussion of related works is given after the Conclusion.

## 2. Sketching Reverse-Mode Automatic Differentiation

2.1. Variance in Stochastic Gradient Descent

We consider the minimization of an objective function:

$$
\min F(\theta) = \mathbb{E}_{\xi}[f(\theta, \xi)], \quad (1)
$$

where $\xi$ denotes a random data sample, and $\theta \in \mathbb{R}^d$ a set of parameters, optimization being done by Stochastic Gradient Descent (SGD) with minibatches (Bottou et al., 1991). At each iteration $t$, a minibatch of size $B$ is sampled and the update then writes:

$$
\theta_{t+1} = \theta_t - \eta g_t, \quad g_t := \frac{1}{B} \sum_{b \in [B]} \nabla_{\theta_b} f(\theta_t, \xi_b). \quad (2)
$$

Under classical assumptions, e.g. F β-smooth, non-
negative, and variance bounded by E[||g_t - ∇F(θ_t)||^2 | θ_t] ≤
σ^2, canonical results show that SGD converges at a rate
proportional to said variance. In particular, with appropri-
ate step-size η, one obtains a bound on the sizes of gradi-
ents after T iterations

$$
\min_{0 \le t \le T-1} \mathbb{E}[\|\nabla F_t\|^2] \le \frac{2(F_0 - F^*)}{\eta T} + \beta \eta \sigma^2, \quad (3)
$$

where $F^*$ denotes the global minimum value of $F$ on $\mathbb{R}^d$.

Crucially, nothing in the proof depends on the specific
scheme beyond unbiasedness and bounded variance. The
bound is conservative, but it mirrors training behavior.

2.2. Stochastic Gradient Surrogate

Suppose that now, instead of $g_t$, we use a modified gradient
surrogate,

$$
\mathbb{E}[\hat{g}_t | g_t] = g_t, \quad (4)
$$

with $\hat{g}_t$ satisfying unbiasedness and bounded variance, i.e. $\mathbb{E}[\|\hat{g}_t - g_t\|^2|g_t] \le V$. Then $\hat{g}_t$ remains an unbiased estimator of $\nabla F(\theta_t)$, and standard SGD analyses apply unchanged. In particular, the convergence bound in (3) still holds, up to replacing the variance $\sigma^2$ by $\sigma^2 + V$. With this additional term, for a target accuracy $\epsilon > 0$ on the right-hand side of (3) obtained under the optimal choice of step-size $\eta$, the required number of iterations scales as

$$
T = O \left( \frac{(\sigma^2 + V)\beta(F_0 - F^*)}{\epsilon^2} \right). \quad (5)
$$

As a consequence, if we denote $\rho(V)$ the per-iteration computation induced by the injected variance $V$, a net computational gain (i.e. wall-clock time) is achieved if

$$
\rho(V)(\sigma^2 + V) \leq \rho(0)\sigma^2. \quad (6)
$$

This condition reflects the variance-efficiency trade-off. Thus, any additional noise that is conditionally unbiased only affects optimization through its contribution to variance. This motivates our approach: we trade exact gradient computations for computationally cheap approximations, at the cost of introducing additional unbiased noise.

2.3. Gradient Computation on a DAG

We now describe how such additional noise arises when
approximating backpropagation.

**Standard Gradient Computations through a Computational Graph.** Consider a computational DAG with leaf inputs and intermediate variables connected by differentiable operations. For each node $i$, the forward computation for a sample $\xi$ is

$$
x_j(\xi) = f_i(\theta_i, (x_i(\xi))_{i \to j}), \quad (7)
$$

where $i \to j$ indicates that $x_i$ is an input to the operation producing $x_j$, and $\theta_i$ are parameters local to node $i$. Given a minibatch $\{\xi_1, \dots, \xi_B\}$, forward pass is applied independently to all $\xi_b$. The empirical loss (scalar objective to minimize) is:

$$
L = \frac{1}{B} \sum_{b=1}^{B} \ell(x_{\text{out}}^{(b)}), \qquad (8)
$$

where $x_{\text{out}}$ is the final node. Define reverse-mode sensitivities (adjoints) $g_i^{(b)} \triangleq \frac{\partial}{\partial x_i} \ell(x_{\text{out}}^{(b)})$ and $h_i^{(b)} \triangleq \frac{\partial}{\partial \theta_i} \ell(x_{\text{out}}^{(b)})$. For convenience, write the local Jacobians

$$
J_{ij}^{(b)} \triangleq \partial_{x_j} f_i(\theta_i, (x_k^{(b)})_{k \to i})^{\top} \\
\text{and } J_i^{(b)} \triangleq \partial_{\theta_i} f_i(\theta_i, (x_k^{(b)})_{k \to i})^{\top}, \qquad (9)
$$

where $(\cdot)^{\top}$ denotes the adjoint.

<footnote n="1">¹Our implementations and experimental framework will be
made publicly available upon publication.</footnote>
Standard backwards message passing, using chain rule, writes, for a sample b

$$ \left\{ \begin{aligned} g_i^{(b)} &= \sum_{j: i \to j} J_{ij}^{(b)} g_j^{(b)}, \\ h_i^{(b)} &= J_i^{(b)} g_i^{(b)}, \end{aligned} \right. \qquad (10) $$

where the recursion is seeded at the outputs with $g_{out}^{(b)} = \frac{\partial}{\partial x_{out}} \ell(x_{out}^{(b)})$. Standard SGD would then use $h_i := B^{-1} \sum_{b=1}^{B} h_i^{(b)}$ as the stochastic gradient for $\theta_i$. We shall also need the notation $g_i := B^{-1} \sum_{b=1}^{B} g_i^{(b)}$.

## 2.4. Gradient Estimation on a DAG

We now study how layerwise variance propagates through the network. To this end, we introduce unbiased local Jacobian estimators $\hat{J}_{ij}^{(b)}$ and propagate them through the layer cascade via

$$ \left\{ \begin{aligned} \hat{g}_i^{(b)} &= \sum_{j: i \to j} \hat{J}_{ij}^{(b)} \hat{g}_j^{(b)}, \\ \hat{h}_i^{(b)} &= \hat{J}_i^{(b)} \hat{g}_i^{(b)}, \end{aligned} \right. \qquad (11) $$

We denote by $\hat{g}_i = \frac{1}{B} \sum_b \hat{g}_i^{(b)}$, $\hat{h}_i = \frac{1}{B} \sum_b \hat{h}_i^{(b)}$ the batch-averaged approximate gradients. We then require

**Assumption 2.1** (Unbiasedness of local VJPs). For every sample b, node i and edge $i \to j$,

$$ \begin{aligned} \mathbb{E}[\hat{J}_{ij}^{(b)} | (\hat{g}_k^{(b')})_{k:i \to k, b' \in B}, (\hat{J}_{ik}^{(b')})_{i \to k, k \neq j, b' \in B}] &= J_{ij}^{(b)}, \\ \mathbb{E}[\hat{J}_i^{(b)} | \hat{g}_i^{(b)}] &= J_i^{(b)}. \end{aligned} \qquad (12) $$

We then have (see Appendix for a proof):

**Proposition 2.2.** Assume the seed at the output node is exact (i.e. $\hat{g}_{out} = g_{out}$) and that Theorem 2.1 holds. Then, for every node i of the DAG,

$$ (i) \text{ (Unbiasedness.) } \mathbb{E}[\hat{g}_i | g_i] = g_i, \mathbb{E}[\hat{h}_i | h_i] = h_i. $$

(ii) (Variance propagation.)

$$ \mathbb{E}[\|\hat{g}_i - g_i\|^2] = \sum_{j: i \to j} \mathbb{E}\left[\left\|\frac{1}{B}\sum_{b=1}^{B}(\hat{J}_{ij}^{(b)} - J_{ij}^{(b)})\hat{g}_j^{(b)}\right\|^2\right] \\ + \mathbb{E}\left[\left\|\frac{1}{B}\sum_{b=1}^{B}\sum_{j: i \to j} J_{ij}^{(b)}(\hat{g}_j^{(b)} - g_j^{(b)})\right\|^2\right]. \qquad (13) $$

The variance at a node i is thus composed of two terms: The first is the variance induced locally from the chosen

local VJP approximation, while the second corresponds to the variance stemming from approximations done at the preceding nodes in backpropagation. The first term in this decomposition motivates criteria for the choice of VJP approximation discussed in the Section 3.

The second term in this decomposition informs whether added variance dampens or increases during backpropagation. Indeed, this second term can be upper-bounded, using Cauchy-Schwarz inequality, by

$$ \frac{|\{j: i \to j\}|}{B} \sum_{b=1}^{B} \sum_{j: i \to j} \mathbb{E}\|J_{ij}^{(b)}\|^2 \|\hat{g}_j^{(b)} - g_j^{(b)}\|^2, $$

where $\|J_{ij}^{(b)}\|$ denotes operator norm. If these operator norms are sufficiently small, the approximation errors will thus dampen along backpropagation.

## 3. Randomized Vector-Jacobian Products in Linear Settings

We now study the construction of estimators such that

$$ \hat{J}_{ij}^{(b)} \hat{g}_j^{(b)} \approx J_{ij}^{(b)} g_j^{(b)}. $$

In the context of VJPs, a natural approach is to introduce a (random) sketching matrix, as defined below.

**Randomized Sketching Framework.** In the proposed setting, for and edge $i \to j$, with $J_{ij} \in \mathbb{R}^{m \times n}$ and $\hat{g}_j \in \mathbb{R}^n$, we consider a random sketching matrix $R \in \mathbb{R}^{n \times n}$ such that $\mathbb{E}[R | \mathcal{F}] = I_n$ where $\mathcal{F}$ captures information relative to all other VJP approximations and batch sampling. Then

$$ \forall b \in 1, \dots, B, \hat{J}_{ij}^{(b)} = J_{ij}^{(b)} R, \qquad (14) $$

so the estimator is locally unbiased, as in Theorem 2.1. As motivated in Section 2.2, we aim at lowering additional variance, i.e. we aim at finding the sketch that minimizes the distortion, measured by $L^2$ cost<fnref n="2" />

$$ \mathcal{L}(R) \triangleq \frac{1}{B} \sum_{b=1}^{B} \mathbb{E}[\| J_{ij}^{(b)} \hat{g}_j^{(b)} - J_{ij}^{(b)} R \hat{g}_j^{(b)} \|^2], \qquad (15) \\ = \frac{1}{B} \sum_{b=1}^{B} \mathbb{E} \mathrm{Tr} \left[ J_{ij}^{(b)\top} J_{ij}^{(b)} (I - R) \hat{g}_j^{(b)} \hat{g}_j^{(b)\top} (I - R)^{\top} \right]. $$

This yields a quadratic objective in the compressed surrogate $R$, which preserves unbiasedness of the estimator for $J$ without adding variance. In the batch-wise regime, $R$ is shared across all $b \in \mathcal{B}$; an exact solution follows from the sketching lemma below.

<footnote n="2">²Strictly speaking we consider conditional expectations with respect to $\mathcal{F}$, the $\sigma$-field capturing information relative to all other approximations and batch sampling. We do not make this explicit to lighten notation.</footnote>
**Optimal Unbiased Low-Rank Matrix Approximation.**
We first give a general result, which we will subsequently apply to solve the previous problem in specific settings:

**Lemma 3.1** (Optimal unbiased random sketch under rank constraint). Let $M$ be a fixed matrix in $\mathbb{R}^{m \times n}$. Let $q = \min(m, n)$ and denote the SVD of $M$ by $M = \sum_{i=1}^q \sigma_i u_i v_i^\top$, where $\sigma_i, u_i, v_i$ are respectively its singular values, left and right singular vectors. Then among matrices $S$ of rank at most $r$ where $r < q$, such that $\mathbb{E}[S] = M$, the choice that minimizes the error in squared Frobenius norm $\mathbb{E}\|M - S\|_F^2$ is obtained by letting

$$S = \sum_{i=1}^{q} \sigma_i \frac{Z_i}{p_i} u_i v_i^\top$$

where the $Z_i$ are correlated Bernoulli random variables with respective parameters $p_i$. The $Z_i$ are such that $\sum_{i=1}^q Z_i \equiv r$, and the $\{p_i\}_{i \in [q]}$ are minimizers of $\sum_{i=1}^q \sigma_i^2/p_i$ among weights $p_i \in [0, 1]$ that sum to $r$.

The proof is detailed in the Appendix; it proceeds by first establishing a lower bound on the expected squared Frobenius norm of any matrix $S$ such that $\mathbb{E}[S] = M$, using Jensen's inequality together with convexity properties of functions of singular values, combined with the so-called weak majorization theory. It then shows that this lower bound is reached by the specific choice described above. The procedure to sample correlated Bernoulli random variables with fixed sum is also detailed.

We recently found out that Barnes et al. (2025) have independently shown Theorem 3.1, establishing the upper-bound and noticing that the lower-bound followed from the supplementary material in Benzing et al. (2019). Our proof is however different and may thus be of independent interest.

**Application to VJPs.** The above Lemma has the following direct application:

**Lemma 3.2** (Distortion in Linear Nodes). Consider $i$ a Linear node,

$$x_j(\xi) = W_{ij}x_i(\xi) + \mathbf{b}, \quad W_{ij} \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}, \mathbf{b} \in \mathbb{R}^{d_{\text{out}}},$$

and $g_i^{(b)}$, $b \in [1, B]$ the gradients of a considered batch $B$. Then the Jacobian w.r.t. activations is the same for element of the batch, and, denoting

$$G := [g^{(1)} \dots g^{(B)}] \in \mathbb{R}^{d_{\text{out}} \times B} \quad (16)$$

the matrix whose columns are the backpropagated gradients, the distortion defined in (15) becomes

$$\mathbb{E}[\text{Tr}(J^\top J(I-R)\Gamma_B(I-R)^\top)], \quad (17)$$

where $\Gamma_B := \frac{1}{B}GG^\top$ denotes the empirical second moment matrix of the batch, and $J = W_{ij}$.

*Proof.* Expanding the expression of $\mathcal{L}$ and using the circularity of the trace, we obtain for a single batch $g_i^{(b)}$

$$\mathbb{E}\|J(I-R)g_i^{(b)}\|^2 = \mathbb{E}\text{Tr}(J^\top J(I-R)g_i^{(b)}g_i^{(b)\top}(I-R)^\top).$$

Now, using the linearity of trace and expectancy, averaging over the batch yields

$$\mathcal{L}(R) = \mathbb{E}[\text{Tr}(J^\top J(I-R)(\frac{1}{B}\sum_{b=1}^B g_i^{(b)}g_i^{(b)\top})(I-R)^\top)],$$

hence the result.

□

This, combined with the previous Lemma, yields

**Proposition 3.3** (Minimal Distortion rank $r$ Unbiased Sketch). Let $J \in \mathbb{R}^{m \times n}$ and $\Gamma_B$ defined as above. Define the symmetric matrix

$$\Gamma_B^{1/2} J^\top J \Gamma_B^{1/2} = U \Sigma U^\top, \quad (18)$$

where $U$ is orthogonal and $\Sigma = \text{diag}(\sigma_1^2, \dots, \sigma_n^2)$ with $\sigma_1 \ge \dots \ge \sigma_n \ge 0$. Among all random matrices $R$ with $\mathbb{E}[R] = I_n$ and whose rank is bounded by $r$, a minimizer of the $\mathbb{L}^2$ cost Equation (15) is attained by sketches that are diagonal in the eigenbasis of $\Gamma_B^{1/2} J^\top J \Gamma_B^{1/2}$:

$$\begin{aligned} R^* &= \Gamma_B^{1/2} U B U^\top \Gamma_B^{-1/2} \\ B &= \text{diag}(\frac{z_1}{p_1^*}, \dots, \frac{z_n}{p_n^*}), \quad \text{with } z_i \sim \mathcal{B}(p_i^*) \end{aligned} \quad (19)$$

with $\mathcal{B}(p)$ being the Bernoulli distribution of parameter $p$, and with selection probabilities $\{p_i^*\}_{i=1}^n$ minimizing $\sum_{i=1}^n \sigma_i^2/p_i$ under the constraint $\sum_i p_i = r$.

*Sketch of proof.* Expanding the quadratic form obtained in Theorem 3.2 and using $\mathbb{E}[R] = I_n$ gives

$$\mathcal{L}(R) = -\text{Tr}(J\Gamma_B J^\top) + \mathbb{E}[\text{Tr}(J R \Gamma_B R^\top J^\top)], \quad (20)$$

so minimizing it reduces to minimizing the second term. Writing $S = J R \Gamma_B^{1/2}$, we are thus looking for a matrix $S$ such that $\mathbb{E}[S] = M = J \Gamma_B^{1/2}$, of rank at most $r$, and with minimal expected squared Frobenius norm. Lemma 3.1 gives us the optimal choice for $S$, from which we can deduce the optimal choice for $R$. Details are given in the Appendix. □

**Diagonal Sketches.** In many neural network layers, restricting sketches to diagonal operators leads to significant implementation and memory advantages hence the following analysis of diagonal sketches (e.g. fully masking coordinates) under the same budget constraints.
**Lemma 3.4** (Diagonal mask with expected size at most r). We now restrict R to be diagonal, R = diag(r₁, ..., rₙ) with

$$r_i = z_i/p_i, \quad z_i \sim \mathcal{B}(p_i) \quad \text{independent, } p_i \in [0, 1],$$

so that $\mathbb{E}[R] = I_n$. For such diagonal mask, under the expected rank constraint $\sum_{i=1}^n p_i \le r$, the minimal L² distortion Equation (15) is obtained by choosing probabilities $p_i$ that minimize $\inf_p \sum_{i=1}^n \frac{a_i}{p_i}$ under the constraint $\sum_i p_i \le r$, where $a_i := (\Gamma_B)_{ii} (J^\top J)_{ii}$.

The details are provided in the Appendix. Note that here we have considered independent Bernoulli random variables, and have thus imposed a constraint on the rank in expectation rather than almost surely. This modification is required for the proof to go through.

# 4. Jacobian Approximations

In this section, we look over different practical Jacobian Approximation methods that elaborate on those in the previous section.

## 4.1. First Strategy: Applying Uniform Masks

We first discuss simple uniform masks, common in the literature, which serve as baselines for approximating vector-Jacobian products. We denote them as M (as opposed to geometrically informed sketches R), and their goal is still to approximate VJPs, i.e. $JMg \approx Jg$.

**Per-Element Masks.** Per-element masking is common in practical deep learning settings (Srivastava et al., 2014; Lee et al., 2019), and is the first considered approximation. Let $p \in (0, 1]$ a parameter so that $M_{kl} = B_{kl}/p$, $B_{kl} \sim \mathcal{B}(p)$, and $\mathbb{E}[M_{kl}] = 1$. Element-wise application $\hat{J} = M \odot J$ gives an unbiased estimator of the Jacobian. It is trivial to verify that $\mathbb{E}[\hat{J}] = J$, hence unbiasedness. For such masks, computational load reduction may stem from the element-wise sparsity of the approximating matrix rather than from a reduced rank property.

**Per-Column Masks.** A second simple strategy is to apply independent Bernoulli gating per column (or channel) of J. For a certain $p \in (0, 1]$, draw $z_i \stackrel{\text{i.i.d.}}{\sim} \mathcal{B}(p)$ and define an unbiased and expected rank r mask as

$$M = \operatorname{diag}(z_1/p, \dots, z_n/p). \qquad (21)$$

**Per-Sample Masks.** The last investigated masking method, inspired by Woo et al. (2024), consists in applying a Bernoulli gate to the whole product at sample level. For a fixed $p \in (0, 1]$, draw a scalar gate $z \sim \mathcal{B}(p)$ and set

$$M = (z/p)I_n. \qquad (22)$$

This M also satisfies unbiasedness and expected low-rank.

## 4.2. Second Strategy: Data Dependent Sketches

We now turn to a class of *Data Dependent* sketching operators that allocate the sketching budget according to data-dependent importance weights. Recall from Section 3 that under unbiasedness and rank constraints, minimizing the distortion of a VJP approximation reduces to a convex optimization problem over the sampling probabilities:

$$\min \sum_{i=1}^{n} \frac{w_i}{p_i}, \quad \text{s.t.} \quad \sum_{i=1}^{n} p_i \le r, \ p_i \in (0, 1]. \qquad (23)$$

Here $p_i$ denote the probabilities defining the sketch, while the weights $w_i$ quantify the relative importance of each direction as induced by the considered sketching constraints. We hereafter describe some possible $w_i$ choices leading to sketching operators.

**Rank-Constraint Sketches (RCS).** We first consider the construction in Theorem 3.3. With solely unbiasedness and rank restrictions imposed to R, the importance weights $w_i$ correspond to the squared singular values $\sigma_i^2$ of the matrix $\Gamma_B^{1/2} J^\top J \Gamma_B^{1/2}$, i.e. $p_i^*$ are obtained by solving (23) with $w_i = \sigma_i^2$, and R is sampled as a diagonal matrix of Bernoulli entries, parameterized by $p_i^*$. The resulting sketching operator uses these probabilities to sample directions in the eigenbasis of $\Gamma_B^{1/2} J^\top J \Gamma_B^{1/2}$ (see $R^*$ construction in Theorem 3.3).

**Diagonal Sketches (DS).** A variant is obtained by restricting the sketching operator R to be diagonal, as studied in Theorem 3.4. Under this constraint, (23) admits importance weights of the form $w_i = (\Gamma_B)_{ii} (J^\top J)_{ii}$.

This estimator generalizes per column masking by allowing non-uniform, data-dependent sampling probabilities.

**Alternative Weights Proxies.** Beyond these two theoretically grounded constructions, (23) provides a more general design principle. Once a set of non-negative importance weights $w_i$ is specified, the resolution gives an unbiased sketch under a budget constraint, by sampling with probabilities proportional to $\sqrt{w_i}$.

Therefore, we explore alternative, computationally cheaper proxies for the weights $w_i$ that do not strictly correspond to the quantities derived in Section 3. In particular, we consider $\ell_1$, $\ell_2$-norms, and empirical variances of gradient coordinates as simple measures of magnitude and dispersion; we denote the corresponding proxy-based sketch $\ell_1$, $\ell_2$ and Var. We also consider importance weights based on the singular value decomposition of the batch gradient matrix G (as defined in Eq. (16)), and refer to this strat-
<figure><img src="image_3.png" /><figcaption>(a) Impact of correlation in Bernoulli sampling</figcaption></figure>

<figure><img src="image_4.png" /><figcaption>(b) Comparison between Masking and Sketching Methods</figcaption></figure>

<caption>Figure 1. Comparison of sampling strategies and scoring methods.</caption>

egy as *G*-Singular Values (G-SV). Finally, since the optimal probabilities scale with $\sqrt{w_i}$, we additionally evaluate the squared versions of these proxies.

# 5. Numerical Analysis

**Experimental Setting.** We first train 4-layer MLPs on MNIST (Deng, 2012): input dimension 784, two hidden layers of width 64, and a 10-way output. Training uses SGD without momentum nor a learning-rate schedule, for 50 epochs with cross-entropy loss, and gradient clipping at norm 1. For each seed, the learning rate is cross-validated over the grid $\{10^{-0.25i} \mid i \in [0, 12]\}$, and we report results for the best-performing value. In the figures, we approximate VJPs at all layers, except for the baseline, obtained with a standard training without any VJP approximation.

**Correlation in Sampling.** As discussed in Theorem 3.1, directions can be sampled either coordinate-wise independently, or by using correlated sampling schemes to ensure a fixed size rank. This latter construction reduces variability in budget allocation across iterations. In practice, we observe that enforcing this rank constraint leads to slightly improved performance, particularly in low-budget regimes, although the effect remains relatively modest. Results are reported in Figure 1a, and we adopt the correlated variant as the default for the subsequent discussions.

**Uniform Masks.** We now compare methods with one another and start by looking at uniform masking methods in opposition to data-dependent sketches (see Section 4). As the former are agnostic to local VJP's characteristics, they provide natural baselines that isolate the effect of injecting noise. Figure 1b shows that the latter consistently outperform the three agnostic methods, indicating that leveraging local information leads to visible empirical gains.

**Weight Proxies Comparison.** In Section 4, importance weights-based methods were introduced. Figure 2a reports a comparison between different weights proxies strategies. All methods yield very similar accuracy curves, indicating that no proxy incontestably dominates the others. Nonetheless, $\ell_1$-based probabilities lie on the upper envelope of the curves. Although the differences remain limited, this trend is consistent, and, therefore, $\ell_1$-based sampling is therefore used as a default choice in the subsequent experiments.

**Coordinate- VS Spectral-Based Strategies.** We now consider strategies that rely on SVD, namely Rank Constrained Sketching and the approach based on singular values of *G* and compare them to methods grounded in the canonic coordinate space. These approaches are computationally more expensive than aforementioned proxies, and are hence expected to reach better accuracy. Figure 2b confirms this intuition: spectral methods consistently outperform simpler proxies across the range of sampling budgets. G-SV strategy yields better results than its square root counterpart and is retained in the following experiments.

**BagNet and ViT.** To assess scalability of our approaches, we conduct experiments with the CIFAR-10 dataset (Krizhevsky et al.), and larger architectures, namely BagNet-17 (Brendel & Bethge, 2019), and Visual Transformers (ViT) (Dosovitskiy et al., 2021) of comparable sizes. The former is conceptually and structurally close to ResNets He et al. (2015), the key difference being that BagNet relies mainly on $1 \times 1$ convolutions, which we assimilate as linear layers and sketch. The latter is a transformer, and thus consists in attention blocks followed by linear feed-forward layers, to which we apply sketching operators as well.

For both architectures, we established a baseline, and then applied our VJP approximation methods to all linear layers,
<figure><img src="image_5.png" /><figcaption>(a) Comparison of simple proxies</figcaption></figure>

<figure><img src="image_6.png" /><figcaption>(b) Comparison of spectral and coordinates based methods</figcaption></figure>

<caption>Figure 2. Comparison of types of weighted methods.</caption>

except the output classification layer. These approximations were performed for budgets $p \in \{0.05, 0.1, 0.2, 0.5\}$, with learning rates cross-validated over five logarithmically spaced values around the baseline setting. Appendix provides training and ViT-specific architectural details.

**Discussion:** Figure 3 reports the results obtained for six retained methods on each architecture. Overall, these results are very encouraging, with Diagonal Sketching consistently emerging as a particularly strong choice. On both architectures, the accuracy degradation of our approaches remains limited, even for relatively small values of the sampling parameter. Furthermore, our plots clearly highlight the benefits of using data- and/or Jacobian-dependent strategies over uniform masking.

Concerning RCS, although we exposed its optimality for local distortion, it does not systematically dominate simpler strategies in practice. We conjecture that this is due to the local nature of the approximation: sketches are constructed independently at each VJP, without accounting for dynamics through subsequent layers.

## 6. Conclusion & Future Work

In this work we initiated the study of noisy unbiased backpropagation in a computational DAG based on approximate VJPs. The analysis of how variance evolves during backpropagation led to novel design principles for unbiased VJPs under budget constraints.

We then derived several VJP approximation techniques elaborating on these principles and conducted an extensive empirical comparison on MLP architectures. After establishing their effectiveness in this controlled setting, we further applied the proposed methods to larger models, where they similarly exhibit solid results and confirm the scalability of our approach.

Our results already show strong performance without any cross-layer coordination between sketches. A natural next step is to design coordinated sketching strategies that account for variance propagation through the DAG and potentially improve general performance. From a more practical point of view, estimating costly statistics intermittently rather than at each step, and/or jointly adapting optimization hyperparameters and sketching policies are promising future work directions.

## 7. Related Work

**Weight Gradient Compression.** Gradient compression in distributed and federated learning primarily acts after backpropagation, directly on parameter gradients:

$$ \hat{h}_i = C(h_i), \quad \mathbb{E}[C(h_i) | h_i] = h_i, $$

covering classical unbiased compressors and their analyses (Wang et al., 2023; Konečný et al., 2016; Horváth & Richtárik, 2020). When the compressor is (or is approximated by) a random linear map, $C(v) = S_i v$, this becomes exactly our parameter-side sketch $\hat{J}_i = S_i J_i$. Sparsification (Stich et al., 2018) corresponds to diagonal $S_i$ (mask-and-rescale), and error-feedback schemes (Richtárik et al., 2021) can be interpreted as adding a stateful correction term to compensate bias induced by repeatedly applying $S_i$. A key distinction with our focus is where the randomness enters: weight-gradient compressors operate on the final gradient signal $h_i$ and are largely agnostic to the internal Jacobian factorization, whereas we target the intermediate VJPs to trade compute/activation-traffic for variance. Note also that a parallel line of work seeks to overlap computation with weight gradients communication (Nabli et al., 2025), that we could combine with the results of this paper.
<figure><img src="image_7.png" /></figure>

**Backward-pass Compression.** Coordinate masking methods such as meProp (Sun et al., 2017) correspond to diagonal sketches $R_{ij}$ that keep a subset of coordinates of $\hat{g}_j$, in the line of sketches analyzed in Lemma 3.4 (although we enforce unbiasedness). More structured policies, e.g. skipping Jacobian applications on selected layers, corresponding to masking on the chosen edges, which covers layer dropping schemes reported for MLP-like architectures (Fagnou et al., 2025; Neiterman & Ben-Artzi, 2024). Pipeline settings that compress backward activations across devices also amount to picking $R_{ij}$ on inter-stage edges (often diagonal masks or low-rank sketches), as in pipeline-oriented designs (Kong et al., 2025; Narayanan et al., 2019; Huang et al., 2019). Batch-level masking of gradients (Bershatsky et al., 2022) is a special case where the same mask is applied to all samples in a mini-batch.

**Quantization and Randomized Arithmetic.** Mixed precision and quantized training accelerate linear algebra in forward/backward passes (Micikevicius et al., 2017; Johnson, 2018; Wiedemann et al., 2020). Unbiased stochastic quantizers can be written as random maps $Q(v)$ with $E[Q(v) | v] = v$; when $Q$ is coordinate-wise and unbiased, it is equivalent to applying a random diagonal operator to $v$ in expectation, hence can be absorbed into our $R_{ij}$ acting on $\hat{g}_j$. In this sense, many “unbiased quantization” schemes are instances of our local-unbiased sketching assumption, differing mainly by the distributional form of $R_{ij}$.

**Alternatives to Backprop.** Methods such as direct feedback alignment or decoupled training replace the exact chain of Jacobians by synthetic or partially decoupled signals (Nøkland, 2016; Rivaud et al., 2025; Belilovsky et al., 2020; Jaderberg et al., 2017). These approaches are sometimes framed as using different operators $\hat{J}_{ij}$ that are not of the form $J_{ij}R_{ij}$ with $E[R_{ij}] = I$; in particular, they generally do not preserve conditional unbiasedness of true

gradients and may change the optimization objective. By contrast, we keep the exact computational graph but randomize VJPs with unbiasedness constraints, so the only effect on SGD is via variance (Section 2.2).

Unbiased alternatives that rely on forward-mode signals (JVPs) propagate random probes through the network and then form gradient estimators (Ren et al., 2022; Fournier et al., 2023). At a high level, these can be viewed as constructing an unbiased estimator of the full-network Jacobian action using random sketches; however, the randomness is injected via forward-mode probes rather than via local VJP replacements. Unbiased truncations for sequences (Tallec & Ollivier, 2017) can be seen as masking temporal edges (setting some Jacobian factors to zero) with appropriate rescaling to keep expectations correct—again matching the “randomly drop factors but preserve unbiasedness” principle, though specialized to time-unrolled graphs.

A related but distinct direction applies masks in the forward pass as well (hence changing the function being differentiated); this can be written as replacing $J_{ij}$ itself by a different Jacobian $\tilde{J}_{ij}$, which in general breaks our conditional-unbiasedness requirement and can yield biased gradient surrogates (e.g. when the mask is not compensated in the backward pass) (Singh et al., 2025). Our framework isolates the “safe” regime: preserve unbiasedness via $R_{ij}$ while controlling the induced variance.

**Sketching for Linear Approximation.** Independent of deep learning, masking and sketching are classical tools for efficient sparse approximation (Kireeva & Tropp, 2024), including kernel subsampling with compensation (Rudi et al., 2017), sketching for ridge regression (Gazagnadou et al., 2022), and communication-efficient federated protocols based on sketches (Shrivastava et al., 2024). Our contribution connects these sketching ideas to where they appear in backpropagation: we study sketches as local un-
biased replacements of VJPs (i.e. choices of $R_{ij}$), and analyze how the induced variance propagates through the DAG (Proposition 2.2), thereby making the *depth and placement* of sketching a first-order design choice rather than a purely systems-level consideration.

## Acknowledgements

EO acknowledges funding from PEPR IA (grant SHARP ANR-23-PEIA-0008). He was granted access to the AI resources of IDRIS under the allocation 2025-AD011015884R1. KB thanks coworkers Pierre Aguié, Loïck Chambon and Alessia Rigonat for insightful discussions and/or technical support. He acknowledges funding from PEPR IA (grant REDEEM ANR-23-PEIA-0005) and was granted access to the AI resources of IDRIS under the allocation AD011017074. LM acknowledges funding from PR[AI]RIE-PSAI – Paris School of Artificial Intelligence, reference ANR-23-IACL-0008.

The authors gratefully acknowledge Christophe Bouder for gracefully facilitating access to SCAI GPU resources.

## Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

## References

Barnes, L. P., Cameron, S., and Howard, B. On unbiased low-rank approximation with minimum distortion, 2025. URL https://arxiv.org/abs/2505.09647.

Baur, W. and Strassen, V. The complexity of partial derivatives. *Theoretical computer science*, 22(3):317–330, 1983.

Baydin, A. G., Pearlmutter, B. A., Radul, A. A., and Siskind, J. M. Automatic differentiation in machine learning: a survey. *Journal of machine learning research*, 18(153):1–43, 2018.

Belilovsky, E., Eickenberg, M., and Oyallon, E. Decoupled greedy learning of cnns. In *International Conference on Machine Learning*, pp. 736–745. PMLR, 2020.

Benzing, F., Gauy, M. M., Mujika, A., Martinsson, A., and Steger, A. Optimal kronecker-sum approximation of real time recurrent learning, 2019. URL https://arxiv.org/abs/1902.03993.

Bershatsky, D., Mikhalev, A., Katrutsa, A., Gusak, J., Merkulov, D., and Oseledets, I. Memory-efficient backpropagation through large linear layers. *arXiv preprint arXiv:2201.13195*, 2022.

Bottou, L. et al. Stochastic gradient learning in neural networks. *Proceedings of Neuro-Nimes*, 91(8):12, 1991.

Brendel, W. and Bethge, M. Approximating cnns with bag-of-local-features models works surprisingly well on imagenet. *arXiv preprint arXiv:1904.00760*, 2019.

Coleman, T. F. and Moré, J. J. Estimation of sparse jacobian matrices and graph coloring blems. *SIAM journal on Numerical Analysis*, 20(1):187–209, 1983.

Deng, L. The mnist database of handwritten digit images for machine learning research. *IEEE Signal Processing Magazine*, 29(6):141–142, 2012.

Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale, 2021. URL https://arxiv.org/abs/2010.11929.

Fagnou, E., Caillon, P., Delattre, B., and Allauzen, A. Accelerated training through iterative gradient propagation along the residual path. *arXiv preprint arXiv:2501.17086*, 2025.

Fournier, L., Rivaud, S., Belilovsky, E., Eickenberg, M., and Oyallon, E. Can forward gradient match backpropagation? In *International Conference on Machine Learning*, pp. 10249–10264. PMLR, 2023.

Gazagnadou, N., Ibrahim, M., and Gower, R. M. Ridgsketch: a fast sketching based solver for large scale ridge regression. *SIAM Journal on Matrix Analysis and Applications*, 43(3):1440–1468, 2022.

Griewank, A. and Walther, A. Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation. *ACM Transactions on Mathematical Software (TOMS)*, 26(1):19–45, 2000.

Griewank, A. and Walther, A. *Evaluating derivatives: principles and techniques of algorithmic differentiation*. SIAM, 2008.

He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition, 2015. URL https://arxiv.org/abs/1512.03385.

Herrmann, J., Beaumont, O., Eyraud-Dubois, L., Hermann, J., Joly, A., and Shilova, A. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory. *arXiv preprint arXiv:1911.13214*, 2019.
Horn, R. A. and Johnson, C. R. *Matrix Analysis*. Cambridge University Press, Cambridge; New York, 2nd edition, 2013. ISBN 9780521839402.

Horváth, S. and Richtárik, P. A better alternative to error feedback for communication-efficient distributed learning. *arXiv preprint arXiv:2006.11077*, 2020.

Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. *Advances in neural information processing systems*, 32, 2019.

Jaderberg, M., Czarnecki, W. M., Osindero, S., Vinyals, O., Graves, A., Silver, D., and Kavukcuoglu, K. Decoupled neural interfaces using synthetic gradients. In *International conference on machine learning*, pp. 1627–1635. PMLR, 2017.

Johnson, J. Rethinking floating point for deep learning. *arXiv preprint arXiv:1811.01721*, 2018.

Kireeva, A. and Tropp, J. A. Randomized matrix computations: Themes and variations. *arXiv preprint arXiv:2402.17873*, 2024.

Konečný, J., McMahan, H. B., Yu, F. X., Richtárik, P., Suresh, A. T., and Bacon, D. Federated learning: Strategies for improving communication efficiency. *arXiv preprint arXiv:1610.05492*, 2016.

Kong, B., Huang, X., Xu, Y., Liang, Y., Wang, B., and Yuan, K. Clapping: Removing per-sample storage for pipeline parallel distributed optimization with communication compression. *arXiv preprint arXiv:2509.19029*, 2025.

Krizhevsky, A., Nair, V., and Hinton, G. Cifar-10 (Canadian institute for advanced research). URL http://www.cs.toronto.edu/~kriz/cifar.html.

LeCun, Y., Bottou, L., Orr, G. B., and Müller, K.-R. Efficient backprop. In *Neural networks: Tricks of the trade*, pp. 9–50. Springer, 2002.

Lee, C., Cho, K., and Kang, W. Mixout: Effective regularization to finetune large-scale pretrained language models. *arXiv preprint arXiv:1909.11299*, 2019.

Li, S., Xue, F., Li, Y., and You, Y. Sequence parallelism: Making 4d parallelism possible. *arXiv preprint arXiv:2105.13120*, 2021.

Linnainmaa, S. Taylor expansion of the accumulated rounding error. *BIT Numerical Mathematics*, 16(2):146–160, 1976.

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization, 2019. URL https://arxiv.org/abs/1711.05101.

Marshall, A. W., Olkin, I., and Arnold, B. C. *Inequalities: Theory of Majorization and its Applications*, volume 143. Springer, second edition, 2011. doi: 10.1007/978-0-387-68276-1.

Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., et al. Mixed precision training. *arXiv preprint arXiv:1710.03740*, 2017.

Nabli, A., Fournier, L., ERBACHER, P., Serrano, L., Belilovsky, E., and Oyallon, E. ACCO: Accumulate while you communicate for communication-overlapped sharded LLM training. In *The Thirty-ninth Annual Conference on Neural Information Processing Systems*, 2025. URL https://openreview.net/forum?id=1qKUVyymXs.

Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. Pipedream: Generalized pipeline parallelism for dnn training. In *Proceedings of the 27th ACM symposium on operating systems principles*, pp. 1–15, 2019.

Neiterman, E. H. and Ben-Artzi, G. Layerdropback: A universally applicable approach for accelerating training of deep networks. *arXiv preprint arXiv:2412.18027*, 2024.

Nøkland, A. Direct feedback alignment provides learning in deep neural networks. *Advances in neural information processing systems*, 29, 2016.

Rahimi, A. and Recht, B. Random features for large-scale kernel machines. *Advances in neural information processing systems*, 20, 2007.

Ren, M., Kornblith, S., Liao, R., and Hinton, G. Scaling forward gradient with local losses. *arXiv preprint arXiv:2210.03310*, 2022.

Richtárik, P., Sokolov, I., and Fatkhullin, I. Ef21: A new, simpler, theoretically better, and practically faster error feedback. *Advances in Neural Information Processing Systems*, 34:4384–4396, 2021.

Rivaud, S., Fournier, L., Pumir, T., Belilovsky, E., Eickenberg, M., and Oyallon, E. PETRA: Parallel end-to-end training with reversible architectures. In *The Thirteenth International Conference on Learning Representations*, 2025. URL https://openreview.net/forum?id=0fhzSFsGUT.

Rojas, E., Kahira, A. N., Meneses, E., Gomez, L. B., and Badia, R. M. A study of checkpointing in large
scale training of deep neural networks. *arXiv preprint arXiv:2012.00825*, 2020.

Rudi, A., Carratino, L., and Rosasco, L. Falkon: An optimal large scale kernel method. *Advances in neural information processing systems*, 30, 2017.

Rumelhart, D. E., Hinton, G. E., and Williams, R. J. Learning representations by back-propagating errors. *nature*, 323(6088):533–536, 1986.

Shrivastava, M., Isik, B., Li, Q., Koyejo, S., and Banerjee, A. Sketching for distributed deep learning: A sharper analysis. *Advances in Neural Information Processing Systems*, 37:6417–6447, 2024.

Singh, V., Khalid, Z., Oyallon, E., and Belilovsky, E. Model parallelism with subnetwork data parallelism. *arXiv preprint arXiv:2507.09029*, 2025.

Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: a simple way to prevent neural networks from overfitting. *The journal of machine learning research*, 15(1):1929–1958, 2014.

Stich, S. U., Cordonnier, J.-B., and Jaggi, M. Sparsified sgd with memory. *Advances in neural information processing systems*, 31, 2018.

Sun, X., Ren, X., Ma, S., and Wang, H. meprop: Sparsified back propagation for accelerated deep learning with reduced overfitting. In *International Conference on Machine Learning*, pp. 3299–3308. PMLR, 2017.

Tallec, C. and Ollivier, Y. Unbiased online recurrent optimization. *arXiv preprint arXiv:1702.05043*, 2017.

Wang, J., Lu, Y., Yuan, B., Chen, B., Liang, P., De Sa, C., Re, C., and Zhang, C. Cocktailsgd: Fine-tuning foundation models over 500mbps networks. In *International Conference on Machine Learning*, pp. 36058–36076. PMLR, 2023.

Wiedemann, S., Mehari, T., Kepp, K., and Samek, W. Dithered backprop: A sparse and quantized backpropagation algorithm for more efficient deep neural network training. In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops*, pp. 720–721, 2020.

Woo, S., Park, B., Kim, B., Jo, M., Kwon, S. J., Jeon, D., and Lee, D. Dropbp: Accelerating fine-tuning of large language models by dropping backward propagation. *Advances in Neural Information Processing Systems*, 37:20170–20197, 2024.
# A. Proofs

## A.1. Proofs of Section 2

**Proposition A.1 (Restatement of Theorem 2.2).** Assume the seed at the output node is exact (i.e. $\hat{g}_{out} = g_{out}$) and that Theorem 2.1 holds. Then, for every node i of the DAG,

(i) (Unbiasedness.) $\mathbb{E}[\hat{g}_i|g_i] = g_i$, $\mathbb{E}[\hat{h}_i|h_i] = h_i$.

(ii) (Variance propagation.)

$$ \mathbb{E}[\|\hat{g}_i - g_i\|^2] = \sum_{j:i \to j} \mathbb{E}\left[\left\|\frac{1}{B}\sum_{b=1}^{B}(\hat{J}_{ij}^{(b)} - J_{ij}^{(b)})\hat{g}_j^{(b)}\right\|^2\right] + \mathbb{E}\left[\left\|\frac{1}{B}\sum_{b=1}^{B}\sum_{j:i \to j} J_{ij}^{(b)}(\hat{g}_j^{(b)} - g_j^{(b)})\right\|^2\right] \quad (24) $$

*Proof.* We proceed by backward induction over the DAG, starting from the output node.

**Unbiasedness.** For each sample b, the approximate backward recursion reads

$$ \hat{g}_i^{(b)} = \sum_{j:i \to j} \hat{J}_{ij}^{(b)} \hat{g}_j^{(b)}, \quad (25) $$

as stated in (11). Taking the conditional expectation given $(\hat{g}_j^{(b)})_{j:i \to j}$ and using Theorem 2.1 yields

$$ \mathbb{E}[\hat{g}_i^{(b)} | (\hat{g}_j^{(b)})_j] = \sum_{j:i \to j} J_{ij}^{(b)} \hat{g}_j^{(b)}. \quad (26) $$

By the induction hypothesis, in reverse order from $g_{out}^{(b)}$, $\mathbb{E}[\hat{g}_j^{(b)}] = g_j^{(b)}$, hence

$$ \mathbb{E}[\hat{g}_i^{(b)}] = \sum_{j:i \to j} J_{ij}^{(b)} g_j^{(b)} = g_i^{(b)}. \quad (27) $$

Averaging over the batch gives $\mathbb{E}[\hat{g}_i | g_i] = g_i$.

The same argument applies to $\hat{h}_i^{(b)} = \hat{J}_i^{(b)} \hat{g}_i^{(b)}$, yielding $\mathbb{E}[\hat{h}_i | h_i] = h_i$.

**Variance propagation.** For each sample b, we decompose

$$ \hat{g}_i^{(b)} - g_i^{(b)} = \sum_{j:i \to j} (\hat{J}_{ij}^{(b)} - J_{ij}^{(b)})\hat{g}_j^{(b)} + \sum_{j:i \to j} J_{ij}^{(b)}(\hat{g}_j^{(b)} - g_j^{(b)}). \quad (28) $$

Taking the squared norm gives

$$ \|\hat{g}_i^{(b)} - g_i^{(b)}\|^2 = \left\| \sum_{j:i \to j} (\hat{J}_{ij}^{(b)} - J_{ij}^{(b)})\hat{g}_j^{(b)} \right\|^2 + \left\| \sum_{j:i \to j} J_{ij}^{(b)}(\hat{g}_j^{(b)} - g_j^{(b)}) \right\|^2 + 2\langle \cdot, \cdot \rangle. \quad (29) $$

The cross-term $\langle \cdot, \cdot \rangle$ has zero expectation by conditional unbiasedness of $\hat{J}_{ij}^{(b)}$ as stated in Theorem 2.1. Local unbiasedness of Jacobians also allows to write

$$ \|\hat{g}_i^{(b)} - g_i^{(b)}\|^2 = \sum_{j:i \to j} \left\| (\hat{J}_{ij}^{(b)} - J_{ij}^{(b)})\hat{g}_j^{(b)} \right\|^2 + \left\| \sum_{j:i \to j} J_{ij}^{(b)}(\hat{g}_j^{(b)} - g_j^{(b)}) \right\|^2. \quad (30) $$

Finally averaging over the batch (and using sample wise independance) and taking expectations yields

$$ \mathbb{E}\|\hat{g}_i - g_i\|^2 = \sum_{j:i \to j} \mathbb{E}\left\|\frac{1}{B}\sum_{b=1}^{B}(\hat{J}_{ij}^{(b)} - J_{ij}^{(b)})\hat{g}_j^{(b)}\right\|^2 + \mathbb{E}\left\|\frac{1}{B}\sum_{b=1}^{B}\sum_{j:i \to j} J_{ij}^{(b)}(\hat{g}_j^{(b)} - g_j^{(b)})\right\|^2. \quad (31) $$

□
## A.2. Proofs of Section 3

**Lemma A.2 (Restatement of Theorem 3.1).** Let $M$ be a fixed matrix in $\mathbb{R}^{m \times n}$. Let $q = \min(m, n)$ and denote the SVD of $M$ by $M = \sum_{i=1}^q \sigma_i u_i v_i^\top$, where $\sigma_i, u_i, v_i$ are respectively its singular values, left and right singular vectors. Then among matrices $S$ of rank at most $r$ where $r < q$, such that $\mathbb{E}[S] = M$, the choice that minimizes the error in squared Frobenius norm $\mathbb{E}\|M - S\|_F^2$ is obtained by letting

$$S = \sum_{i=1}^{q} \sigma_i \frac{Z_i}{p_i} u_i v_i^\top$$

where the $Z_i$ are correlated Bernoulli random variables with respective parameters $p_i$. The $Z_i$ are such that $\sum_{i=1}^q Z_i \equiv r$, and the $\{p_i\}_{i \in [q]}$ are minimizers of $\sum_{i=1}^q \sigma_i^2/p_i$ among weights $p_i \in [0, 1]$ that sum to $r$.

Let us first state the following Lemma:

**Lemma A.3.** For fixed $i \in [q]$, the function $g_i$ defined on $\mathbb{R}^{m \times n}$ by

$$g_i(M) := \sum_{j=1}^{i} \sigma_j(M),$$

that is the function returning the sum of the top-i singular values is convex.

This follows for instance from Theorem 4.3.27, p. 194 in Horn & Johnson (2013).

*Proof of Theorem 3.1.* For all $i \in [q]$, we have by Jensen's inequality:

$$\mathbb{E}g_i(S) = \mathbb{E}\sum_{j=1}^{i} \sigma_j(S) \geq g_i(\mathbb{E}(S)) = g_i(M) = \sum_{j=1}^{i} \sigma_j. \qquad (32)$$

Now, let $i_0$ be an integer such that

$$\sigma_{i_0} > \frac{\sum_{j=i_0+1}^{q} \sigma_j}{r - i_0} \geq \sigma_{i_0+1}, \qquad (33)$$

where if $\sigma_1 \leq r^{-1} \sum_{j=1}^q \sigma_j$, we let $i_0 = 0$.

Let us denote, for $i \in [i_0]$, $\mathbb{E}\sigma_i(S)$ by $s_i$. We thus have

$$s_1 \geq \sigma_1, s_1 + s_2 \geq \sigma_1 + \sigma_2, \dots, \sum_{j=1}^{i_0} s_j \geq \sum_{j=1}^{i_0} \sigma_j.$$

This property is also known as the fact that vector $(\sigma_j)_{j \in [i_0]}$ is weakly majorized by vector $(s_j)_{j \in [i_0]}$, also denoted by $\sigma \prec_w s$, see Definition (11), p. 12 in Marshall et al. (2011).

Split then the singular value decomposition of $S$ into the first $i_0$ singular values, $\sigma_1(S), \dots, \sigma_{i_0}(S)$ and into the last non-zero singular values, $\sigma_{i_0+1}(S), \dots, \sigma_r(S)$, and let $X$ be a uniformly random selection among these last $r - i_0$ singular values of $S$. Use now Inequality (32) with $i = q$ to get

$$s_1 + \dots + s_{i_0} + (r - i_0)\mathbb{E}(X) \geq \sum_{j=1}^{q} \sigma_j,$$

so that

$$\mathbb{E}(X) \geq \frac{1}{r - i_0} \left[ \sum_{j=1}^{q} \sigma_j - \sum_{i=1}^{i_0} s_i \right].$$
This implies

$$ \mathbb{E}\|S\|_F^2 = \sum_{i=1}^{r} \mathbb{E}[\sigma_i(S)^2] \geq \sum_{i=1}^{i_0} s_i^2 + (r - i_0)(\mathbb{E}(X))^2 \quad (34) $$

$$ \geq \sum_{i=1}^{i_0} s_i^2 + \frac{1}{r - i_0} \left[ \sum_{j=1}^{q} \sigma_j - \sum_{i=1}^{i_0} s_i \right]^2. \quad (35) $$

Introduce the function

$$ h((x_i)_{i \in [i_0]}) := \sum_{i=1}^{i_0} x_i^2 + \frac{1}{r - i_0} \left[ \sum_{j=1}^{q} \sigma_j - \sum_{i=1}^{i_0} x_i \right]^2, $$

and consider the convex set

$$ A := \{ (x_i)_{i \in [i_0]} \in \mathbb{R}^{i_0} : h((x_i)_{i \in [i_0]}) < h((\sigma_i)_{i \in [i_0]}) \}. $$

Let us show that A does not intersect with the set

$$ B := \{ (x_1 \ge \cdots \ge x_{i_0} : \sigma <_w x) \}. $$

To that end, note that both sets are convex. Note that, at $x = \sigma$, the gradient of $h$ reads

$$ \nabla h(\sigma) = 2 \left( \sigma_i - \frac{1}{r - i_0} \sum_{j=i_0+1}^{q} \sigma_j \right)_{i \in [i_0]} . $$

Clearly, for all $x \in A$,

$$ \langle \nabla h(\sigma), x - \sigma \rangle < 0. $$

Let now $x \in B$. To show that A and B are disjoint, it will suffice to prove that

$$ \forall x \in B, \langle \nabla h(\sigma), x - \sigma \rangle \geq 0. $$

Let

$$ \alpha := \frac{1}{r - i_0} \sum_{j=i_0+1}^{q} \sigma_j. $$

Write then

$$ \begin{aligned} \langle \nabla h(\sigma), x - \sigma \rangle &= \sum_{i=1}^{i_0} (x_i - \sigma_i) [\sigma_i - \alpha] \\ &= \sum_{i=1}^{i_0} (x_i - \sigma_i) [\sigma_{i_0} - \alpha + \sum_{j=i}^{i_0-1} (\sigma_j - \sigma_{j+1})] \\ &= (\sigma_{i_0} - \alpha) \sum_{i=1}^{i_0} (x_i - \sigma_i) \\ &\quad + \sum_{j=i_0+1}^{i_0-1} (\sigma_j - \sigma_{j+1}) \sum_{i=1}^{j} (x_i - \sigma_i). \end{aligned} $$

The property $\sigma <_w x$ implies $\sum_{i=1}^{j} (x_i - \sigma_i) \ge 0$ for all $j \in [i_0]$. The terms $(\sigma_j - \sigma_{j+1})$ are non-negative by assumption. Finally, the assumption (33) ensures that $\sigma_{i_0} - \alpha > 0$. Thus all terms in the last display are non-negative, so that for all $x \in B$,

$$ \langle \nabla h(\sigma), x - \sigma \rangle \geq 0. $$

We thus have, since $\sigma <_w s$, that $h(s) \ge h(\sigma)$. However, (34) ensures that $\mathbb{E}\|S\|_F^2 \ge h(s)$.

By explicating $h(\sigma)$, we get the lower bound:

$$ \mathbb{E}\|S\|_F^2 \geq \sum_{i=1}^{i_0} \sigma_i^2 + \frac{1}{r - i_0} \left[ \sum_{j=i_0+1}^{q} \sigma_j \right]^2. \quad (36) $$
We now describe the construction of the random matrix $S$ that achieves equality in this lower bound. Recall that the SVD of $M$ be given by

$$M = \sum_{i=1}^{q} \sigma_i u_i v_i^\top.$$

Define for $i = i_0 + 1, \dots, q$ the quantity

$$p_i := (r - i_0) \frac{\sigma_i}{\sum_{j=i_0+1}^{q} \sigma_j}.$$

By the second inequality in the definition (33) of $i_0$, it follows that

$$p_i \in [0, 1], i \in \{i_0 + 1, \dots, q\}.$$

Let now $U_1$ be uniformly distributed over $[0, 1]$, and let

$$U_j = U_1 + j - 1, j = 2, \dots, r - i_0.$$

Finally, we define

$$\forall i \in \{i_0 + 1, \dots, q\}, \quad Z_i = \begin{cases} 1 & \text{if } \exists j \in [r - i_0] : U_j \in \left[ \sum_{l=i_0+1}^{i-1} p_l, \sum_{l=i_0+1}^{i} p_l \right], \\ 0 & \text{otherwise.} \end{cases}$$

By construction, we have $\sum_{i=i_0+1}^{q} p_i = r - i_0$, so that each $Z_i$ is Bernoulli with parameter $p_i$, and moreover

$$\sum_{i=i_0+1}^{q} Z_i = r - i_0 \text{ almost surely.}$$

We then let

$$S := \sum_{i=1}^{i_0} \sigma_i u_i v_i^\top + \sum_{i=i_0+1}^{q} Z_i \frac{\sigma_i}{p_i} u_i v_i^\top.$$

The fact that $\mathbb{E}S = M$ follows directly from the fact that the $Z_i$ are Bernoulli with parameter $p_i$. Finally, we have

$$\begin{aligned} \mathbb{E}\|S\|_F^2 &= \sum_{i=1}^{i_0} \sigma_i^2 + \sum_{i=i_0+1}^{q} \frac{\sigma_i^2}{p_i} \\ &= \sum_{i=1}^{i_0} \sigma_i^2 + \sum_{i=i_0+1}^{q} \sigma_i^2 \frac{\sum_{j=i_0+1}^{q} \sigma_j}{(r - i_0)\sigma_i} \\ &= \sum_{i=1}^{i_0} \sigma_i^2 + \frac{1}{r - i_0} \left[ \sum_{j=i_0+1}^{q} \sigma_j \right]^2 \end{aligned}$$

We have thus exhibited a random matrix $S$ of rank at most $r$, satisfying $\mathbb{E}[S] = M$, that achieves the lower bound (36). Consequently, this bound is tight and characterizes the minimum possible value of $\mathbb{E}\|M - S\|_F^2$ among all unbiased rank-$r$ sketches.

Moreover, for any sketch of the form

$$S = \sum_{i=1}^{q} \sigma_i \frac{Z_i}{p_i} u_i v_i^\top, \quad \mathbb{P}(Z_i = 1) = p_i, \quad \sum_{i=1}^{q} Z_i = r$$

we have

$$\mathbb{E}\|S\|_F^2 = \sum_{i=1}^{q} \frac{\sigma_i^2}{p_i}.$$
Therefore, minimizing $\mathbb{E}\|M-S\|_F^2$ over unbiased rank-$r$ sketches is equivalent to solving the convex optimization problem

$$ \min \sum_{i=1}^{q} \frac{\sigma_i^2}{p_i} \quad \text{subject to} \quad \sum_{i=1}^{q} p_i = r, \quad 0 < p_i \le 1. $$

The weights ($p_i$) defined above attain this minimum. Since the objective is convex and the constraints are affine, optimality follows from the Karush-Kuhn-Tucker conditions, which yield the thresholding structure

$$ p_i^* = \min \left( 1, \sqrt{\frac{\sigma_i^2}{\lambda}} \right), $$

for some $\lambda > 0$ chosen so that $\sum_i p_i^* = r$. This exactly corresponds to the definition of $i_0$ and the probabilities constructed above, completing the proof.

□

**Proposition A.4 (Restatement of Theorem 3.3).** Let $J \in \mathbb{R}^{m \times n}$ and $\Gamma_B$ defined as above (i.e. the empirical second moment matrix of the batch $\mathcal{B}$). Define the symmetric matrix

$$ \Gamma_B^{1/2} J^\top J \Gamma_B^{1/2} = U \Sigma U^\top, \qquad (37) $$

where $U$ is orthogonal and $\Sigma = \text{diag}(\sigma_1^2, \dots, \sigma_n^2)$ with $\sigma_1 \ge \dots \ge \sigma_n \ge 0$. Among all random matrices $R$ with $\mathbb{E}[R] = I_n$ and whose rank is bounded by $r$, a minimizer of the $\mathbb{L}^2$ cost Equation (15) is attained by sketches that are diagonal in the eigenbasis of $\Gamma_B^{1/2} J^\top J \Gamma_B^{1/2}$:

$$ \begin{aligned} R^* &= \Gamma_B^{1/2} U B U^\top \Gamma_B^{-1/2} \\ B &= \text{diag}(\frac{z_1}{p_1^*}, \dots, \frac{z_n}{p_n^*}), \quad \text{with } z_i \sim \mathcal{B}(p_i^*) \end{aligned} \qquad (38) $$

with $\mathcal{B}(p)$ being the Bernoulli distribution of parameter $p$, and with selection probabilities $\{p_i^*\}_{i=1}^n$ minimizing $\sum_{i=1}^n \sigma_i^2/p_i$ under the constraint $\sum_i p_i = r$.

*Proof.* We want to minimize the distortion

$$ \mathcal{L}(R) = \frac{1}{B} \sum_{b=1}^{B} \mathbb{E}[\|J_{ij}^{(b)} \hat{g}_j^{(b)} - J_{ij}^{(b)} R \hat{g}_j^{(b)}\|^2], \qquad (39) $$

as defined in Equation (15). As we stated in the main text, via the sketch of proof, the use of Theorem 3.2 for linear settings and the constraint $\mathbb{E}[R] = I_n$ yield

$$ \mathcal{L}(R) = \mathbb{E}[\text{Tr}(J R \Gamma_B R^\top J^\top)]. \qquad (40) $$

Leveraging the fact that, as a second moment matrix $\Gamma_B := \frac{1}{B}GG^\top$ is positive definite, so $\Gamma_B^{1/2}$ exists and we write $S = JR\Gamma_B^{1/2}$. We then have

$$ \mathcal{L}(R) = \mathbb{E}[\text{Tr}(SS^T)] = \mathbb{E}[\|S\|^2]. \qquad (41) $$

Since we look at sketches $R$ such as $\text{rank}(R) \le r$, then $S$ is of rank at most $r$ as well. Also, since $\mathbb{E}[R] = I_n$, we define $M := \mathbb{E}[S] = J\Gamma^{1/2}$, and the minimization of $\mathcal{L}$ is exactly the problem stated in Theorem 3.3. Applying Theorem 3.1 to $M$ we get that an optimal rank-$r$ unbiased sketch $S$ is obtained by sampling in its singular vector basis. Let $M = P \Sigma^{1/2} U^\top$ be an SVD of $M$, equivalently

$$ M^\top M = U \Sigma U^\top \quad \text{with} \quad M^\top M = \Gamma_B^{1/2} J^\top J \Gamma_B^{1/2}, $$

which matches (18). Therefore, an optimal choice is

$$ S^* = P \Sigma^{1/2} U^\top B \quad \text{with} \quad B = \text{diag}\left(\frac{z_1}{p_1^*}, \dots, \frac{z_n}{p_n^*}\right), \quad z_i \sim \mathcal{B}(p_i^*), \quad \sum_{i=1}^n z_i = r \qquad (42) $$
where the probabilities $p_i^*$ solve the convex program stated in Theorem 3.3.

Finally, since $S^* = JR^*\Gamma_B^{1/2}$ and $M = J\Gamma_B^{1/2} = P\Sigma^{1/2}U^\top$, we have

$$S^* = MB = J\Gamma_B^{1/2}B. \qquad (43)$$

Identifying both expressions yields a choice for $R^*$

$$R^*\Gamma_B^{1/2} = \Gamma_B^{1/2}B, \qquad (44)$$

hence

$$R^* = \Gamma_B^{1/2}B\Gamma_B^{-1/2}. \qquad (45)$$

Writing $B$ in the eigenbasis of $\Gamma_B^{1/2}J^\top J\Gamma_B^{1/2}$ finally gives

$$R^* = \Gamma_B^{1/2}U \text{diag}\left(\frac{z_1}{p_1^*}, \dots, \frac{z_n}{p_n^*}\right) U^\top \Gamma_B^{-1/2}, \qquad (46)$$

which is exactly the form (19).

<figure><img src="image_8.png" /></figure>

**Lemma A.5 (Restatement of Theorem 3.4).** We now restrict $R$ to be diagonal, $R = \text{diag}(r_1, \dots, r_n)$ with

$$r_i = z_i/p_i, \quad z_i \sim \mathcal{B}(p_i) \quad \text{independent, } p_i \in [0, 1],$$

so that $\mathbb{E}[R] = I_n$. For such diagonal mask, under the expected rank constraint $\sum_{i=1}^n p_i \le r$, the minimal $L^2$ distortion Equation (15) is obtained by choosing probabilities $p_i$ that minimize $\inf_p \sum_{i=1}^n \frac{a_i}{p_i}$ under the constraint $\sum_i p_i \le r$, where $a_i := (\Gamma_B)_{ii} (J^\top J)_{ii}$.

*Proof.* With $R = \text{diag}(r_1, \dots, r_n)$ (hence $(I-R)^\top = (I-R)$), expand

$$\begin{aligned} \text{Tr}(J^\top J(I-R)\Gamma_B(I-R)) &= \sum_{i,j=1}^n (J^\top J)_{ij} (1-r_i)(1-r_j)(\Gamma_B)_{ij} \\ &= \sum_{i=1}^n (J^\top J)_{ii} (1-r_i)^2 (\Gamma_B)_{ii} + \sum_{i \ne j} (J^\top J)_{ij} (1-r_i)(1-r_j)(\Gamma_B)_{ij}. \end{aligned} \qquad (47)$$

Under the diagonal mask model $r_i = z_i/p_i$ with $z_i \sim \mathcal{B}(p_i)$ independent of $G$, we have:

$$\begin{aligned} \mathbb{E}[1-r_i] &= 0, \\ \mathbb{E}[(1-r_i)(1-r_j)] &= \mathbb{E}[(1-r_i)]\mathbb{E}[(1-r_j)] = 0 \quad (i \ne j), \\ \mathbb{E}[(1-r_i)^2] &= \frac{1}{p_i} - 1 \quad (p_i > 0). \end{aligned} \qquad (48)$$

Taking expectations yields

$$\mathbb{E}[\text{Tr}(J^\top J(I-R)\Gamma_B(I-R))] = \sum_{i=1}^n (J^\top J)_{ii} \left( \frac{(\Gamma_B)_{ii}}{p_i} - (\Gamma_B)_{ii} \right). \qquad (49)$$

Since we aim at minimizing with $p_i$ as the variable, we get rid of the second term of the difference, hence the minimization objective

$$\sum_{i=1}^{n} \frac{a_i}{p_i}, \quad \text{with } a_i := (\Gamma_B)_{ii} (J^\top J)_{ii}. \qquad (50)$$
Therefore, minimizing $\mathcal{L}$ over diagonal masks with budget $\sum_{i=1}^{n} p_i \le r$ is equivalent to

$$ \min_{p \in [0,1]^n, \sum_i p_i \le r} \sum_{i=1}^{n} \frac{a_i}{p_i} \qquad (51) $$

This matches the constraints stated in Proposition 3.3 and concludes the proof. $\square$

<figure><img src="image_9.png" /></figure>
<figure><img src="image_10.png" /><figcaption>(a) Masking Methods</figcaption></figure>

<figure><img src="image_11.png" /><figcaption>(b) Coordinate-based Sketching Methods</figcaption></figure>

<figure><img src="image_12.png" /><figcaption>(c) Spectral-based Sketching Methods</figcaption></figure>

<caption>Figure 4. Impact of VJP Approximation Location in MLPs.</caption>

# B. Additional Numerical Details

## B.1. Additional Variance Location Result on MLP

In the main body of the paper, we exclusively considered experiments in which our methods are implemented for all layers of the MLP. In this appendix, we complement these results by comparing full-layer VJP approximation with variants, where it is performed either only to the first layer or only to the last layer. Results are reported in Figure 4. As one might intuitively expect, approximating local VJPs solely in the last layer leads to a more pronounced degradation in accuracy than when noise is injectes only in the first layer. Beyond this empirical observation, these results may provide practical insights for distributed training settings and especially, in straggler-mitigation scenarios where one may wish to apply VJP approximations selectively at a slower compute node rather than uniformly across the network.

## B.2. Details on ViT and BagNet

We provide here more details about the training of the larger architectures we used in Section 5 to perform experiments on the CIFAR-10 (Krizhevsky et al.) dataset. For both models, the data was processed through the same simple augmentation process consisting in random 32 × 32 cropping and horizontal flipping. Both training used a batch-size of 128.
**BagNet-17.** The first architecture we looked at was BagNet-17 (Brendel & Bethge, 2019). To establish a baseline, we used an SGD optimizer, with a momentum of 0.9, a learning rate of $10^{-1.5} \approx 0.032$, and a weight decay of $10^{-3}$. We also used a cosine scheduler, that reduces the learning rate to $10^{-5}$ over the course of the training. As mentioned in the main article, VJP approximation methods were employed on every linear (or equivalent, like $1 \times 1$ convolutions) layer, except the final classification layer and the initial input projection.

**Visual Transormer (ViT).** We also trained some ViT architecture (Dosovitskiy et al., 2021). Regarding the architecture details, we used an embedding dimension of 192, a MLP size of 1024, a network depth of 9, and 12 attention heads. Images were processed with a patch size of 4. Finally, to establish a baseline, we trained our model using AdamW (Loshchilov & Hutter, 2019) optimizer, with a learning rate of $3 \times 10^{-4}$ and a weight decay of 0.05. We also used a dropout parameter of 0.1, and a cosine scheduler with a 10 epochs warmup.
## C. Pseudo-code

### C.1. Practical considerations

We provide here the pseudocode for the different masking procedures implemented and discussed in the main paper. Note that, in most practical deep learning frameworks, operations are performed on tensors that represent batches of data. Moreover, in these frameworks, row vectors are typically used, in contrast to the usual theoretical setup that considers column vectors.

For example, in the usual mathematical setup, a linear layer is written as:

$$y = Wx + b, \qquad (52)$$

with $x \in \mathbb{R}^{d_{in}}$, $W \in \mathbb{R}^{d_{out} \times d_{in}}$, $b \in \mathbb{R}^{d_{out}}$, and $y \in \mathbb{R}^{d_{out}}$, where $d_{in}$ and $d_{out}$ denote the input and output dimensions of the layer, respectively.

In the practical setup, however, the linear layer is implemented as:

$$y = xW^{\top} + b, \qquad (53)$$

with $x \in \mathbb{R}^{B \times d_{in}}$, $W \in \mathbb{R}^{d_{out} \times d_{in}}$, $b \in \mathbb{R}^{d_{out}}$ (broadcasted across the batch dimension), and $y \in \mathbb{R}^{B \times d_{out}}$, where B denotes the batch size.

### C.2. Utils

This subsection shows code to find optimal probabilities $p_i^*$ for data-dependent sketches (Section 4.2), by solving the convex program stated in (23) as well as the algorithm that performs exact-r Bernoulli correlated sampling, as described in Theorem 3.3.

<caption>**Algorithm 1** Convex program solving for optimal probabilities (23)</caption>

**Require:** Importance weights $w_i > 0$
**Ensure:** Optimal sampling probabilities $p_i^* \in (0, 1]$ with $\sum_{i=1}^n p_i^* \le r$.
1: $t_i \leftarrow \sqrt{w_i}$ for all $i$
2: Sort $t$ in decreasing order: $t_{(1)} \ge \cdots \ge t_{(n)}$
3: Compute suffix sums $S_k \leftarrow \sum_{i=k}^n t_{(i)}$
4: **for** $k = 0, \dots, n-1$ **do**
5: &nbsp;&nbsp;&nbsp;&nbsp;remainder $\leftarrow r-k$
6: &nbsp;&nbsp;&nbsp;&nbsp;**if** remainder $\le 0$ **then**
7: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;**break**
8: &nbsp;&nbsp;&nbsp;&nbsp;**end if**
9: &nbsp;&nbsp;&nbsp;&nbsp;$$\sqrt{\lambda} \leftarrow S_{k+1}/\text{remainder}$$
10: **if** $k=0$ **or** $t_{(k)} \ge \sqrt{\lambda}$ **then**
11: &nbsp;&nbsp;&nbsp;&nbsp;**if** $t_{(k+1)} \le \sqrt{\lambda}$ **then**
12: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$k^* \leftarrow k$
13: &nbsp;&nbsp;&nbsp;&nbsp;**break**
14: &nbsp;&nbsp;&nbsp;&nbsp;**end if**
15: &nbsp;&nbsp;&nbsp;&nbsp;**end if**
16: **end for**
17: **for** $i = 1, \dots, n$ **do**
18: &nbsp;&nbsp;&nbsp;&nbsp;$p_{(i)}^* \leftarrow \min(1, t_{(i)}/\sqrt{\lambda})$
19: &nbsp;&nbsp;&nbsp;&nbsp;**end for**
20: Undo the sorting permutation
21: **return** $p^*$
**Algorithm 2 Correlated exact-r sampling**

**Require:** Probabilities $p \in (0, 1]^n$ with $\sum_{i=1}^n p_i = r$

**Ensure:** Exactly r distinct sampled indices

1. Compute cumulative sums:

2. $P_j \leftarrow \sum_{i=1}^{j} p_i$ for $j = 1, \dots, n$

3. $P_n \leftarrow r$ {numerical safety}

4. Sample $u \sim \text{Uniform}(0, 1]$

5. **for** $\ell = 0, \dots, r-1$ **do**

6. $t \leftarrow u + \ell$

7. Find smallest $j$ such that $P_j \ge t$

8. Select index $j$

9. **end for**

10. **return** the r selected indices

### C.3. Masking Methods

This subsection contains pseudocode for the three masking methods presented in Section 4.1.

**Algorithm 3 Per-Element Masking Backward Pass in a Linear Layer**

**Require:** Output gradient $G \triangleq \frac{\partial L}{\partial \mathbf{X}} \in \mathbb{R}^{B \times d_{\text{out}}}$

**Require:** Cached tensors $\mathbf{X} \in \mathbb{R}^{B \times d_{\text{in}}}$, $\mathbf{W} \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$, $\mathbf{b} \in \mathbb{R}^{d_{\text{out}}}$

**Require:** Sampling probability $p \in (0, 1]$

**Ensure:** Gradients $\frac{\partial L}{\partial \mathbf{X}}$, $\frac{\partial L}{\partial \mathbf{W}}$, $\frac{\partial L}{\partial \mathbf{b}}$

1. **// Generate element-wise masks**

2. $\mathbf{M}_W \sim \text{Bernoulli}(p) \in \{0, 1\}^{d_{\text{out}} \times d_{\text{in}}}$ (weight mask)

3. $\mathbf{M}_X \sim \text{Bernoulli}(p) \in \{0, 1\}^{B \times d_{\text{in}}}$ (input mask)

4. **// Gradient with respect to the input**

5. $\hat{\mathbf{W}} \leftarrow \mathbf{W} \odot \mathbf{M}_W$

6. $\frac{\partial L}{\partial \mathbf{X}} \leftarrow \frac{1}{p} \mathbf{G} \hat{\mathbf{W}}$

7. **// Gradient with respect to the weights**

8. $\hat{\mathbf{X}} \leftarrow \mathbf{X} \odot \mathbf{M}_X$

9. $\frac{\partial L}{\partial \mathbf{W}} \leftarrow \frac{1}{p} \mathbf{G}^{\top} \hat{\mathbf{X}}$

10. **// Gradient with respect to the bias**

11. $\frac{\partial L}{\partial \mathbf{b}} \leftarrow \sum_{i=1}^{B} \mathbf{G}[i, :]$

12. **return** $(\frac{\partial L}{\partial \mathbf{X}}, \frac{\partial L}{\partial \mathbf{W}}, \frac{\partial L}{\partial \mathbf{b}})$

Technical note regarding Algorithm 4: The mask generated is column-wise on the gradients, in a *practical* implementation, where notations are transposed compared to usual theoretical frameworks. Hence the mask zeroes-out *rows* of the gradients in mathematical notation... Which is equivalent to zeroing columns of the Jacobian.

### C.4. Sketching Methods

We now give the general sketching algorithm. Note that weight can be computed in various different ways.

This algorithm (Algorithm 6) can be used for all sketching methods that operate in the canonical basis (e.g. $\ell_1$, $\ell_2$, etc.), score computation just needs to be adapted.
**Algorithm 4** Per-Sample Masking Backward Pass in a Linear Layer

**Require:** Output gradient $G \triangleq \frac{\partial L}{\partial \mathbf{X}} \in \mathbb{R}^{B \times d_{\text{out}}}$

**Require:** Cached tensors $\mathbf{X} \in \mathbb{R}^{B \times d_{\text{in}}}, \mathbf{W} \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}, \mathbf{b} \in \mathbb{R}^{d_{\text{out}}}$

**Require:** Sampling probability $p \in (0, 1]$

**Ensure:** Gradients $\frac{\partial L}{\partial \mathbf{X}}, \frac{\partial L}{\partial \mathbf{W}}, \frac{\partial L}{\partial \mathbf{b}}$

1: // Generate row-wise mask (sample-level masking)

2: $\mathbf{m} \sim \text{Bernoulli}(p) \in \{0, 1\}^B$

3: $\mathbf{M}_G \leftarrow \mathbf{m} \mathbf{1}_{d_{\text{out}}}^{\top} \in \{0, 1\}^{B \times d_{\text{out}}}$

4: // Apply mask and rescale (unbiased estimator)

5: $\hat{\mathbf{G}} \leftarrow \frac{1}{p} (\mathbf{G} \odot \mathbf{M}_G)$

6: // Backward pass with masked gradient

7: $\frac{\partial L}{\partial \mathbf{X}} \leftarrow \hat{\mathbf{G}} \mathbf{W}$

8: $\frac{\partial L}{\partial \mathbf{W}} \leftarrow \hat{\mathbf{G}}^{\top} \mathbf{X}$

9: $\frac{\partial L}{\partial \mathbf{b}} \leftarrow \sum_{i=1}^{B} \hat{\mathbf{G}}[i, :]$

10: **return** $(\frac{\partial L}{\partial \mathbf{X}}, \frac{\partial L}{\partial \mathbf{W}}, \frac{\partial L}{\partial \mathbf{b}})$

**Algorithm 5** Per-Column Masking Backward Pass in a Linear Layer

**Require:** Output gradient $G \triangleq \frac{\partial L}{\partial \mathbf{Y}} \in \mathbb{R}^{B \times d_{\text{out}}}$

**Require:** Cached tensors $\mathbf{X} \in \mathbb{R}^{B \times d_{\text{in}}}, \mathbf{W} \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}, \mathbf{b} \in \mathbb{R}^{d_{\text{out}}}$

**Require:** Sparsity probability $p \in (0, 1]$

**Ensure:** Gradients $\frac{\partial L}{\partial \mathbf{X}}, \frac{\partial L}{\partial \mathbf{W}}, \frac{\partial L}{\partial \mathbf{b}}$

1: // Generate column-wise mask

2: $\mathbf{m} \sim \text{Bernoulli}(p) \in \{0, 1\}^{d_{\text{out}}}$

3: $\mathbf{M}_G \leftarrow \mathbf{1}_B \mathbf{m}^{\top} \in \{0, 1\}^{B \times d_{\text{out}}}$

4: // Apply mask and rescale (unbiased estimator)

5: $\hat{\mathbf{G}} \leftarrow \frac{1}{p} (\mathbf{G} \odot \mathbf{M}_G)$

6: // Backward pass with masked gradient

7: $\frac{\partial L}{\partial \mathbf{X}} \leftarrow \hat{\mathbf{G}} \mathbf{W}$

8: $\frac{\partial L}{\partial \mathbf{W}} \leftarrow \hat{\mathbf{G}}^{\top} \mathbf{X}$

9: $\frac{\partial L}{\partial \mathbf{b}} \leftarrow \sum_{i=1}^{B} \hat{\mathbf{G}}[i, :]$

10: **return** $(\frac{\partial L}{\partial \mathbf{X}}, \frac{\partial L}{\partial \mathbf{W}}, \frac{\partial L}{\partial \mathbf{b}})$
**Algorithm 6 Coordinate Sketch Backward Pass for a Linear Layer (Example with L1 Proxy)**

**Require:** Gradient w.r.t. output (matrix) $\mathbf{G} \triangleq \frac{\partial \mathcal{L}}{\partial \mathbf{Y}} \in \mathbb{R}^{B \times d_{\text{out}}}$

**Require:** Cached tensors $\mathbf{X} \in \mathbb{R}^{B \times d_{\text{in}}}$, $\mathbf{W} \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$, (optional) $\mathbf{b} \in \mathbb{R}^{d_{\text{out}}}$

**Require:** Target number of selected columns $r$

**Require:** Subroutines: PSTARFROMSCORES(·) (Alg. 1) and CORRELATEDEXACTRSAMPLE(·) (Alg. 2)

**Ensure:** Gradients $\frac{\partial \mathcal{L}}{\partial \mathbf{X}}$, $\frac{\partial \mathcal{L}}{\partial \mathbf{W}}$, $\frac{\partial \mathcal{L}}{\partial \mathbf{b}}$

1: // Compute L1 based importance weights per output column

2: for $j = 1, \dots, d_{\text{out}}$ do

3: $s_j \leftarrow \|\mathbf{G}[:, j]\|_1^2$

4: end for

5: $\mathbf{s} \leftarrow (s_1, \dots, s_{d_{\text{out}}})$

6: // Compute optimal marginal probabilities and sample exactly $r$ columns

7: $\mathbf{p}^* \leftarrow \text{PSTARFROMWEIGHTS}(\mathbf{s}, r)$

8: $\mathcal{I} \leftarrow \text{CORRELATEDEXACTRSAMPLE}(\mathbf{p}^*, r)$

($|\mathcal{I}| = r$)

9: // Form unbiased masked estimator of the output gradient

10: $\mathbf{m} \leftarrow \mathbf{0} \in \{0, 1\}^{d_{\text{out}}}$

11: for each $j \in \mathcal{I}$ do

12: $m_j \leftarrow 1$

13: end for

14: $\hat{\mathbf{G}} \leftarrow \mathbf{G} \odot (\mathbf{1}_B \mathbf{m}^{\top})$

15: for each $j \in \mathcal{I}$ do

16: $\hat{\mathbf{G}}[:, j] \leftarrow \hat{\mathbf{G}}[:, j]/p_j^*$

// Scaling to ensure unbiasedness

17: end for

18: // Backward pass using the sketched output gradient

19: $\frac{\partial \mathcal{L}}{\partial \mathbf{X}} \leftarrow \hat{\mathbf{G}} \mathbf{W}$

20: $\frac{\partial \mathcal{L}}{\partial \mathbf{W}} \leftarrow \hat{\mathbf{G}}^{\top} \mathbf{X}$

21: if $\mathbf{b}$ is present then

22: $\frac{\partial \mathcal{L}}{\partial \mathbf{b}} \leftarrow \sum_{i=1}^{B} \hat{\mathbf{G}}[i, :]$

23: end if

24: return $(\frac{\partial \mathcal{L}}{\partial \mathbf{X}}, \frac{\partial \mathcal{L}}{\partial \mathbf{W}}, \frac{\partial \mathcal{L}}{\partial \mathbf{b}})$