Generalizing Stochastic Smoothing for Differentiation and Gradient Estimation

TL;DR

This paper introduces a generalized stochastic smoothing framework for gradient estimation of non-differentiable functions, reducing assumptions and improving variance reduction techniques, with applications in sorting, graph algorithms, rendering, and simulations.

Contribution

It develops a theory for stochastic smoothing without requiring differentiable densities or full support, and proposes variance reduction strategies for diverse applications.

Findings

Benchmarking of 6 distributions and 24 variance reduction strategies.

Effective gradient estimation for sorting, ranking, and graph shortest-paths.

Improved performance in differentiable rendering and cryo-ET simulations.

Abstract

We deal with the problem of gradient estimation for stochastic differentiable relaxations of algorithms, operators, simulators, and other non-differentiable functions. Stochastic smoothing conventionally perturbs the input of a non-differentiable function with a differentiable density distribution with full support, smoothing it and enabling gradient estimation. Our theory starts at first principles to derive stochastic smoothing with reduced assumptions, without requiring a differentiable density nor full support, and we present a general framework for relaxation and gradient estimation of non-differentiable black-box functions $f:\mathbb{R}^n\to\mathbb{R}^m$. We develop variance reduction for gradient estimation from 3 orthogonal perspectives. Empirically, we benchmark 6 distributions and up to 24 variance reduction strategies for differentiable sorting and ranking, differentiable shortest-paths on graphs, differentiable rendering for pose estimation, as well as differentiable cryo-ET simulations.