Smoothing Methods for Automatic Differentiation Across Conditional Branches

TL;DR

This paper introduces a novel approach combining smooth interpretation with automatic differentiation to efficiently compute gradients of programs with discontinuous control flow, enabling improved optimization in complex models.

Contribution

It presents a new method integrating smooth interpretation with AD and a Monte Carlo estimator to handle discontinuities in program optimization tasks.

Findings

Monte Carlo estimator avoids assumptions of SI, providing robust gradient estimates.

The combined approach outperforms existing methods in high-dimensional problems.

The method enables direct gradient-based optimization of discontinuous programs.

Abstract

Programs involving discontinuities introduced by control flow constructs such as conditional branches pose challenges to mathematical optimization methods that assume a degree of smoothness in the objective function's response surface. Smooth interpretation (SI) is a form of abstract interpretation that approximates the convolution of a program's output with a Gaussian kernel, thus smoothing its output in a principled manner. Here, we combine SI with automatic differentiation (AD) to efficiently compute gradients of smoothed programs. In contrast to AD across a regular program execution, these gradients also capture the effects of alternative control flow paths. The combination of SI with AD enables the direct gradient-based parameter synthesis for branching programs, allowing for instance the calibration of simulation models or their combination with neural network models in machine learning pipelines. We detail the effects of the approximations made for tractability in SI and propose a novel Monte Carlo estimator that avoids the underlying assumptions by estimating the smoothed programs' gradients through a combination of AD and sampling. Using DiscoGrad, our tool for automatically translating simple C++ programs to a smooth differentiable form, we perform an extensive evaluation. We compare the combination of SI with AD and our Monte Carlo estimator to existing gradient-free and stochastic methods on four non-trivial and originally discontinuous problems ranging from classical simulation-based optimization to neural network-driven control. While the optimization progress with the SI-based estimator depends on the complexity of the program's control flow, our Monte Carlo estimator is competitive in all problems, exhibiting the fastest convergence by a substantial margin in our highest-dimensional problem.