Functorial String Diagrams for Reverse-Mode Automatic Differentiation (2107.13433v1)

Abstract: We enhance the calculus of string diagrams for monoidal categories with hierarchical features in order to capture closed monoidal (and cartesian closed) structure. Using this new syntax we formulate an automatic differentiation algorithm for (applied) simply typed lambda calculus in the style of [Pearlmutter and Siskind 2008] and we prove for the first time its soundness. To give an efficient yet principled implementation of the AD algorithm we define a sound and complete representation of hierarchical string diagrams as a class of hierarchical hypergraphs we call hypernets.

• The paper introduces a novel reverse-mode AD algorithm using enriched, hierarchical string diagrams to efficiently represent and differentiate higher-order functions.
• The paper proves the soundness of the algorithm by mapping it to reverse derivative categories, ensuring rigorous formal correctness.
• The paper presents practical implications by implementing hierarchical hypergraphs (hypernets) that enable efficient automated graph rewriting in AD.

Functorial String Diagrams for Reverse-Mode Automatic Differentiation

The paper presents an exploration into the integration of functorial string diagrams with reverse-mode automatic differentiation (AD), contributing to the formal understanding and practical implementation of AD in higher-order programming languages. The approach leverages the graphical syntax of string diagrams, enriched with hierarchical features, to model complex categorical structures and simplify the conceptual and formal analysis of AD algorithms.

String diagrams provide a powerful alternative to textual formalisms, offering a visual syntax for expressing morphisms in monoidal categories. The authors extend this graphical notation to support closed monoidal (and cartesian closed) structures, enabling the representation of the simply typed lambda calculus with explicit substitutions—a necessary foundation for expressing higher-order functions.

Contributions

1. Hierarchical String Diagrams and Automatic Differentiation: The primary contribution lies in formulating an AD algorithm using these enriched string diagrams. By structuring the diagrammatic calculus to support hierarchical aspects like abstraction, the paper enables an efficient representation and manipulation of higher-order functions. This is pivotal in reverse-mode AD, where the gradient of functions involving closures must be accurate and computationally feasible.
2. Soundness of the AD Algorithm: The authors address a gap in the literature by proving the soundness of the AD algorithm formulated within this new syntactic framework. This involves the complex task of demonstrating that the algorithm behaves correctly with higher-order constructs, a challenge accentuated by the need for effective track-back mechanisms in reverse-mode AD.
3. Hierarchical Hypergraphs (Hypernets): To achieve a concrete implementation of the AD algorithm, the paper introduces hypernets—a class of hierarchical hypergraphs that represent hierarchical string diagrams. These structures enable efficient graph rewriting operations akin to double-pushout (DPO) rewriting in categorical graph theory, facilitating the automated transformation processes central to AD.

Technical Overview

The algorithm's soundness is established through a mapping to reverse derivative categories, which are shown to align with the intended calculus-based notions of differentiation. A crucial insight is the formulation of the AD transformation rules and their decomposition into complementary forward and reverse passes that reflect common AD practice but are here formalized within a graphical syntax.

Practical and Theoretical Implications

The results have significant implications for both the theory and practice of automatic differentiation in functional programming languages:

• Theoretical Implications: The insights provided into the relationship between string diagrams and categorical structures may inform new developments in understanding computational effects in category theory and their applications in programming language design.
• Practical Implications: From a practical perspective, these findings could be embedded into functional language compilers, enhancing the robust support for AD in environments using higher-order constructs, such as those found in ML or Haskell.

Future Directions

The paper hints at further opportunities to explore implementations, particularly considering effects, broader computational models, and more efficient evaluation strategies in non-strict languages. Additionally, the connection between string diagrams and other graphical formalisms like proof nets suggests further cross-pollination of ideas might result in new insights into the handling of state and coinductive structures within differentiable programming.

Overall, this research presents a compelling case for adopting graphical techniques in the design and analysis of complex algorithms in computer science, positioning string diagrams as both a theoretical and a practical tool in the evolving landscape of differentiable programming.