Reproduce “Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagation” on OpenResearch

About this paper

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

The authors introduce unbiased randomized sketches for neural-network backpropagation, retaining and rescaling selected gradient directions so each backward signal remains unbiased. The method analyzes how sketching variance propagates through layers and derives information-guided designs. On MNIST and CIFAR-10 image-classification models, informed sketches generally preserve accuracy better than uniform masks at lower tested sampling budgets; runtime savings are not measured.

Key insights

• Using conditionally unbiased gradient approximations in Stochastic Gradient Descent only affects convergence through increased variance, allowing for a principled variance-efficiency trade-off.
• Data-dependent and Jacobian-aware sketching methods consistently outperform uniform masking strategies by leveraging local information to prioritize gradient components.
• While theoretically optimal rank-constrained sketches are derived, simpler and computationally less expensive methods like Diagonal Sketching prove highly effective in practice on larger models.

Summary by alphaXiv.