PaperlayerView exact revision on arXiv

Source-linked AI summary

The simple essence of automatic differentiation

Conal Elliott

arXiv:1804.00746v4cs.PL

TL;DR

Common RAD implementations are complex and stateful, complicating understanding and parallel execution. This paper derives a generalized AD algorithm from a simple specification and specializes it through different derivative representations. The resulting RAD algorithms are simpler than previously known, while detailed performance analysis remains future work.

  • Problem

    Common RAD algorithms such as backpropagation are complex and stateful, hindering understanding, improvement, and parallel execution.

    Source: §1 · p. 1

  • Method

    The paper specifies AD using functoriality, derives an algorithm from that specification, and specializes it using categorical constructions and different derivative representations.

    Source: §1 · p. 2, §17 · p. 25

  • Results

    The specialized RAD implementations are far simpler than previously known and use no graphs, tapes, variables, partial derivatives, or mutation.

    Source: §1 · p. 2, §17 · p. 25

  • Takeaways & Limitations

    The algorithms are correct by construction, naturally parallel, and usable through an AD-agnostic compiler plugin without new data types or programming style.

    Source: §1 · p. 2, §6 · p. 14

    Detailed performance analysis against backpropagation and other conventional AD algorithms remains future work.

    Source: §16 · p. 25, §17 · p. 25

Abstract

from arXiv · show

Automatic differentiation (AD) in reverse mode (RAD) is a central component of deep learning and other uses of large-scale optimization. Commonly used RAD algorithms such as backpropagation, however, are complex and stateful, hindering deep understanding, improvement, and parallel execution. This paper develops a simple, generalized AD algorithm calculated from a simple, natural specification. The general algorithm is then specialized by varying the representation of derivatives. In particular, applying well-known constructions to a naive representation yields two RAD algorithms that are far simpler than previously known. In contrast to commonly used RAD implementations, the algorithms defined here involve no graphs, tapes, variables, partial derivatives, or mutation. They are inherently parallel-friendly, correct by construction, and usable directly from an existing programming language with no need for new data types or programming style, thanks to use of an AD-agnostic compiler plugin.

1 Introduction

Reverse-mode AD is valuable for high-dimensional optimization but is commonly implemented with complex, stateful machinery that complicates understanding and parallel execution. The paper builds AD compositionally from a simple specification, using derivative representations and standard composition rules to derive simpler algorithms.

1 Introduction

  • RAD and backpropagation handle high-dimensional, scalar-output optimization efficiently, but common implementations use mutation, graphs, and tapes that add overhead and hinder parallel execution.The paper proposes differentiable functional programming as an alternative to graph-centered machine-learning workflows, using a Haskell compiler plugin.Source: §1 · pp. 1–2

2 What’s a Derivative?

  • A derivative is a linear map, a general representation that avoids increasingly complex derivative representations and chain rules for higher-dimensional and higher-order cases.Partial derivatives and Jacobians are component-based representations; the paper instead treats derivatives as linear maps from input to output spaces.Source: §2 · p. 3

3.1 Sequential Composition

Theorem 1 (compose/“chain” rule)

  • Pairing a function with its derivative makes differentiation compositional, but computing both independently repeats f a and can make nested compositions impractically expensive.The improved D+ combines them as D+ f a = (f a, D f a), avoiding that redundant work.Source: Theorem 1 (compose/“chain” rule) · p. 4

3.2 Parallel Composition

  • For functions applied independently to paired inputs, the derivative of their parallel composition is the parallel composition of their derivatives.D (f × g) (a, b) = D f a × D g b; the two functions have no data dependencies.Source: §3.2 · p. 4, Theorem 2 (cross rule) · p. 5

Theorem 2 (cross rule)

  • D+ composes across parallel operations by combining the outputs and pairing the two derivative maps.The construction is valid because parallel composition preserves linearity.Source: Theorem 2 (cross rule) · p. 5

3.3 Linear Functions

  • A linear function is its own derivative everywhere, so D+ f can return the function’s output together with f itself.This expresses linear-map derivatives directly, without requiring a matrix representation.Source: §3.3 · p. 5

4 Putting the Pieces Together

This section puts compositional differentiation together into a correct-by-construction implementation, deriving familiar categorical structures for expressing and checking programs translated from conventional functional languages. It also develops fork/unfork and join/unjoin operations and extends structure-preserving differentiation to nonlinear numeric primitives, separating primitive differentiation from its context.

4.1 Categories

  • Representing D+ as a functor into a category of derivative pairs yields a Category instance; hiding D’s constructor lets category laws follow mechanically from functor laws.The proofs use only functor laws and the corresponding laws for functions, so they need not be repeated when category instances are specified via a functor.Source: §4.1 · pp. 6, 8

4.2 Monoidal Categories

  • Specifying the derivative representation as a monoidal functor yields a correct-by-construction Monoidal instance for parallel composition over products.The resulting operation combines two morphisms by applying them to paired inputs and pairing their outputs and derivatives.Source: §4.2 · pp. 8–9

4.3 Cartesian Categories

  • A cartesian-functor specification yields a correct-by-construction Cartesian instance, adding projection and duplication operations to the derivative representation.The instance defines exl, exr, and dup by applying linearD to each corresponding function.Source: §4.3 · p. 9

4.4 Cocartesian Categories

  • Because the function category lacks the needed Cocartesian instance, the paper uses additive functions, with coproducts coinciding with products, to represent derivatives in a Cocartesian category.The paper identifies this structure as a biproduct category.Source: §4.4 · p. 10

4.5 Derived Operations

  • △ combines two morphisms sharing a domain via duplication, while ▽ combines two sharing a codomain via jamming; △ is particularly useful for translating the λ-calculus to categorical form, and both are invertible in uncurried form.Source: §4.5 · pp. 10–11
  • The derived operations pair fork with unfork and join with unjoin, defining each operation alongside its inverse-style decomposition.fork(f, g) = f △ g and unfork(h) = (exl ◦ h, exr ◦ h); join(f, g) = f ▽ g and unjoin(h) = (h ◦ inl, h ◦ inr).Source: §4.5 · p. 11

4.6 Numeric Operations

  • Numeric interfaces generalize operations beyond functions, allowing non-function operations while defining numeric primitives through structure-preserving homomorphisms.The interfaces roughly follow Haskell’s numeric type-class hierarchy and un-curry operations to demand less of supporting categories.Source: §4.6 · p. 11
  • Conventional differentiation rules implicitly mix each numeric primitive with the chain rule and leave unclear whether their arguments are functions or non-functions.The paper instead differentiates primitive operations in isolation, leaving contextual composition to the chain rule.Source: §4.6 · p. 11
  • Multiplication’s derivative maps (da, db) to da · b + a · db, a linear map represented using generalized scalar multiplication and the derived ▽ operation.The product rule is cited to [Spivak, 1965, Theorem 2-3 (2)].Source: §4.6 · p. 12
  • Preserving NumCat structure yields a correct-by-construction differentiable-functions instance, with negation and addition handled linearly and multiplication using the product rule.The passage also gives analogous derivative definitions for sin, cos, and exp.Source: §4.6 · p. 12

5 Examples

The examples apply differentiation to two functions, magSqr and cosSinProd, then visualize their undifferentiated and differentiated forms. The derivative graphs show linear functions, with work shared between the primal result and derivative in cosSinProd.

5 Examples

  • magSqr squares and adds the two components of its input, while cosSinProd multiplies its inputs and returns the cosine and sine of their product.Source: §5 · p. 12
  • Figures 2 and 3 visualize the example functions before differentiation; Figures 4 and 5 visualize their differentiable versions.The paper describes the graphs as visualizations only; they play no role in the programming interface or differentiation implementation.Source: §5 · p. 13, Figure 2 · p. 13, Figure 3 · p. 13
  • The derivatives are linear functions, and the cosSinProd derivative shares work between the function’s result and its derivative.The paper identifies the function’s result as the “primal.”Source: §5 · p. 13

6 Programming as Defining and Solving Algebra Problems

The section presents programming as an algebra problem: choose a compositional representation and specification, then solve for an implementation whose laws follow from its homomorphic structure.

6 Programming as Defining and Solving Algebra Problems

  • The section frames this general procedure as a recipe for defining and solving algebra problems.Source: §6 · p. 13
  • The process starts with a possibly expensive or non-computable specification and encodes the desired result in a new data type, here combining a function with its derivative.It then tests whether conversion from simpler forms is compositional with respect to familiar algebraic abstractions.Source: §6 · p. 13
  • If the initial specification is not compositional, examine its failure and iteratively augment the specification until it becomes compositional.Unadorned differentiation D is given as an example of a failure of compositionality.Source: §6 · p. 13
  • The algebra problem is to make conversion into the chosen representation a homomorphism for well-understood abstractions, such as a category or functor.The problem is solved using compositionality properties.Source: §6 · pp. 13–14
  • Keeping the new data type abstract ensures that the solution satisfies the required laws, by virtue of the homomorphic specification.Source: §6 · p. 14
  • The recipe yields a computable alternative to a potentially non-computable specification when programs use the chosen algebraic vocabulary, which can then be reinterpreted in the new data type.Automatically converting functional programs into that vocabulary lets them be reinterpreted to compute the specification.Source: §6 · p. 14

7 Generalizing Automatic Differentiation

The paper generalizes AD by replacing linear maps with other categories, since operations on differentiable functions are defined through corresponding operations on linear maps.

7 Generalizing Automatic Differentiation

  • Replacing linear maps with a broader variety of underlying categories yields a greatly generalized notion of AD.Operations on differentiable functions are defined using the corresponding operations on linear maps.Source: §7 · p. 14
  • The generalized definitions parameterize the new category with a category k that replaces linear maps in D.Source: §7 · p. 14
  • The generalized linearD function takes two arrows, and the functionality required of the underlying category is made explicit.Source: §7 · p. 14
  • The object constraint for Dk combines Additive, needed for the Cocartesian instance, with Obj k, needed for all instances.Source: §7 · p. 14

8 Matrices

The section relates generalized linear-map operations to matrix shapes and gives categorical definitions for the derivative-augmented function D.

8 Matrices

  • Horizontal juxtaposition adds matrix widths while preserving height, and vertical juxtaposition adds heights while preserving width.Source: §8 · p. 14
  • The three shape constraints establish and preserve rectangularity.Source: §8 · p. 14
  • Three operations in the generalized linear-map vocabulary correspond to the possible forms of nonempty matrices, with types encoding width and height constraints.For matrices representing linear maps, domain and codomain types follow from matrix width and height, respectively, under the stated column-vector convention.Source: §8 · p. 14
  • The Dk category requires Additive and Obj k; its identity and composition are defined through linearD and the corresponding linear-map operations.Source: §8 · p. 15
  • The product of two derivative-augmented functions is defined by pairing their outputs and pairing their linear maps.Source: §8 · p. 15
  • The definitions include linearD-based forms for projections, duplication, injections, and joining.Source: §8 · p. 15

9 Extracting a Data Representation

The section moves from costly function-based Jacobian extraction to composable generalized-matrix representations, then uses continuation transformations to left-associate compositions for reverse-mode AD. These representations trade extraction cost against construction cost while preserving categorical composition.

9 Extracting a Data Representation

  • Function-based Jacobian extraction requires m passes for an m-dimensional input, making it expensive as dimension grows, especially for high-dimensional optimization.Each pass applies the derivative map to one vector from a basis of the input space.Source: §9 · p. 16

10 Generalized Matrices

  • Representable functors provide composable, type-safe data representations that save recomputation, as an alternative to multidimensional arrays.The representations are described as functional memoization.Source: §10 · p. 16

11 Efficiency of Composition

  • Generalized matrices avoid expensive matrix extraction but make construction operations used throughout more expensive.Function representations have simple, efficient composition, whereas generalized matrices shift cost toward construction.Source: §11 · p. 16
  • For reverse-mode AD, fully left-associated compositions are more efficient, while continuation transformations reconcile that association with modular program organization.The paper describes fully right-associated compositions as forward mode and other associations as mixed modes.Source: §11 · p. 17, §12 · p. 17

12 Reverse-Mode Automatic Differentiation

  • Continuation-based Contr k provides a simple reverse-mode AD algorithm by representing derivatives as linear functions that map from codomain to domain.The continuation representation converts composition patterns into fully left-associated form.Source: §12 · p. 17

13 Gradients and Duality

The paper derives a dual representation for reverse-mode AD by converting continuation representations using the scalar-field dual isomorphism. It also describes category transformers that reverse association to yield forward-mode AD, preferable for low-dimensional domains.

13 Gradients and Duality

  • Dual representations are derived by converting continuation representations through an isomorphism between scalar-field duals and the original spaces.The resulting duality reverses sequential compositions and swaps corresponding Cartesian and Cocartesian operations.Source: §13 · p. 19
  • The (△) and (▽) operations mutually dualize, corresponding to horizontal and vertical juxtaposition, respectively, in the matrix perspective.In that perspective, duality is transposition; Dualk need not perform actual matrix computations.Source: §13 · p. 19

14 Forward-Mode Automatic Differentiation

  • Turning the Cont and Dual techniques around yields full right-association and forward-mode AD, preferable for low-dimensional domains rather than codomains.The construction chooses the scalar field for r, using the isomorphism (s ⊸ a) ∼= a.Source: §14 · p. 21

15 Scaling Up

Indexed products provide a practical way to scale categorical AD to high-dimensional spaces, while the related-work discussion situates the paper’s graph-free, functorial approach and its memory and parallelism properties.

15 Scaling Up

  • Indexed products avoid unwieldy nested binary-product representations and support a homomorphic AD construction for indexed biproducts.The paper says the collected morphisms must share a domain and codomain, a restriction chosen to accommodate Haskell’s type system.Source: §15 · pp. 21–22, Figure 13
  • The indexed categorical operations include collections of projections and injections, replication and combination, plus indexed variants of fork and unfork.These operations can derive further operations, including fmap from crossI, replI, and composition.Source: §15 · p. 21

16 Related Work

  • The paper’s approach is specified through functoriality without manipulating graphs or other syntactic representations, and its calculated forms are correct by construction.It complements work that formulates backpropagation as a functor but remains tied to its machine-learning application.Source: §16 · pp. 23–24
  • The categorical account makes reverse-mode fanout explicit as dup, avoiding additional complexity in the translation from Haskell code.A-normal form is another way to make fanout lexically apparent in straight-line code.Source: §16 · pp. 23–24
  • Building on the compiler plugin, the paper adds Contr k and Dualk and derives reverse-mode AD implementations from functorial specifications.The duality-based gradient algorithm uses neither traditional matrices nor generalized matrices represented by representable functors.Source: §16 · p. 24
  • Unlike typical reverse-mode implementations that retain mutation-updated accumulators until completion, these implementations can reuse memory as they run.The paper says detailed time- and space-use comparisons remain worthwhile, and that lack of mutation makes the algorithms naturally parallel, with speedup only a potential benefit.Source: §16 · pp. 24–25

17 Conclusions

The paper derives a mode-independent AD algorithm from functorial specifications and specializes it to reverse-mode algorithms. It presents explicit parallel composition and a compiler-based view of differentiation as symbolic work, while identifying detailed performance analysis as future work.

17 Conclusions

  • Two categorical specializations of a generalized AD algorithm yield reverse-mode AD for general derivatives and for gradients, with implementations described as far simpler than previous ones.The algorithms avoid graphs, tapes, variables, partial derivatives, and mutation, and use an AD-agnostic compiler plugin with an existing language.Source: §17 · p. 25
  • Detailed performance analysis against backpropagation and conventional AD, efficient higher-order differentiation, and applications to derivative-like notions remain future work.The listed applications include subdifferentiation and automatic incrementalization.Source: §17 · p. 25
  • Explicit general support for parallel composition puts operations on equal footing regardless of arity, alongside the usual rule for sequential composition.The paper also gives a rule for linear operations.Source: §17 · p. 25
  • The paper recasts automatic differentiation as symbolic differentiation performed by a compiler, which preserves sharing in computations.It argues that this addresses expression blowup and the need for symbolic manipulation in a computer algebra system.Source: §17 · p. 25
  • The specification and implementation make a step toward differentiable functional programming, allowing programmers to define functions in a typed functional language and differentiate where desired.The paper presents this as a step toward the stated vision, not as a completed realization of it.Source: §17 · p. 25

A Terminal and Initial Objects

The appendix develops categorical preliminaries and checks that differentiation respects composition, parallel products, and the continuation transformation. It also establishes the hom-set additive structure used by these constructions.

A Terminal and Initial Objects

  • In the paper’s biproduct setting, terminal and initial objects coincide and can be represented by the unit type, a canonical zero-dimensional vector space.For functions, the terminal map returns unit and the initial map sends unit to zero.Source: A Terminal and Initial Objects · p. 26

B Abelian Categories

  • Morphisms with a shared domain and codomain form an abelian group, with zero built by composing initial and terminal morphisms and addition defined through duplication and combination.The appendix also gives identities expressing the derived operations △ and ▽ using projections and injections.Source: B Abelian Categories · p. 26

C.1 Corollary 1.1

  • The derivative of a composition is the composition of the derivatives, and the appendix refactors the rule to expose this compositional structure.The refactoring shares the intermediate value produced by the inner function.Source: §C.1 · p. 26

C.2 Corollary 2.1

  • For parallel products, the derivative of f × g is the product of the derivatives of f and g, evaluated at their respective inputs.The derivation expresses the result by applying D+ separately to each function.Source: §C.2 · p. 27

C.3 Theorem 4

  • The continuation transformation cont is shown to preserve identity and composition, with composition order reversed in its transformed form.The appendix strengthens the composition condition into a solved form used to derive the construction.Source: §C.3 · p. 27

= Cont

The appendix derives continuation-based categorical instances from functorial specifications and verifies linearity properties needed for the dual construction. It also establishes that invertible linear maps have linear inverses.

= Cont

  • The continuation transformation’s categorical instances follow by solving its functorial, Cartesian, and Cocartesian requirements; the resulting optimized forms match Figure 7.The derivation uses the join/unjoin isomorphism to split a continuation over a product into component continuations.Source: = Cont · pp. 28–29

C.4 Theorem 5

  • The appendix proves that dot and dot−1 are linear, using bilinearity of the uncurried dot product and the inverse relationship between the two maps.These properties support the construction of Dualk instances.Source: §C.4 · p. 29, §2 · p. 30

2. Invertible linear functions have linear inverses. In particular,

  • An invertible linear function has a linear inverse, as shown for dot−1 through preservation of addition and scalar multiplication.The derivation uses dot ◦ dot−1 = id and dot−1 ◦ dot = id.Source: §2 · pp. 29–30

3. Noting that the argument of both sides is a pair,

The section establishes identities for derivatives of paired functions, then constructs the dual continuation representation and verifies its categorical operations. The resulting representation maps standard structural operations to their dual counterparts.

3. Noting that the argument of both sides is a pair,

  • For paired inputs, the derivative of a pair separates into the product of the component derivatives, using linearity.The proof also derives the corresponding identity involving ▽ from the pair-derivative identity.Source: §3 · p. 30, §5 · p. 30

5. Noting again that the argument of both sides is a pair

  • The dual continuation representation is defined by applying Dual to onDot, which conjugates a function by dot and its inverse.onDot f = dot^-1 ◦ f ◦ dot.Source: §5 · p. 30

asDual

  • For the Category instance of Dualk, identity preservation is verified by simplifying asDual id = asDual (Cont id) = Dual (dot−1 ◦ id ◦ dot) = Dual id.Source: §5 · p. 30
  • Under asDual, exl and inl map to their dual forms, while dup maps to jam and jam maps to dup.The derivations use the paired-derivative identities and dot^-1 ◦ dot = id.Source: asDual · pp. 31–32
  • The construction also maps inr to Dual exr and preserves scaling; its product operation is shown to respect the monoidal structure.The product-preservation proof simplifies both sides to Dual (onDot f × onDot g).Source: asDual · pp. 31–32

C.5 Corollary 5.1

The appendix derives product and indexed-product rules for the derivative transformations. It shows compositionality for indexed cross operations and gives a supporting property proved in the function category.

C.5 Corollary 5.1

  • Dual f △ Dual g equals Dual (f × g) ◦ dup, while Dual f ▽ Dual g equals Dual (f △ g).These identities follow by expanding the definitions of △, ▽, and the product operations for Dualk.Source: §C.5 · pp. 32–33

C.6 Theorem 6

  • D distributes over paired functions: D (f × g) (a, b) = D f a × D g b.The theorem is also expressed as cross ◦ (D f × D g).Source: §C.6 · p. 33

Theorem 8 (indexed cross rule)

  • D+ is compositional with respect to crossI, and its indexed rule is expressed using second crossI, unzip, and crossI over fmap D+ fs.The derivation rewrites the paired value and derivative through Theorem 8 and Lemma 9.Source: Theorem 8 (indexed cross rule) · p. 33

Lemma 10

  • The supporting transpose property is proved in the function category, which the passage says suffices for the paper's purpose.The passage raises, but does not resolve, whether the property holds for general biproduct categories.Source: Lemma 10 · p. 34

Related papers

Loading 1804.00746v4…