Stabilizing Backpropagation Through Time to Learn Complex Physics

TL;DR

This paper introduces a modified backpropagation through time method tailored for physics simulations, which stabilizes gradient flow and improves learning in complex, long-horizon control tasks.

Contribution

It proposes a novel stabilization technique for backpropagation through time that maintains minima positions and enhances training stability in physics-based recurrent models.

Findings

Improved control performance on complex physics tasks

Stabilized gradient flow reduces exploding and vanishing updates

Method scales well with task complexity

Abstract

Of all the vector fields surrounding the minima of recurrent learning setups, the gradient field with its exploding and vanishing updates appears a poor choice for optimization, offering little beyond efficient computability. We seek to improve this suboptimal practice in the context of physics simulations, where backpropagating feedback through many unrolled time steps is considered crucial to acquiring temporally coherent behavior. The alternative vector field we propose follows from two principles: physics simulators, unlike neural networks, have a balanced gradient flow, and certain modifications to the backpropagation pass leave the positions of the original minima unchanged. As any modification of backpropagation decouples forward and backward pass, the rotation-free character of the gradient field is lost. Therefore, we discuss the negative implications of using such a rotational vector field for optimization and how to counteract them. Our final procedure is easily implementable via a sequence of gradient stopping and component-wise comparison operations, which do not negatively affect scalability. Our experiments on three control problems show that especially as we increase the complexity of each task, the unbalanced updates from the gradient can no longer provide the precise control signals necessary while our method still solves the tasks. Our code can be found at https://github.com/tum-pbs/StableBPTT.