💻 computer science

Higher Order Automatic Differentiation of Higher Order Functions

This paper presents semantic correctness proofs for forward-mode automatic differentiation in a higher-order language with algebraic data types by characterizing the method as a unique structure-preserving macro and establishing its validity through a gluing construction on diffeological spaces that extends to higher-order derivatives via Taylor approximation.

Original authors: Mathieu Huot, Sam Staton, Matthijs Vákár

Published 2026-05-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mathieu Huot, Sam Staton, Matthijs Vákár

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a complex recipe for a cake. This recipe doesn't just list ingredients; it includes instructions for other recipes (like "make the frosting first," then "use that frosting here"). In the world of computer science, this is called a higher-order function: a function that takes other functions as ingredients or creates new functions as results.

Now, imagine you want to know exactly how changing one tiny ingredient (like adding a pinch more sugar) changes the final taste of the cake. In math, this is called finding a derivative. In the world of machine learning and AI, this process is called Automatic Differentiation (AD). It's the engine that teaches computers how to learn by tweaking their internal settings to minimize errors.

This paper tackles a very tricky problem: How do we mathematically prove that our computer code for calculating these "taste changes" is correct, even when the recipe itself is made of other recipes?

Here is a breakdown of their solution using simple analogies:

1. The Problem: The "Black Box" of Higher-Order Functions

Usually, if you have a simple function (like $f(x) = x^2$ ), calculus gives us a clear rule to find its slope (derivative). But when you have a function that takes another function as an input (like a "recipe maker"), standard calculus gets confused. It's like trying to measure the slope of a machine that builds other machines. There isn't a single, agreed-upon mathematical rule for this in traditional geometry.

The authors ask: If we write a program that automatically calculates these complex slopes, how do we know it's not lying to us?

2. The Solution: A New Kind of Map (Diffeological Spaces)

To solve this, the authors needed a new way to visualize the "space" where these programs live.

Old Map (Manifolds): Think of this as a standard map of the Earth. It's great for smooth hills and valleys, but it breaks down if you try to map a "space of all possible maps." It can't handle the idea of a function being an object you can hold and manipulate.
New Map (Diffeological Spaces): The authors use a concept called diffeological spaces. Imagine this as a "super-map" that doesn't just look at points on a surface, but looks at all the possible paths (curves) you could draw on that surface.
- If you can draw a smooth path on a shape, that shape is "smooth."
- This approach is flexible enough to handle not just simple numbers, but also lists, choices (like "if this, then that"), and even functions that take other functions as arguments.

3. The Method: The "Dual Number" Translator

The paper describes a specific tool called a macro. Think of this macro as a translator that takes your original program and rewrites it.

Original Program: "Calculate the cost of this neural network."
Translated Program: "Calculate the cost and how the cost changes if you wiggle every single number slightly."

The authors prove that this translator works correctly by using a technique called Logical Relations.

The Analogy: Imagine you have a "Shadow World" (the original program) and a "Double-Shadow World" (the program with derivatives). The authors create a rulebook (a relation) that says: "For every move you make in the Shadow World, there must be a corresponding, mathematically correct move in the Double-Shadow World."
They prove that no matter how complex the nesting of functions gets, the translator always keeps the shadows aligned. If the original program is smooth, the translated program correctly calculates the smooth changes.

4. The "Gluing" Trick

To make this proof rigorous, they use a mathematical construction called Gluing.

The Analogy: Imagine you are building a 3D model out of flat pieces of paper. You have the "original" paper and the "derivative" paper. The "gluing" is the tape that holds them together, ensuring that the derivative paper is always attached to the original paper in the correct way.
This "glued" space allows them to treat the original function and its derivative as a single, unified object. They show that their translator is the only way to build this glue that preserves the structure of the language.

5. The Surprise: Derivatives Aren't Always Unique

One of the most interesting findings is that for these complex "function-building" machines, there isn't always just one correct derivative.

The Analogy: Imagine you are driving a car. The "derivative" is your speed. If you are driving on a straight road, your speed is clear. But if you are driving a car that builds other cars, there might be two different ways to define "speed" that both work perfectly fine for the final result, even if they look different on the dashboard.
The authors show that their method picks one specific, simple, and efficient version of this derivative. It doesn't matter which "valid" version you pick, as long as it correctly calculates the changes for the final, simple numbers (the first-order functions) that actually get used in the real world.

Summary

In short, this paper builds a mathematical safety net for advanced automatic differentiation.

They created a new type of mathematical space (diffeological spaces) that can hold complex, nested functions.
They proved that their code translator (the macro) correctly calculates derivatives for these complex functions by showing it aligns perfectly with the rules of this new space.
They demonstrated that while there might be multiple ways to define a derivative for a "function-maker," their method is a valid, consistent, and correct choice that ensures the final calculations for AI and machine learning are accurate.

They didn't invent a new way to train AI, but they provided the proof that the current ways of training AI using these complex tools are mathematically sound.

Technical Summary: Higher-Order Automatic Differentiation of Higher-Order Functions

1. Problem Statement

Automatic Differentiation (AD) is a fundamental technique in machine learning and scientific computing, used to compute derivatives of functions defined by programs. While AD is well-understood for first-order functions (functions mapping real numbers to real numbers), a significant gap exists in the formal theory when dealing with higher-order functions (functions that take other functions as arguments or return them) within a higher-order programming language.

The core challenge addressed in this paper is twofold:

Semantic Gap: Traditional differential geometry (based on manifolds) does not naturally support spaces of functions, making it difficult to define a denotational semantics for higher-order differentiable programs.
Correctness of Higher-Order AD: Existing AD methods (specifically forward-mode) are often implemented as syntactic macros that transform code. While these macros work compositionally, there has been no rigorous semantic proof establishing their correctness for languages that include higher-order functions and algebraic data types (like lists and variants). Specifically, it was unclear what the "derivative" of a higher-order function should mean semantically, and whether the syntactic transformation preserves the correct mathematical properties.

2. Methodology

The authors employ a denotational semantics approach combined with logical relations and categorical gluing to formalize and prove the correctness of higher-order AD.

2.1. Semantic Model: Diffeological Spaces

To resolve the lack of function spaces in standard manifold theory, the paper utilizes diffeological spaces.

Definition: A diffeological space is a set equipped with a chosen set of "plots" (smooth maps from open subsets of $\mathbb{R}^n$ to the set) satisfying specific axioms (constant maps, smooth pre-composition, and gluing).
Advantage: This framework supports function spaces (the set of smooth maps between two diffeological spaces is itself a diffeological space) and inductive types (like lists), allowing for a compositional interpretation of the entire language.
Smoothness: The authors prove that every definable expression in their language interprets to a smooth function within this model.

2.2. Syntactic Translation: Taylor-Mode AD Macros

The paper defines a syntactic macro translation, denoted $\vec{D}_{(k,R)}$ , which transforms a program $t$ into a new program $\vec{D}(t)$ .

Higher-Order Derivatives: Instead of just computing first derivatives (dual numbers), the macro computes $(k, R)$ -Taylor approximations. This involves tracking partial derivatives up to order $R$ in $k$ dimensions.
Mechanism: The translation is defined inductively on the syntax of the language (variables, operations, $\lambda$ -abstractions, applications, tuples, pattern matching, and inductive types like lists).
Basic Operations: For primitive operations (e.g., $+$ , $\times$ , $\sigma$ ), the macro replaces them with their $(k, R)$ -Taylor representations (derived from the Faà di Bruno formula), which compute the value and all partial derivatives up to order $R$ .

2.3. Proof Technique: Logical Relations and Gluing

To prove correctness, the authors use a logical relations argument, which they later reframe using categorical gluing.

Logical Relations: They define a relation $S_\tau$ between the semantic interpretation of a type $\tau$ and the interpretation of its transformed type $\vec{D}(\tau)$ . This relation links a smooth function $f$ with its Taylor representation.
Fundamental Lemma: They prove that if inputs are related by $S$ , then the output of the program and the output of the transformed program are also related by $S$ . This establishes that the macro correctly computes derivatives.
Categorical Gluing: The logical relation is formalized as a category of "glued spaces" ( $Gl_k$ ), where objects are triples $(X, X', S)$ representing a space, its derivative space, and the relation between them. The AD macro is shown to be a unique structure-preserving functor in this setting.

3. Key Contributions

Denotational Semantics for Higher-Order AD: The paper provides the first denotational semantics for a higher-order language with algebraic data types using diffeological spaces, proving that all definable functions are smooth.
Correctness Proof for Higher-Order AD: It presents a rigorous proof that the forward-mode AD macro $\vec{D}_{(k,R)}$ correctly computes the $(k, R)$ -Taylor representation of any definable function, even when the function involves higher-order constructs.
Categorical Characterization: The authors characterize the AD macro as the unique structure-preserving functor (a "gluing" construction) given a choice of derivatives for basic operations. This provides a universal property for AD.
Extension to Inductive Types: The methodology is extended to cover variant types and inductive types (specifically lists), demonstrating that AD can be applied to programs manipulating complex data structures.
Analysis of Non-Uniqueness: The paper investigates the derivatives of higher-order functions and demonstrates that, unlike first-order functions, semantic derivatives of higher-order functions are not necessarily unique. The syntactic macro chooses one valid derivative, but others exist that are equally correct with respect to the semantics.

4. Results

Theorem 4.3 (Limited Correctness): For any term $t$ of type real (even if defined using higher-order functions), the interpretation of the transformed term $\vec{D}(t)$ is the $(k, R)$ -Taylor representation of the interpretation of $t$ .
Theorem 6.6 (Full Correctness): This result is generalized to all first-order types (e.g., list(real), real -> real). The syntactic translation coincides with the $(k, R)$ -jet bundle functor on manifolds, modulo canonical isomorphisms.
Non-Uniqueness (Section 7): The authors show via a counterexample (the evaluation function) that multiple distinct semantic derivatives can exist for a higher-order function. The AD macro selects one specific derivative, but the semantics does not force a unique choice, provided the chosen derivative yields correct results for first-order functions.

5. Significance and Claims

The paper claims to achieve two primary goals:

Compositional Understanding: It offers a compositional, denotational understanding of differentiable programming, bridging the gap between the syntactic operations of AD and the mathematical theory of differential geometry.
Formal Correctness: It provides a formal explanation of why higher-order AD works, moving beyond empirical observation to a semantic proof based on logical relations and gluing constructions.

The authors emphasize that their approach is modest regarding the uniqueness of higher-order derivatives. They argue that in programming practice, the specific choice of a derivative for a higher-order function is less critical than ensuring that the derivative correctly propagates to the first-order functions that are ultimately computed. The fact that the semantics allows for non-unique derivatives is framed as a feature that models the pragmatic flexibility of programming, rather than a flaw.

The work extends previous research on first-order AD (e.g., [HSV20a]) to higher-order derivatives and higher-order functions, providing a theoretical foundation for tools that compute higher-order derivatives (Taylor-mode AD) in functional languages. It does not claim to solve implementation efficiency issues (like reverse-mode optimization) but establishes the semantic correctness of the forward-mode approach.