Differentiable Causal Computations via Delayed Trace

Abstract

We investigate causal computations, which take sequences of inputs to sequences of outputs such that the nth output depends on the first n inputs only. We model these in category theory via a construction taking a Cartesian category \mathbbC to another category St(\mathbbC) with a novel trace-like operation called “delayed trace”, which misses yanking and dinaturality axioms of the usual trace. The delayed trace operation provides a feedback mechanism in St(\mathbbC) with an implicit guardedness guarantee. When \mathbbC is equipped with a Cartesian differential operator, we construct a differential operator for St (\mathbbC) using an abstract version of backpropagation through time, a technique from machine learning based on unrolling of functions. This obtains a swath of properties for backpropagation through time, including a chain rule and Schwartz theorem. Our differential operator is also able to compute the derivative of a stateful network without requiring the network to be unrolled.