Reference Categorical Foundations of Gradient-Based Learning [bruno-etal-categorical-learning-2021]

Abstract

  • July 13, 2021
  • G. S. H. Cruttwell, Bruno Gavranović, Neil Ghani, Paul Wilson, Fabio Zanasi

We propose a categorical semantics of gradient-based machine learning algorithms in terms of lenses, parametrised maps, and reverse derivative categories. This foundation provides a powerful explanatory and unifying framework: it encompasses a variety of gradient descent algorithms such as ADAM, AdaGrad, and Nesterov momentum, as well as a variety of loss functions such as as MSE and Softmax cross-entropy, shedding new light on their similarities and differences. Our approach to gradient-based learning has examples generalising beyond the familiar continuous domains (modelled in categories of smooth maps) and can be realized in the discrete setting of boolean circuits. Finally, we demonstrate the practical significance of our framework with an implementation in Python. 

 @article
{bruno-etal-categorical-learning-2021, title={Categorical Foundations of Gradient-Based Learning}, url={http://arxiv.org/abs/2103.01931}, DOI={10.48550/arXiv.2103.01931}, abstractNote={We propose a categorical semantics of gradient-based machine learning algorithms in terms of lenses, parametrised maps, and reverse derivative categories. This foundation provides a powerful explanatory and unifying framework: it encompasses a variety of gradient descent algorithms such as ADAM, AdaGrad, and Nesterov momentum, as well as a variety of loss functions such as as MSE and Softmax cross-entropy, shedding new light on their similarities and differences. Our approach to gradient-based learning has examples generalising beyond the familiar continuous domains (modelled in categories of smooth maps) and can be realized in the discrete setting of boolean circuits. Finally, we demonstrate the practical significance of our framework with an implementation in Python.}, number={arXiv:2103.01931}, publisher={arXiv}, author={Cruttwell, G. S. H. and Gavranović, Bruno and Ghani, Neil and Wilson, Paul and Zanasi, Fabio}, year={2021}, month={Jul} }

Backlinks

Example [efr-37F7]

Let \mathsf {CartSp} denote the category of Cartesian spaces \mathbb {R}^n and smooth maps between them, which acts on itself by Cartesian product. Then there is a functor \mathsf {SmMfd} \to \mathsf {Optic}(\mathsf {SmMfd}) (= \mathsf {Lens}(\mathsf {SmMfd})) which carries \mathbb {R}^n to the pair {\mathbb {R}^n \choose \mathbb {R}^n} and a smooth map f: \mathbb {R}^n \to \mathbb {R}^m to the pair (f, Rf) where Rf: \mathbb {R}^n \times \R ^m \to \mathbb {R}^n is the reverse derivative, where Rf(x,-) : \mathbb {R}^m \to \mathbb {R}^n is the linear map given by the transpose of the Jacobian.

Moreover this is monoidal for the Hancock tensor on \mathsf {Optic}(\mathsf {CartSp}). Since \mathsf {Para}(-) is functorial, we obtain a functor \mathsf {Para}(\mathsf {CartSp}) \to \mathsf {Para}(\mathsf {Optic}(\mathsf {CartSp})). Given some parametrized (smooth) map X \times P \to Y, we obtain an optic {X \choose X} \otimes {P \choose P} \to {Y \choose Y}. After choosing x \in X and a section d: Y \to Y---for example, by sampling x,y^* from a dataset and letting d(y) be the gradient of (y-y^*)^2 at y---the resulting gradient P \to P describes how to update the parameters to move y in the desired direction. The fact that this is functorial is essentially the mechanism behind the backpropagation algorithm---see Reference [backprop-as-functor], Reference [bruno-etal-categorical-learning-2021] for more.

(Note that this construction can also be generalized to smooth manifolds, but one must pass to dependent lenses and use the cotangent manifold instead.)

Example Machine Learning as a bisystem [efr-000E]

It is an observation which goes back at least to Reference [backprop-as-functor], and was more thoroughly developed in Reference [bruno-etal-categorical-learning-2021], that the process of training a machine learning algorithm by gradient descent can be abstracted in the following way:

  • Wanting to learn a function X \to Y, we choose a parametrized function f: X \times P \to Y, where all these are, in the simplest case, Euclidean spaces \mathbb {R}^k
  • We take the backwards derivative of f, obtaining a lens: (f,Df): {TX \choose X} \otimes {TP \choose P} \leftrightarrows {TY \choose Y}
  • For each datum (x_n,y_n), we compute the loss gradient \nabla L(-,y_n) : Y \to TY, and combining this with x_n \in X and the current parameter p, we get a gradient on the parameter space which we can use to update

Thus a machine learning algorithm is a sort of bisystem. Indeed our bisystems are essentially an abstracted version of the learners of Fong--Spivak--Tuyeras. The functoriality of this assignment is the main point of the above-mentioned papers.

Markov Fibrations › Introduction [efr-HQ73]

In applying category theory to fields as diverse as game theory (Reference [hedges-etal-comp-gametheory],Reference [hedges-etal-bayesian-games], Reference [hedges-etal-graph-games]), machine learning (Reference [bruno-etal-categorical-learning-2021], Reference [bruno-thesis-fundamental-components]), dynamical systems (Reference [myers-cst], Reference [lynch-myers-rischel-staton-stoch-clocks]), and server design (Reference [videla-capucci-servers]), people have found it useful to study categories whose morphisms describe processes or functions that take place in two "stages", where the second is dependent on the first, but composes in the other direction---so-called lenses. In many of these domains, stochastic phenomena play a role. There is a useful generalization of lenses to categories of stochastic maps---the so-called optics of Riley, Reference [riley-optics]---but this does not accommodate another useful generalization, so called dependent lenses where not only the backwards process but the set it takes values in is indexed over the base. It has been a long-standing problem to develop a suitable common generalization, "dependent optics", of these two ideas. In this thesis, we solve this problem by developing a theory of Markov fibrations (§ [efr-O088]).

A Markov fibration is a weakening of the notion of (Grothendieck) fibration to include (subject to some assumptions) Markov categories of indexed families of objects (given by deterministic functions E \to X) and compatible stochastic maps (given by commutative squares). The main point of Markov fibrations is that they admit fiberwise opposites, which generalize the fiberwise opposites of ordinary categories. Just as the fiberwise opposite of the codomain fibration of a finitely complete category describes dependent lenses (see § [efr-ZCTD]), the fiberwise opposite of these codomain Markov fibrations give a good notion of stochastic lens (see Theorem [efr-K6NM]).

We will apply our theory chiefly to two problems. First, we will use them to generalize open games (Hedges, Reference [hedges-towards-compositional-thesis]) to a larger class of interfaces (namely, indexed families of sets) while at the same time considering possibly-stochastic maps. This will allow us to construct the so-called external choice operator on open games, which describes branching.

Secondly, we will generalize Myers' categorical dynamical systems theory to allow for stochastic maps in the base. When combined with another generalization of these systems, to general parametrized lenses, this gives a more natural way of modeling certain stochastic dynamical systems, such as those associated to the training dynamics of machine learning models, see eg Example [efr-IMZ3].

A brief introduction to lenses

If \mathcal {C} is a category with finite products, the category \mathsf {Lens}(\mathcal {C}) of lenses in \mathcal {C} has objects pairs \binom {A}{X} of objects in \mathcal {C}, and morphisms \binom {A}{X} \leftrightarrows \binom {B}{Y} given by pairs f: X \to Y, f^\#: B \times X \to A. (Note that it is of course the morphisms, not the objects, that are called lenses). One way to generalize this is to ask for \mathcal {C} to have pullbacks, and consider an object given by a more general map A \to X, and let a map be f: X \to Y, f^\#: B \times _Y X \to A (so that the triangle over X commutes). (Note that this recovers the "simple lenses" when the objects are of the form A \times X \to X). Under the interpretation of a map A \to X as a family of objects A_x indexed over the elements of X, we see this as a "dependent lens" (since A is a "dependent type"). There are various variations of this idea, which ultimately all fit the pattern of taking a Grothendieck fibration \mathcal {D} \to \mathcal {C} and forming the fiberwise opposite---for the dependent lenses as above, this is the codomain fibration \mathcal {C}^\to \to \mathcal {C}.

Meanwhile, another wide-ranging generalization of \mathsf {Lens}(\mathcal {C}) are the optics of Riley (Reference [riley-optics]). For \mathcal {C} a monoidal category, the objects of \mathsf {Optic}(\mathcal {C}) are again pairs \binom {A}{X} of objects, but now the morphisms are elements of the coend \int ^{M \in \mathcal {C}}\mathcal {C}(X, M \otimes Y) \times \mathcal {C}(M \otimes B, A). That is, to give an optic \binom {A}{X} \leftrightarrows \binom {B}{Y} is to give an object M \in \mathcal {C} and morphisms l: X \to M \otimes Y, r: M \otimes B \to A, up to the equivalence relation generated by, whenever s: N \to M, l: X \to N \otimes Y, r: M \otimes B \to A, identifying the two optics given by (M, (s \otimes 1_Y) \circ l, r) and (N, l, r \circ (s \otimes 1_B)). One can show that, in the case where \mathcal {C} is a Cartesian monoidal category, this set can be identified with the set of (simple) lenses.

A vector field on a smooth manifold X is simply a (smooth) section of the tangent bundle TX. A section is a bundle map from the trivial bundle X \to X. This gives the idea of considering a parametrized dynamical system as consisting of a map X \to A to some other smooth manifold, a bundle E \to A (the points of E are the parameters) and a bundle map E \times _A X \to TX. As we will see, such a bundle map is exactly a dependent lens in a category of bundles. Building on this idea, Myers (Reference [myers-cst]) described a highly abstract theory of open dynamical systems, given by lenses out of tangent bundles (where the notion of space, bundle and tangent bundle are generalized to any fibration).

It has been observed that, although a lens from the tangent bundle TS \leftrightarrows A may indeed be said to describe an open dynamical system with state space S, the same can be said for a lens of type TS \otimes A \leftrightarrows I---although a dynamical system of a different kind (essentially, the first type are the Moore machines, the second the Mealy machines---see the introduction to § [efr-000D] for more on this). It is natural to consider parametrized maps TS \otimes A \leftrightarrows B as a common generalization of these two concepts---and in fact, special cases of this idea, lenses parametrized by tangent bundles, have already been considered many times, for example in the semantics of gradient descent (see Reference [bruno-etal-categorical-learning-2021], Reference [towards-cybercat], Reference [backprop-as-functor]).

Since parametrized maps compose in an obvious way, we now have three different notions of morphisms between bundles---lenses, charts, and "generalized systems". These should form some sort of symmetric monoidal triple category, but the right axiomatisation of this concept is somewhat elusive.

Lenses, in their various incarnations, have been widely used in applied category theory. For example, Hedges and his collaborators have developed a compositional approach to game theory (Reference [hedges-string-games], Reference [hedges-etal-comp-gametheory], Reference [hedges-etal-institutions], Reference [hedges-etal-graph-games], others). Recent work by Hedges and Sakamoto bring this idea to reinforcement learning (Reference [sakamoto-reinforcement-lenses]). We briefly mentioned Myers' categorical systems theory above, and we will see much more of it later. As we discussed above, there is also a literature using lenses to study gradient descent in an abstract sense. This is without even mentioning their original role in the theory of functional programming (this is the origin of the somewhat confusing term lens), or their prehistory in Gödel's dialectica interpretation---see eg Reference [jules-lenses-blogpost] for a survey of this.

Dependent Optics

An open game, in the sense of Hedges, has as its interface two objects in the category of lenses (of sets) \binom {S}{X}, \binom {R}{Y} (that is, two pairs of sets). An open game with this interface consists of a set \Sigma of strategies equipped with a function \Sigma \to \mathsf {Lens}(\mathsf {Set})(\binom {S}{X},\binom {R}{Y}), and a subset Eq(x, k) \subset \Sigma of equilibrium strategies for each x \in X and k: Y \to R---note that such an x is exactly a lens \binom {*}{*} \leftrightarrows \binom {S}{X}, and such a k is exactly a lens \binom {R}{Y} \to \binom {*}{*}. This view of open games makes the composition rule much simpler to define, although we will not delve into the details here, see eg Reference [hedges-etal-comp-gametheory]. The idea is that the open game represents the strategies and preferences of a player or set of players---\sigma \in \Sigma are the possible strategies, x \in X is the information revealed to the player before they make their decision of which moves to make, the resulting y \in Y is the choice the player "sends" in response to x, and the r \in R is the "utility", the value they eventually learn, depending on their choice y, which they have some preferences over.

The above description gives a theory of deterministic games, but of course, it is completely essential for game theory to model both random decisions (mixed strategies) and decisions made under uncertainty. This motivates the replacement of lenses in this definition with optics in a category of probability kernels (the form of the equilibrium relation must be modified as well, see Reference [hedges-etal-bayesian-games], Reference [towards-cybercat])

In the abstract study of open games, it is desirable to define a "choice" operation \oplus on the objects (the pairs of sets) so that maps into \binom {R}{X} \oplus \binom {R'}{X'} represent players who are faced with some binary choice between X and X', after which play may proceed for a time in one of two branches. However, this is immediately problematic because the category of lenses do not have coproducts---inside the category of dependent lenses, however, we can form the coproduct simply as R \times X + R' \times X' \to X + X'---that is, we get a family which is R over X, and R' over X'. This indeed does the job, but for game theory the extension to stochastic maps---optics---is absolutely essential. This motivates the search for a theory of dependent optics, a common generalization of optics and dependent lenses.

There have been a number of attempts at this---see eg Reference [fibre-optics-2021], Reference [hedges-braithwaite-dep], Reference [milewski-polylens], Reference [vertechi-dep-optics]. While these efforts have to some extent succeeded in describing a theory that is general enough to contain the desired examples, it is generally very ad hoc. The construction of optics as a special case of Vertechi's dependent optics (Reference [vertechi-dep-optics]) gives them as fibre optics with the indexing bicategory being (the delooping of) a monoidal category \mathcal {M}, but embeds lenses using a bicategory of spans. Thus simple lenses admit two distinct encodings in this theory. Moreover, Vertechi's example of dependent monoidal optics fails to describe a supercategory of \mathsf {Optic}(\mathcal {C}) when applied to Markov categories, and are thus not suitable for, for example, game theory.

Categorical Cybernetics

The observation that both open games and machine learning systems can be fruitfully described using lenses, naturally led to the idea that this could be used to develop a more thorough analogy between the two. After all, a machine learning system is also a player in a kind of game (with the objective of minimizing loss). Hedges coined the phrase categorical cybernetics to describe this locus of ideas (Reference [towards-cybercat], also used in Reference [sakamoto-reinforcement-lenses]).

Myers' categorical systems theory also describes dynamical systems as lenses. However, where a cybernetic system in the sense of Reference [towards-cybercat] is a parametrized lens \Sigma \otimes {\bar {X} \choose X} \leftrightarrows {\bar {Y} \choose Y}, and thus has two different inputs - x \in X, which is the information on which basis they are allowed to choose their decision y \in Y, and y^\sharp \in \bar {Y}, which is their "utility", the information they are allowed to care about.

It is straightforward to describe the machine learning part of the categorical cybernetics analogy in terms of Myers' theory---a "learner" is simply a lens of the above form with \Sigma replace with a tangent bundle TS. This idea (described in these terms, although not with the analogy to machine learning) has already been described by Reference [energy-driven-systems]. However, for game theory, it is necessary that the output of a given system can be stochastically chosen (a player must be able to randomize their strategy). At the same time, for a machine learning system, the x \in X is usually some sort of sample from a training distribution---ie, it is random. Hence to really describe the training dynamics of machine learning systems using this theory, we are again naturally drawn to consider systems theories (that is, fibrations) which have stochastic maps in the base, not merely the fiber.

Overview and contributions

We begin the thesis in § [efr-POD2] with a review of some preliminary material. This chapter can largely be skipped for readers already familiar with the material (Markov categories, double categories, lenses and optics, fibrations, and categorical dynamical systems theory). The chapter on Markov categories introduces a few novel auxiliary notions, but these can be referred back to as necessary (they are primarily regularity conditions which hold in most Markov categories of interest).

In § [efr-O088], we give the main technical contribution of the thesis by developing a theory of Markov fibrations. Their fiberwise opposites generalize both the fiberwise opposites of codomain fibrations (dependent lenses in a classical sense) and optics in Markov categories. Theorem [efr-K6NM] provides the statement of the latter (the former is straightforward). We also give a description of monoidal structures on Markov fibrations (which pass to monoidal structures on the resulting categories of optics).

In § [efr-000D], we give a construction of the double category of parametrized morphisms for a category action (or actegory)---in fact, we do this internally to any 2-category. This has the advantage of allowing us to construct more highly-structured versions of this double category by carrying out the construction internally to higher-structured actions. (Bicategorical versions of \mathsf {Para} are a relatively old idea, and the double categorical version is a straightforward extension which has existed for some time as folklore, but the fully-internal construction here is novel).

In § [efr-GFG0], we apply Markov fibrations to compositional game theory, and solve a longstanding problem by giving a category of stochastic open games with a general "external choice" operator. This construction is largely abstract over the particular category, and so can be applied to different Markov fibrations to describe games defined with different notions of stochasticity. We also review some previous work with Capucci, Hedges, and Gavranovic on abstract constructions of categories of open games (Reference [towards-cybercat]).

In § [efr-ZRUZ], we describe how to extend Myers' categorical systems theory to Markov fibrations, to give theories of systems which may have stochastic maps not merely in the fiber (as was already described by Myers in Reference [myers-cst]) but also in the base. We also leverage our construction from § [efr-000D] to give a description of a symmetric monoidal triple category whose three types of morphisms, are lenses, charts, and what we call bisystems, that is morphisms TS \otimes A \leftrightarrows B - these are open dynamical systems which have two directions of interaction with the environment. We also describe a particular systems theory of smooth manifolds and smooth stochastic maps between them.

Acknowledgements [efr-0008]

I would like to extend a great thanks to the Mathematically Structured Programming group at the University of Strathclyde, my fellow PhD students in particular, for making my time there such a great one. A good chunk of this work was completed while I worked at the Compositional Systems Laboratory at the Tallinn University of Technology---I owe everyone there a great thanks as well.

For conversations which contributed to this work, and to my mathematical development in general, I would like to thank Dylan Braithwaite, Matteo Capucci, Elena di Lavore, Tobias Fritz, Davidad, Tomáš Gonda, Bruno Gavranović, Diana Kessler, Owen Lynch, Jade Master, David Jaz Myers, Riu Nakamura, Chad Nester, Evan Patterson, Paolo Perrone, Mario Roman, Brandon Shapiro, Toby Smithe, Pawel Sobocinski, David Spivak, and Andre Videla. A special thanks to Dylan, Matteo, and Bruno, for serving their time with me in the dependent optics mines.

To Neil Ghani, Radu Mardare, and Jules Hedges, a massive thanks for serving as my advisors during the somewhat tumultuous process that was the production of this thesis.

Finally, I thank my wonderful family. Now you can stop asking me about it.

The \mathsf {Para} construction in generic 2-categories › Introduction [efr-003P]

Myers' theory of categorical systems theory (see § [efr-0023]) gives a rich categorical structure to a wide variety of types of dynamical system. The central idea can be summarized by saying that there are two different, but tightly related, notions of morphism at play in dynamical systems. Open dynamical systems themselves involve a bidirectional information flow, captured by the notion of lens {X' \choose X} \leftrightarrows {Y' \choose Y}, and composition of such lenses describes the composition of subsystems into systems. But morphisms between systems are unidirectional, captured by the notion of chart {X' \choose X} \rightrightarrows {Y' \choose Y},. Algebraically, the relationship between these two notions is that they assemble into a double category, which indexes the category of systems.

There is another double category involving lenses which has been considered in the categorical study of systems. That is the double category of parametrized morphisms, applied to the monoidal category of lenses. These have been studied as an abstraction for gradient descent in several papers, see eg Reference [bruno-etal-categorical-learning-2021], Reference [bruno-thesis-fundamental-components], (the idea goes back to Reference [backprop-as-functor]). Given any action of a monoidal category \mathcal {M} on another category \mathcal {C} (most simply, if \mathcal {C} is monoidal it acts on itself,) we obtain a category of parametrized morphisms f: X \cdot P \to Y (where P \in \mathcal {M}, X, Y \in \mathcal {C},) and these turn out to be extremely useful. We will mention two applications:

  1. A parametrized lens {P \choose P} \otimes {X \choose X} \to {Y \choose Y} is essentially what is called a learner in Reference [backprop-as-functor]---it contains the information necessary to compute a new parameter value p' given an existing p \in P and a sample pair x \in X, y \in Y, as well as the additional information required to compose such things. Thus the functoriality of backpropagation can be derived from two facts: the reverse derivative defines a monoidal functor \mathsf {Euc} \to \mathsf {Lens}(\mathsf {Euc}), and the construction \mathsf {Para}(-), taking a category to its category of parametrized maps, is itself functorial. This viewpoint has been significantly developed in Reference [bruno-etal-categorical-learning-2021] and other papers.
  2. An open game, in the sense of Hedges Reference [hedges-etal-comp-gametheory], is almost the same thing as a parametrized lens {\Sigma ' \choose \Sigma } \otimes {S \choose X} \leftrightarrows {R \choose Y}, equipped with a subset E \subset \Sigma \times \mathsf {Set}(\Sigma ,\Sigma '). Given a context for the game---that is, a state x \in X (describing the state of information when the decision is made) and a continuation Y \to R (describing how decisions y \in Y map to outcomes r \in R,) we obtain a function k: \Sigma \to \Sigma ', and we say \sigma \in \Sigma is an equilibrium strategy if (\sigma ,k) \in E. Since \mathsf {Set}(\Sigma ,\Sigma ') = \mathsf {Lens}(\mathsf {Set})({\Sigma ' \choose \Sigma },I), and \Sigma = \mathsf {Lens}(\mathsf {Set})(I,{\Sigma ' \choose \Sigma }), this neatly captures the extra data of an open game in terms of the category of lenses. The potential of this idea as a generalized approach to "cybernetic systems" is explored in Reference [towards-cybercat].

In this chapter, we will develop the theory of the category \mathsf {Para} of parametrized maps. We will begin by reviewing the existing literature briefly. We will describe a double categorical version of this category---this does not seem to have appeared in the literature yet, although it has been folklore for at least a few years (and there is nothing complicated about this construction, certainly). The remainder of this chapter will be dedicated to lifting this construction to a generic 2-category \mathbb {C} (with the above being the specialization to \mathbb {C} = \mathsf {Cat} ). We will derive this lifting using the machinery of 2-category theory. We will see how this generalization accounts for much structure which can be seen to exist on \mathsf {Para}, such as its symmetric monoidal structure (assuming \mathcal {M},\mathcal {C} are symmetric monoidal). But the true application of this will be in § [efr-ZRUZ], where we use this to construct a triple category of open dynamical systems.

The definition of the double category \mathsf {\mathbb Para} involves two types of categorical structure with which the reader may be unfamiliar---actegories, which are the input to the construction, and pseudo double categories, which are the output. Since we will shortly introduce the abstract internal versions of these, internal pseudomonoid actions and internal pseudocategories, we will not give a separate introduction here. The reader who is unfamiliar with these should refer to Reference [actegories-amthematician-capucci-gavranovic] for actegories, and Reference [shulman-monfibs] for (pseudo) double categories. A reader who simply needs a definition may look at Definition [efr-OFNV] and Example [efr-ZRUY].