Generalization Gradient Descent

## Abstract

Abstract Cartesian reverse derivative categories (CRDCs) provide an axiomatic generalization of the reverse derivative, which allows generalized analogues of classic optimization algorithms such as gradient descent to be applied to a broad class of problems.
In this paper, we show that generalized gradient descent with respect to a given CRDC induces a hypergraph functor from a hypergraph category of optimization problems to a hypergraph category of dynamical systems. The domain of this functor consists of objective functions that are 1) general in the sense that they are defined with respect to an arbitrary CRDC, and 2) open in that they are decorated spans that can be composed with other such objective functions via variable sharing. The codomain is specified analogously as a category of general and open dynamical systems for the underlying CRDC. We describe how the hypergraph functor induces a distributed optimization algorithm for arbitrary composite problems specified in the domain. To illustrate the kinds of problems our framework can model, we show that parameter sharing models in multitask learning, a prevalent machine learning paradigm, yield a composite optimization problem for a given choice of CRDC. We then apply the gradient descent functor to this composite problem and describe the resulting distributed gradient descent algorithm for training parameter sharing models.

## 1 Introduction

Figure: Figure 1: An illustration of multitask learning in our framework. The left string diagram shows the composite objective function representing the multitask learning objective. The right string diagram is the result of applying the gradient descent functor to the morphism on the left, resulting in a distributed optimization algorithm.

Recently, Cartesian reverse derivative categories (CRDCs) were used to define generalized analogues of classic optimization algorithms, such as gradient descent, that minimize generalized objective functions [, ]. This has allowed techniques from machine learning such as backpropagation to be applied not only to artificial neural networks, but to a broad class of models including boolean circuits.
The generality of the framework makes future expansion to other problem types likely.
Central to the definition of a CRDC is the Cartesian reverse derivative combinator R, which must obey axioms mirroring the behavior of the standard directional derivative in the Euclidean domain. This combinator is *compositional* in that it satisfies a chain rule, allowing one to define a general version of the backpropagation algorithm [, ].

In prior work, CRDCs have been used primarily to compose parameterized morphisms, i.e to build a learning model, all the while leveraging R to ensure the parameters can be updated with respect to the generalized gradient of an objective. In this way the induced optimization problem of maximizing/minimizing the given objective with respect to the parameters has a compositional structure by virtue of the *model* being compositional. We propose that in addition to compositional structure in the *model*, compositional structure in the *objective* is also of interest. In particular, in machine learning it is common for parameters to be optimized for more than one objective. Simple examples of compositional objectives are those arising from various regularization methods, for example $\ell_{1}$ or $\ell_{2}$ penalties [], wherein one term in the objective incentivizes “good performance” on the given task while the other enforces some generically desirable property like sparseness. Other examples with more sophisticated composition patterns include parameter sharing methods wherein subsets of parameters may be optimized for one or more objective functions. Parameter sharing arises within the setting of multitask learning, which we will discuss in this paper. Figure 1 illustrates this example.

To capture composition of objective functions that may have shared parameters, we turn to decorated spans []. Our approach closely follows that of [], where algebras of undirected wiring diagrams are used to compose objective functions with shared decision variables, and a variety of first-order methods including (sub)gradient descent and primal-dual methods are shown to be functorial in the standard Euclidean domain. This paper extends the results on gradient descent to the general setting of Cartesian reverse derivative categories and generalized optimization introduced in []. Specifically, our contributions are as follows.

- •
We show how to produce hypergraph categories of generalized open objectives and optimizers over a given optimization domain.
- •
We prove that generalized gradient descent yields a hypergraph functor from open objectives to open optimizers.
- •
We highlight how functoriality allows the gradient descent solution algorithms for composite problems to be implemented in a distributed fashion.
- •
We show that multitask learning [] with hard parameter sharing is an example of objective composition occurring in the machine learning literature, and we leverage our functorial formulation of gradient descent to recover a distributed optimization scheme.

We will begin by setting up the categories of open generalized objectives and open dynamical systems. We then specify the generalized gradient descent functor between them before turning to the multitask learning example.

## 2 Generalized Gradient Descent as a Hypergraph Functor

In order to proceed we need to recall many definitions from prior work. For convenience, they have been included in Appendix A. The reader who is unfamiliar with Cartesian reverse derivative categories, linear maps in CRDCs, optimization domains, and generalized gradients may wish to peruse the appendix before proceeding.

The fundamental component for generalized optimization is the notion of an *optimization domain*, which generalizes the role of the real numbers in defining the “cost” or “value” of a given configuration of decision variables. Specifically, an optimization domain consists of a pair $(\mathcal{C},R)$ where $\mathcal{C}$ is a Cartesian reverse derivative category and $R\in\mathcal{C}$ has the structure of an ordered commutative ring (for full details, see Definition A.5). To construct hypergraph categories of objectives and optimizers, we will also require that the subcategory of linear maps in $\mathcal{C}$, denoted $\mathsf{Lin}(\mathcal{C})$, has finite limits. Going forward, we use $\times$ to denote the product and $*$ to denote the terminal object in an arbitrary CRDC.

### 2.1 The Hypergraph Category of Open Objectives

The goal of this section is to define a hypergraph category of generalized optimization problems. In standard optimization, a generic unconstrained minimization problem has the form

$$ $\text{minimize }f(x),$ $$

where $f\colon\mathbb{R}^{n}\to\mathbb{R}$ is known as the objective function, and $x\in\mathbb{R}^{n}$ is referred to as the decision or optimization variable. When $f(x)$ has the form $\sum_{i=1}^{N}f_{i}(x_{i})$, where the $x_{i}$ are disjoint subvectors of $x$, $f$ is referred to as *separable*. Separable objectives can be solved entirely in parallel as each optimal $x_{i}$ has no effect on the rest of the objective. However, in many important problems, the objective function is *partially separable*. For example, the problem

$$ $\text{minimize }f(w,x)+g(u,w,y)+h(u,w,z)$ $$

consists of a sum of objectives with *complicating* variables $u$ and $w$. Such a sum of objectives with shared variables is the type of compositional structure that we wish to capture and generalize in our hypergraph category of optimization problems.

A standard method for defining such a hypergraph category is to leverage the theory of decorated spans(^1^11The literature typically uses *covariant* decorating functors to define decorated cospan categories. However, all the same results hold for the dual construction of *contravariant* decorating functors and decorated spans, which we leverage in this paper.) [, ]. Decorated spans turn a closed system into an open system by specifying a boundary for the system in addition to morphisms that relate the domain of the system to its boundary. We now leverage this general pattern to define a notion of *open objectives*.

###### Definition 2.1 .

Given an optimization domain $(\mathcal{C},R)$, an open objective with domain boundary $X$ and codomain boundary $Y$ consists of a span $X\leftarrow N\rightarrow Y$ in $\mathsf{Lin}(\mathcal{C})$ together with a choice of objective on $N$, i.e., a map $f\colon N\to R$ in $\mathcal{C}$.

There is a straightforward interpretation of open objectives when the legs are both projections: the boundary objects are subobjects of the overall decision space which can be shared with other objectives. This intuition of sharing parts of the decision space is still helpful even when the legs are general linear maps. We now arrange open objectives into a hypergraph category by specifying the following decorating functor.

theoremobj
Given an optimization domain $(\mathcal{C},R)$, there is a contravariant lax symmetric monoidal functor $O_{\mathcal{C}}^{R}\colon(\mathsf{Lin}(\mathcal{C})^{\mathrm{op}},\times)\to(
\mathrm{\textnormal{{Set}}},\times)$ defined by the following maps (where we let $O\coloneqq O_{\mathcal{C}}^{R}$ for the remainder of the theorem and proof).

- •
Given an object $X$, $O(X)$ is the set of generalized objectives with domain $X$, i.e., the hom-set $\mathcal{C}(X,R)$.
- •
Given a linear map $\phi\colon X\to Y$ and objective $f\colon Y\to R$, $O(\phi)(f)\colon X\to R$ is the objective $f\circ\phi$ obtained by precomposition.
- •
Given objects $X,Y$, the product comparison $\varphi_{X,Y}\colon O(X)\times O(Y)\to O(X\times Y)$ is defined by
$\varphi_{X,Y}(f,g)\coloneqq f\circ\pi_{0}+g\circ\pi_{1},$
(1)
where $\pi_{0}\colon X\times Y\to X$ and $\pi_{1}\colon X\times Y\to Y$ are the natural projections.
- •
The unit comparison $\varphi_{0}\colon\{*\}\to O(*)$ selects the additive identity $0_{R}$ for the hom-set $\mathcal{C}(*,R)$ supplied by the left additive structure.

###### Proof.

$O$ is plainly functorial as it is just the contravariant hom functor $\mathcal{C}(-,R)$ restricted to act only on the linear maps of $\mathcal{C}$. We still need to verify the symmetric lax monoidal axioms. For naturality of the product comparison, we need the following diagram to commute for all objects $X,X^{\prime},Y,Y^{\prime}$ and linear maps $\phi\colon X^{\prime}\to X$ and $\psi\colon Y^{\prime}\to Y$:

$$ ${{O(X)\times O(Y)}}$ ${{O(X^{\prime})\times O(Y^{\prime})}}$ ${{O(X\times Y)}}$ ${{O(X^{\prime}\times Y^{\prime})}}$ $\scriptstyle{\varphi_{X,Y}}$ $\scriptstyle{O(\phi)\times O(\psi)}$ $\scriptstyle{\varphi_{X^{\prime},Y^{\prime}}}$ $\scriptstyle{O(\phi\times\psi)}$ $$

Letting $f\colon X\to R$ and $g\colon Y\to R$ be arbitrary, following the top path yields the morphism

$$ $f\circ\phi\circ\pi_{0}+g\circ\psi\circ\pi_{1}$ (2) $$

while following the bottom path yields

$$ $(f\circ\pi_{0}+g\circ\pi_{1})\circ(\phi\times\psi)=f\circ\pi_{0}\circ(\phi \times\psi)+g\circ\pi_{1}\circ(\phi\times\psi),$ (3) $$

where the equality comes from $\textbf{LA.1}.$ Now let $x\in X^{\prime}$ and $y\in Y^{\prime}$ be generalized elements. Then applying (2) to $(x,y)$ gives $f(\phi(x))+g(\psi(y))$ while applying (3) to $(x,y)$ gives

$$ $f(\pi_{0}(\phi(x),\psi(y)))+g(\pi_{1}(\phi(x),\psi(y)))=f(\phi(x))+g(\psi(y)),$ (4) $$

as desired. Commutativity of the symmetry, associativity, and unitality diagrams all follow from the fact that the hom-set $\mathcal{C}(X^{\prime}\times Y^{\prime},R)$ is a commutative monoid.
∎

We can now use our decorated spans of open objectives to construct a hypergraph category.

###### Corollary 2.1.1 (Open Generalized Optimization Problems) .

Given an optimization domain $(\mathcal{C},R)$ such that $\mathsf{Lin}(\mathcal{C})$ has finite limits, there is a hypergraph category $\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}$ defined by the following data.

- •
Objects are the same as those of $\mathcal{C}$.
- •
Morphisms are (isomorphism classes of) open generalized objectives.
- •
Given two open objectives $(X\xleftarrow{l_{1}}N\xrightarrow{r_{1}}Y,f\in O_{\mathcal{C}}^{R}(N))$ and $(Y\xleftarrow{l_{2}}M\xrightarrow{r_{2}}Z,g\in O_{\mathcal{C}}^{R}(M))$, their composite consists of the following span computed by pullback
${{N\times_{Y}M}}$${N}$${{N\times M}}$${M}$${X}$${Y}$${Z}$$\scriptstyle{l_{1}}$$\scriptstyle{r_{1}}$$\scriptstyle{l_{2}}$$\scriptstyle{r_{2}}$$\scriptstyle{\pi_{1}}$$\scriptstyle{\pi_{0}}$$\scriptstyle{b_{1}}$$\scriptstyle{b_{0}}$$\scriptstyle{\langle b_{0},b_{1}\rangle}$
(5)
together with the objective obtained from the composite morphism
$\{*\}\cong\{*\}\times\{*\}\xrightarrow{f\times g}O_{\mathcal{C}}^{R}(N)\times O
_{\mathcal{C}}^{R}(M)\xrightarrow{\varphi_{N,M}}O_{\mathcal{C}}^{R}(N\times M)
\xrightarrow{O_{\mathcal{C}}^{R}\langle b_{0},b_{1}\rangle}O_{\mathcal{C}}^{R}
(N\times_{Y}M).$
(6)
- •
The identity on $X$ is the pair $(X=X=X,\{*\}\xrightarrow{\varphi_{0}}O_{\mathcal{C}}^{R}(*)\xrightarrow{O_{
\mathcal{C}}^{R}(!)}O_{\mathcal{C}}^{R}(X))$.
- •
The monoidal product on objects is their product in $\mathcal{C}$ while the monoidal product of morphisms $(X\leftarrow N\to Y,f)$ and $(W\leftarrow M\to Z,g)$ is
$(X\times W\leftarrow N\times M\to Y\times Z,\varphi_{N,M}(f,g)).$
(7)
- •
The hypergraph maps are inherited from those of $\mathrm{Span}(\mathsf{Lin}(\mathcal{C}))$ together with the identity decoration.

###### Proof.

This follows by applying Proposition 3.2 and Theorem 3.4 in [] to Theorem 2.1.
∎

###### Definition 2.1 .

###### Proof.

###### Corollary 2.1.1 (Open Generalized Optimization Problems) .

###### Proof.

### 2.2 The Hypergraph Category of Open Optimizers

Having defined a hypergraph category of open objective functions, we will now use the same machinery of decorated spans to define a hypergraph category of open dynamical systems over a given CRDC.

###### Definition 2.2 .

Given a CRDC $\mathcal{C}$, an open dynamical system with domain boundary $X$ and codomain boundary $Y$ consists of a span $X\leftarrow N\rightarrow Y$ in $\mathsf{Lin}(\mathcal{C})$ together with a choice of system on $N$, i.e., a map $v\colon N\to N$ in $\mathcal{C}$.

Endomaps in a CRDC are also referred to as optimizers in [], thus we use the terms optimizer and dynamical system interchangeably. Similar to open objectives, the boundaries of open dynamical systems specify which parts of a system’s state space can be influenced by other systems. The decorating functor is given as follows.

theoremopt
Given a CRDC $\mathcal{C}$, there is a contravariant lax symmetric monoidal functor
$D_{\mathcal{C}}\colon(\mathsf{Lin}(\mathcal{C})^{\mathrm{op}},\times)\to(
\mathrm{\textnormal{{Set}}},\times)$ defined by the following maps (where we let $D\coloneqq D_{\mathcal{C}}$ for the remainder of the theorem and proof).

- •
Given an object $X$, $D(X)$ is the set of generalized dynamical systems with state space $X$, i.e., the hom-set $\mathcal{C}(X,X)$.
- •
Given a linear map $\phi\colon X\to Y$ and a system $v\colon Y\to Y$, $D(\phi)(v)$ is the system $\phi^{\dagger}\circ v\circ\phi\colon X\to X$.
- •
Given objects $X,Y$, the product comparison $\varphi_{X,Y}\colon D(X)\times D(Y)\to D(X\times Y)$ is defined by
$\varphi_{X,Y}(v,w)\coloneqq\pi_{0}^{\dagger}\circ v\circ\pi_{0}+\pi_{1}^{
\dagger}\circ w\circ\pi_{1}.$
(8)
- •
The unit comparison $\varphi_{0}\colon\{*\}\to D(*)$ is uniquely determined as $\mathcal{C}(*,*)$ is singleton.

###### Proof.

For functoriality of $D$, let $\phi\colon X\to Y,\psi\colon Y\to Z$, and $v\colon Z\to Z$ be arbitrary. For preservation of composition, we have

$$ $D(\phi\circ\psi)(v)\coloneqq(\phi\circ\psi)^{\dagger}\circ v\circ(\phi\circ \psi)=(\psi^{\dagger}\circ\phi^{\dagger})\circ v\circ(\phi\circ\psi)\\ =\psi^{\dagger}\circ(\phi^{\dagger}\circ v\circ\phi)\circ\psi=(D(\psi)\circ D( \phi))(v),$ (9) $$

where the first equality comes from contravariant functoriality of $(-)^{\dagger}$ and the second equality comes from associativity of composition. Similarly, for preservation of identities, we have

$$ $D(\mathrm{id}_{Z})(v)\coloneqq\mathrm{id}_{Z}^{\dagger}\circ v\circ\mathrm{id} _{Z}=\mathrm{id}_{Z}\circ v\circ\mathrm{id}_{Z}=v,$ (10) $$

where the first equality again comes from functoriality of $(-)^{\dagger}$.

To verify the naturality of the product comparison, we need the following diagram to commute for all $\phi\colon X^{\prime}\to X$ and $\psi\colon Y^{\prime}\to Y$:

$$ ${{D(X)\times D(Y)}}$ ${{D(X^{\prime})\times D(Y^{\prime})}}$ ${{D(X\times Y)}}$ ${{D(X^{\prime}\times Y^{\prime})}}$ $\scriptstyle{\varphi_{X,Y}}$ $\scriptstyle{D(\phi)\times D(\psi)}$ $\scriptstyle{\varphi_{X^{\prime},Y^{\prime}}}$ $\scriptstyle{D(\phi\times\psi)}$ $$

For this, let $(v,w)\in D(X)\times D(Y)$ be arbitrary. Then, following the top path yields

$$ $\pi_{0}^{\dagger}\circ\phi^{\dagger}\circ v\circ\phi\circ\pi_{0}+\pi_{1}^{ \dagger}\circ\psi^{\dagger}\circ w\circ\psi\circ\pi_{1},$ (11) $$

while following the bottom path yields

$$ $\begin{array}[]{ll@{}ll}(\phi\times\psi)^{\dagger}\circ(\pi_{0}^{\dagger}\circ v \circ\pi_{0}+\pi_{1}^{\dagger}\circ w\circ\pi_{1})\circ(\phi\times\psi)&\\ =(\phi\times\psi)^{\dagger}\circ(\pi_{0}^{\dagger}\circ v\circ\pi_{0}\circ( \phi\times\psi)+\pi_{1}^{\dagger}\circ w\circ\pi_{1}\circ(\phi\times\psi))& \text{ Left additivity}\\ =(\phi\times\psi)^{\dagger}\circ\pi_{0}^{\dagger}\circ v\circ\pi_{0}\circ(\phi \times\psi)+(\phi\times\psi)^{\dagger}\circ\pi_{1}^{\dagger}\circ w\circ\pi_{1 }\circ(\phi\times\psi)&\text{ Additivity of linear maps}\\ =(\phi\times\psi)^{\dagger}\circ\pi_{0}^{\dagger}\circ v\circ\phi\circ\pi_{0}+ (\phi\times\psi)^{\dagger}\circ\pi_{1}^{\dagger}\circ w\circ\psi\circ\pi_{1}& \text{ Lemma \ref{lem:helpers}}\\ =(\pi_{0}\circ\phi\times\psi)^{\dagger}\circ v\circ\phi\circ\pi_{0}+(\pi_{1} \circ\phi\times\psi)^{\dagger}\circ w\circ\psi\circ\pi_{1}&\text{ Contravariant functoriality}\\ =(\phi\circ\pi_{0})^{\dagger}\circ v\circ\phi\circ\pi_{0}+(\psi\circ\pi_{1})^{ \dagger}\circ w\circ\psi\circ\pi_{1}&\text{ Lemma \ref{lem:helpers}}\\ =\pi_{0}^{\dagger}\circ\phi^{\dagger}\circ v\circ\phi\circ\pi_{0}+\pi_{1}^{ \dagger}\circ\psi^{\dagger}\circ w\circ\psi\circ\pi_{1},\end{array}$ (12) $$

as desired. The symmetric monoidal coherence axioms again follow from the commutative monoid structure of hom-sets.
∎

This is a generalization of the dynamics decorating functor defined in []. We can now define an analogous hypergraph category of open generalized dynamical systems.

###### Corollary 2.2.1 (Open Generalized Dynamics) .

Given a CRDC $\mathcal{C}$ such that $\mathsf{Lin}(\mathcal{C})$ has finite limits, there is a hypergraph category $\mathrm{\textnormal{{Dynam}}}_{\mathcal{C}}$ defined by the following data.

- •
Objects are the same as those of $\mathcal{C}$.
- •
Morphisms are (isomorphism classes of) open dynamical systems.
- •
Given two open systems $(X\xleftarrow{l_{1}}N\xrightarrow{r_{1}}Y,v\in D_{\mathcal{C}}(N))$ and $(Y\xleftarrow{l_{2}}M\xrightarrow{r_{2}}Z,w\in D_{\mathcal{C}}(M))$, their composite consists of the pullback span as in (5) together with the system obtained from the composite morphism
$\{*\}\cong\{*\}\times\{*\}\xrightarrow{v\times w}D_{\mathcal{C}}(N)\times D_{
\mathcal{C}}(M)\xrightarrow{\varphi_{N,M}}D_{\mathcal{C}}(N\times M)
\xrightarrow{D_{\mathcal{C}}\langle b_{0},b_{1}\rangle}D_{\mathcal{C}}(N\times
_{Y}M).$
(13)
- •
The identity on $X$ is the pair $(X=X=X,\{*\}\xrightarrow{\varphi_{0}}D_{\mathcal{C}}(*)\xrightarrow{D_{
\mathcal{C}}(!)}D_{\mathcal{C}}(X))$.
- •
The monoidal product on objects is their product in $\mathcal{C}$ while the monoidal product of morphisms $(X\leftarrow N\to Y,v)$ and $(W\leftarrow M\to Z,w)$ is
$(X\times W\leftarrow N\times M\to Y\times Z,\varphi_{N,M}(v,w)).$
(14)
- •
The hypergraph maps are inherited from those of $\mathrm{Span}(\mathsf{Lin}(\mathcal{C}))$ together with the identity decoration.

###### Proof.

This follows by applying Proposition 3.2 and Theorem 3.4 in [] to Theorem 2.2.
∎

###### Definition 2.2 .

###### Proof.

###### Corollary 2.2.1 (Open Generalized Dynamics) .

###### Proof.

### 2.3 Functoriality of Generalized Gradient Descent

The pieces are now in place for our main result. We would like to show that one can functorially relate a composite objective function defined in $\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}$ to a composite optimizer in $\mathrm{\textnormal{{Dynam}}}_{\mathcal{C}}$ that performs gradient descent on the given objective function. The framework of decorated spans allows the construction of such a hypergraph functor by specifying a monoidal natural transformation between the underlying decorating functors defining the domain and codomain hypergraph categories. Note that we work with continuous gradient flow systems; however, these results easily extend to discrete gradient descent systems as Euler’s method is known to be functorial for resource sharing dynamical systems [].

###### Theorem 2.3 .

Given an optimization domain $(\mathcal{C},R)$, there is a monoidal natural transformation
$\mathsf{gd}(\mathcal{C})\colon O^{R}_{\mathcal{C}}\Rightarrow D_{\mathcal{C}}$ with components $\mathsf{gd}(\mathcal{C})_{X}\colon O_{\mathcal{C}}^{R}(X)\to D_{\mathcal{C}}(X)$ defined by the generalized gradient descent optimization scheme, i.e.,

$$ $f\mapsto-\mathrm{\textnormal{{R}}}[f]_{1}.$ (15) $$

###### Proof.

Let $\mathsf{gd}\coloneqq\mathsf{gd}_{\mathcal{C}},O\coloneqq O_{\mathcal{C}}^{R},$ and $D\coloneqq D_{\mathcal{C}}$. For naturality, we need to verify that the following diagram commutes for all objects $X,Y\in\mathsf{Lin}(\mathcal{C})$ and linear maps $\phi\colon Y\to X$:

$$ ${{O(X)}}$ ${{O(Y)}}$ ${{D(X)}}$ ${{D(Y)}}$ $\scriptstyle{O(\phi)}$ $\scriptstyle{\mathsf{gd}_{X}}$ $\scriptstyle{\mathsf{gd}_{Y}}$ $\scriptstyle{D(\phi)}$ $$

Let $f\colon X\to R$ be an arbitrary objective. Then, following the top path yields the optimizer

$$ $-\mathrm{\textnormal{{R}}}[f\circ\phi]_{1}\coloneqq-\mathrm{\textnormal{{R}}}[ f\circ\phi]\circ\langle\mathrm{id}_{Y},1_{YR}\rangle.$ (16) $$

Likewise, following the bottom path yields

$$ $-\phi^{\dagger}\circ\mathrm{\textnormal{{R}}}[f]_{1}\circ\phi\coloneqq-\phi^{ \dagger}\circ\mathrm{\textnormal{{R}}}[f]\circ\langle\mathrm{id}_{X},1_{XR} \rangle\circ\phi.$ (17) $$

Now consider applying (16) to a generalized element $y\in Y$:

$$ $\begin{array}[]{ll@{}ll}-\mathrm{\textnormal{{R}}}[f\circ\phi]\circ\langle \mathrm{id}_{Y},1_{YR}\rangle(y)&=-\mathrm{\textnormal{{R}}}[f\circ\phi](y,1_{ R})\\ &=-\mathrm{\textnormal{{R}}}[\phi](y,\mathrm{\textnormal{{R}}}[f](\phi(y),1_{R }))&\text{ Chain rule}\textbf{ RD.5}\\ &=-\phi^{\dagger}(\mathrm{\textnormal{{R}}}[f](\phi(y),1_{R}))&\text{ Linearity.}\\ \end{array}$ (18) $$

Finally, applying (17) to the same element gives

$$ $-\phi^{\dagger}(\mathrm{\textnormal{{R}}}[f](\langle\mathrm{id}_{X},1_{XR} \rangle(\phi(y))))=-\phi^{\dagger}(\mathrm{\textnormal{{R}}}[f](\phi(y),1_{R}))$ (19) $$

as desired. To verify the transformation is monoidal, we must show that the following diagrams commute:

$$ ${{O(X)\times O(Y)}}$ ${{D(X)\times D(Y)}}$ ${{*}}$ ${{O(*)}}$ ${{O(X\times Y)}}$ ${{D(X\times Y)}}$ ${{D(*)}}$ $\scriptstyle{\mathsf{gd}_{X}\times\mathsf{gd}_{Y}}$ $\scriptstyle{\varphi_{X,Y}}$ $\scriptstyle{\varphi_{X,Y}}$ $\scriptstyle{\mathsf{gd}_{X\times Y}}$ $\scriptstyle{\varphi_{0}}$ $\scriptstyle{\varphi_{0}}$ $\scriptstyle{\mathsf{gd}_{*}}$ $$

Let $g\colon Y\to R$ be another arbitrary objective. Following the top path of the product comparison diagram yields the system

$$ $-\pi_{0}^{\dagger}\circ\mathrm{\textnormal{{R}}}[f]_{1}\circ\pi_{0}-\pi_{1}^{ \dagger}\circ\mathrm{\textnormal{{R}}}[g]_{1}\circ\pi_{1}$ (20) $$

while following the bottom path gives

$$ $\begin{array}[]{ll@{}ll}-\mathrm{\textnormal{{R}}}[f\circ\pi_{0}+g\circ\pi_{1} ]_{1}&=-\mathrm{\textnormal{{R}}}[f\circ\pi_{0}+g\circ\pi_{1}]\circ\langle \mathrm{id}_{X\times Y},1_{X\times Y,R}\rangle\\ &=-(\mathrm{\textnormal{{R}}}[f\circ\pi_{0}]+\mathrm{\textnormal{{R}}}[g\circ \pi_{1}])\circ\langle\mathrm{id}_{X\times Y},1_{X\times Y,R}\rangle&\textbf{ RD.1}\\ &=-\mathrm{\textnormal{{R}}}[f\circ\pi_{0}]\circ\langle\mathrm{id}_{X\times Y} ,1_{X\times Y,R}\rangle-\mathrm{\textnormal{{R}}}[g\circ\pi_{1}]\circ\langle \mathrm{id}_{X\times Y},1_{X\times Y,R}\rangle&\textbf{ LA.1}\\ &=-\mathrm{\textnormal{{R}}}[f\circ\pi_{0}]_{1}-\mathrm{\textnormal{{R}}}[g \circ\pi_{1}]_{1}.\end{array}$ (21) $$

These can be shown equivalent following the same reasoning as for naturality, namely by applying the chain rule. The unit comparison diagram commutes trivially as $D(*)$ is terminal.
∎

###### Corollary 2.3.1 (Functoriality of Generalized Gradient Flow) .

Given an optimization domain $(\mathcal{C},R)$ such that $\mathsf{Lin}(\mathcal{C})$ has finite limits, there is an identity on objects hypergraph functor $\mathsf{GD}_{\mathcal{C}}\colon\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}
\to\mathrm{\textnormal{{Dynam}}}_{\mathcal{C}}$ which takes an open objective $(X\xleftarrow{l}N\xrightarrow{r}Y,f\colon N\to R)$ to the open optimizer $(X\xleftarrow{l}N\xrightarrow{r}Y,-\mathrm{\textnormal{{R}}}[f]_{1}\colon N\to
N)$.

###### Proof.

This follows by applying Theorem 4.1 of [] to the monoidal natural transformation $\mathsf{gd}$.
∎

Taking stock, there are a many implications of the categories and functor just defined. For one, corollary 2.3.1 says we have the following equality given any composable pair $F\coloneqq(X\leftarrow N\to Y,f\colon N\to R),G\coloneqq(Y\leftarrow M\to Z,g
\colon M\to R)$ of open objectives:

$$ $\mathsf{GD}_{\mathcal{C}}(G\circ F)=\mathsf{GD}_{\mathcal{C}}(G)\circ\mathsf{ GD}_{\mathcal{C}}(F).$ (22) $$

In particular, this yields the following equivalence of open optimizers:

$$ $-\mathrm{\textnormal{{R}}}[f\circ\pi_{0}\circ\phi+g\circ\pi_{1}\circ\phi]_{1}= -\phi^{\dagger}\circ(\iota_{0}\circ\mathrm{\textnormal{{R}}}[f]_{1}\circ\pi_{0 }+\iota_{1}\circ\mathrm{\textnormal{{R}}}[g]_{1}\circ\pi_{1})\circ\phi,$ (23) $$

where we are using $\phi$ to denote the pairing $\langle b_{0},b_{1}\rangle$ from (5). Although these two optimizers are *extensionally* equivalent (i.e., they compute the same output for a given input), they are not *computationally* equivalent. In particular, the optimizer on the RHS of (23) is far better suited to implementation in a distributed computing environment as it can be interpreted with message passing semantics. Specifically, the distribute step is given by applying the linear map $\phi$ to the input followed by the application of the projections to send the desired components to the subsystems. The parallel computation step is given by computing the gradients of each objective with respect to their current local values, and the collect step is given by injecting the results into a single vector and applying the linear map $\phi^{\dagger}$.

Another benefit of open objectives and optimizers is the ability to easily specify composite objective functions using the graphical syntax of string diagrams. In the following section we will model hard parameter sharing for mulit-task learning in our category $\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}$. We do so to show that, rather than existing as an exercise in abstraction, our framework is applicable to a prevalent machine learning paradigm.

###### Theorem 2.3 .

###### Proof.

###### Corollary 2.3.1 (Functoriality of Generalized Gradient Flow) .

###### Proof.

## 3 Multitask Learning

Multitask learning (MTL) [] is a machine learning paradigm in which learning signals from different tasks are aggregated in order to train a single model. Precise definitions of what a task is vary, but for our purposes we will consider a task $\mathcal{T}_{i}\coloneqq(\mathcal{D}_{i},\mathcal{L}_{i})$ to be a pair of data set $\mathcal{D}_{i}$ and loss function $\mathcal{L}_{i}$. Usually (see []) the tasks $\mathcal{T}_{i}$ are supervised, so $\mathcal{D}_{i}\coloneqq\mathcal{X}_{i}\times\mathcal{Y}_{i}$ for input data $\mathcal{X}_{i}$ and labels $\mathcal{Y}_{i}$. The primary (or at least original) motivation of MTL is to exploit additional learning signals for improved model performance on a given “main” task, beyond what might be possible if the model were trained on any single task [].

A common approach (see []) to multitask learning is to leverage parameterised models where at least some parameters are trained on all tasks. Such an approach is called *hard parameter sharing* because the shared parameters are updated *directly* using gradient information from multiple tasks. This is in contrast to *soft parameter sharing* where the parameters are only *penalized* for differing from counterparts trained on other tasks. Theoretically, hard parameter sharing helps to avoid overfitting of shared parameters []. When the underlying model is a neural network, hard parameter sharing is often achieved by sharing the parameters of the first several layers while training a unique set of predictive final layers per task. However, recent work suggests that more flexible schemes may lead to improved performance []. With this motivation, we will now frame hard parameter sharing as a composite optimization problem that allows arbitrary parameters to be shared amongst arbitrary tasks.

For each task $\mathcal{T}_{i}$ in a collection of $N$ tasks, consider a neural network $f_{i}$ that has task-specific parameters $P_{i}$ and shared parameters $P_{0}$. We will call an assignment of parameters to real numbers $W_{i}\in\mathbb{R}^{P_{i}}$ the *weights* of a neural network, though note that the machine learning literature often uses the terms “weights” and “parameters” interchangeably. The composite optimization problem given by such a multitask learning setup (cf. [], (1)) is

$$ $\begin{array}[]{lll}\min\limits_{W_{0},\ldots,W_{N}}&\sum_{i=1}^{N}&\mathcal{L }^{D}_{i}(f_{i};W_{0},W_{i})\\ \end{array}$ (24) $$

where

$$ $\begin{array}[]{ll}\mathcal{L}^{D}_{i}(f_{i};W_{0},W_{i})&\coloneqq\frac{1}{| \mathcal{D}_{i}|}\sum_{d=1}^{|\mathcal{D}^{i}|}\mathcal{L}_{i}(f_{i}(\mathcal{ X}^{d}_{i};W_{0},W_{i}),\mathcal{Y}_{i}^{d})\\ \end{array}$ (25) $$

with $f_{i}(\mathcal{X}_{i}^{d};W_{0},W_{i})$ denoting the prediction of the network $f_{i}$ weighted by $(W_{0},W_{i})$ on input datum $\mathcal{X}_{i}^{d}$. We will now describe how problem 24 is a composite of open objectives, and therefore how to mechanically recover a distributed gradient descent algorithm for it.

### 3.1 Compositional Formulation

We will take our CRDC $\mathcal{C}$ to be Euc with objects euclidean spaces $\mathbb{R}^{N}$ and morphisms $f:\mathbb{R}^{N}\rightarrow\mathbb{R}^{M}$ the smooth maps between spaces. The monoidal product is given by $\times$, the standard product of Euclidean spaces. Each task’s loss is a smooth map $\mathcal{L}^{D}_{i}\colon\mathbb{R}^{P_{i}}\times\mathbb{R}^{P_{0}}\rightarrow
\mathbb{R}$, meaning our optimization domain is $(\mathrm{\textnormal{{Euc}}},\mathbb{R})$. Note that we have avoided some difficulty by treating each loss as an atomic map, as opposed to the composition of a neural network applied to data followed by a loss function. In the latter case, we would need to consider the set of parameters for the network in addition to the set of data points, and handle the subtleties of updating parameters while ignoring gradients w.r.t the input data. See [] for a thorough discussion of such concerns. Because the loss $\mathcal{L}^{D}_{i}$ is computed with respect to the entire data set $\mathcal{D}_{i}$ for a task, we have restricted ourselves to optimization methods that perform true gradient descent as opposed to the stochastic or mini-batch variants which are more common in the machine learning literature. We leave an extension of our work to the SGD setting for future efforts.

With the goal of specifying problem 24 as a morphism of $\mathrm{\textnormal{{Opt}}}_{\mathrm{\textnormal{{Euc}}}}^{\mathbb{R}}$, we first consider the following span that relates shared and task specific parameters for some $i$.

$$ $\mathbb{R}^{P_{0}}\xleftarrow{\pi_{0}}\mathbb{R}^{P_{0}}\times\mathbb{R}^{P_{i }}\xrightarrow{\pi_{1}}\mathbb{R}^{P_{i}}$ (26) $$

To promote 26 to a decorated span, we need only to specify an objective function $\mathbb{R}^{P_{0}}\times\mathbb{R}^{P_{i}}\rightarrow\mathbb{R}$. With $\mathcal{L}^{D}_{i}$ being the obvious choice, we construct a decorated span $F^{i}\coloneqq\left(\mathbb{R}^{P_{0}}\xleftarrow{\pi_{0}}\mathbb{R}^{P_{0}}
\times\mathbb{R}^{P_{i}}\xrightarrow{\pi_{1}}\mathbb{R}^{P_{i}},\mathcal{L}^{D
}_{i}\right)$. Each $F^{i}$ is a morphism in the hypergraph category $\mathrm{\textnormal{{Opt}}}_{\mathrm{\textnormal{{Euc}}}}^{\mathbb{R}}$, meaning that we may take the monoidal product $\bigotimes_{i=1}^{N}F^{i}$. By Cor. 2.1.1, specifically Eq. 7, we know that the monoidal product of our task specific decorated spans is as follows.

$$ $\bigotimes_{i=1}^{N}F^{i}=\left(\prod_{i=1}^{N}\mathbb{R}^{P_{0}}\xleftarrow{ \prod_{i=1}^{N}\pi^{i}_{0}}\prod_{i=1}^{N}(\mathbb{R}^{P_{0}}\times\mathbb{R}^ {P_{i}})\xrightarrow{\prod_{i=1}^{N}\pi_{1}^{i}}\prod_{i=1}^{N}\mathbb{R}^{P_{ i}},\varphi(\mathcal{L}^{D}_{1},\ldots,\mathcal{L}^{D}_{N})\right)$ (27) $$

Note that by Thm. 2.1 the decorating objective function is given by the comparison map:

$$ $\varphi(\mathcal{L}^{D}_{1},\ldots,\mathcal{L}^{D}_{N})\colon\prod_{i=1}^{N}O( \mathbb{R}^{P_{0}}\times\mathbb{R}^{P_{i}})\rightarrow O(\prod_{i=1}^{N} \mathbb{R}^{P_{0}}\times\mathbb{R}^{P_{i}}),$ (28) $$

where $\varphi(\mathcal{L}^{D}_{1},\ldots,\mathcal{L}^{D}_{N})$ maps the tuple of task specific losses to their sum:

$$ $(\mathcal{L}^{D}_{1},\ldots,\mathcal{L}^{D}_{N})\mapsto\sum_{i=1}^{N}\mathcal{ L}^{D}_{i}\circ\pi^{i},$ (29) $$

where $\pi^{i}$ denotes the $i$’th projection of $\prod_{i=1}^{N}(\mathbb{R}^{P_{0}}\times\mathbb{R}^{P_{i}})$. We now wish to specify that the parameter space $\mathbb{R}^{P_{0}}$ is shared among all $N$ loss functions. This is accomplished as follows.

Recall that in a hypergraph category like $\mathrm{\textnormal{{Opt}}}_{\mathrm{\textnormal{{Euc}}}}^{\mathbb{R}}$, we have comultiplication maps

$$ $\delta_{p}\coloneqq\left(\mathbb{R}^{p}\xleftarrow{\mathrm{id}_{p}}\mathbb{R}^ {p}\xrightarrow{\Delta}\mathbb{R}^{p}\times\mathbb{R}^{p},0\colon\mathbb{R}^{p }\rightarrow\mathbb{R}\right),$ (30) $$

where $0\colon\mathbb{R}^{p}\to\mathbb{R}$ is the “empty” decoration, i.e., the constant 0 objective function $0(v)=0$. The map $\Delta\colon\mathbb{R}^{p}\to\mathbb{R}^{p}\times\mathbb{R}^{p}$ is the duplication map $v\mapsto(v,v)$ built-in to any Cartesian category. Thus, to copy the single parameter space $\mathbb{R}^{P_{0}}$ to all tasks, we simply apply $\delta_{P_{0}}$ $N$ times. We denote this $N$-fold application $\delta_{P_{0}}^{N}$. To recover problem 24, all we must do is precompose our monoidal product of task specific objectives with the $N$-fold copy, yielding the morphism in $\mathrm{\textnormal{{Opt}}}_{\mathrm{\textnormal{{Euc}}}}^{\mathbb{R}}$

$$ $\begin{array}[]{ll}\Bigl{(}\left(\bigotimes_{i=1}^{N}F^{i}\right)\circ\delta_{ P_{0}}^{N},O(\langle b_{0},b_{1}\rangle)\circ\left(0+\sum_{i=1}^{N}\mathcal{L} ^{D}_{i}\circ\pi^{i}\right)\Bigr{)}&=\\ \Bigl{(}\left(\bigotimes_{i=1}^{N}F^{i}\right)\circ\delta_{P_{0}}^{N},\mathcal {L}^{D}\colon(W_{0},W_{1},\dots,W_{N})\mapsto\sum_{i=1}^{N}\mathcal{L}^{D}_{i} (W_{0},W_{i})\Bigr{)}\par\end{array}$ (31) $$

which can be pictured diagramatically as:

$$ ${{\mathbb{R}^{P_{0}}\times\prod_{i=1}^{N}\mathbb{R}^{P_{i}}}}$ ${{\mathbb{R}^{P_{0}}}}$ ${{\mathbb{R}^{P_{0}}\times\left(\prod_{i=1}^{N}\mathbb{R}^{P_{0}}\times\mathbb {R}^{P_{i}}\right)}}$ ${{\prod_{i=1}\mathbb{R}^{P_{0}}\times\mathbb{R}^{P_{i}}}}$ ${{\mathbb{R}^{P_{0}}}}$ ${{\prod_{i=1}^{N}\mathbb{R}^{P_{0}}}}$ ${{\prod_{i=1}\mathbb{R}^{P_{i}}}}$ $\scriptstyle{\mathrm{id}}$ $\scriptstyle{\Delta}$ $\scriptstyle{\prod_{i=1}^{N}\pi_{0}^{i}}$ $\scriptstyle{\prod_{i=1}^{N}\pi_{1}^{i}}$ $\scriptstyle{{\hat{\pi}_{1}}}$ $\scriptstyle{{\hat{\pi}_{0}}}$ $\scriptstyle{{b_{1}}}$ $\scriptstyle{{b_{0}}}$ $\scriptstyle{{\langle b_{0},b_{1}\rangle}}$ $$

In Eq. 31, the computation of the pullback has effectively filtered the multiple copies of $\mathbb{R}^{P_{0}}$ and left us with the loss we want. Note that the objective function 0 0 of the comultiplication morphism does not contribute to the overall objective.

### 3.2 Distributed Optimization via Functoriality of Gradient Descent

Figure: Algorithm 1 Distributed MTL

Having represented problem 24 as a morphism in $\mathrm{\textnormal{{Opt}}}_{\mathrm{\textnormal{{Euc}}}}^{\mathbb{R}}$, we can create a gradient descent based optimizer by applying Cor. 2.3.1. In particular we aim for a distributed gradient descent scheme. Exploiting the functoriality of $\mathsf{GD}_{\mathrm{\textnormal{{Euc}}}}$, we prefer to compute the gradient for each morphism separately and then compose systems. The generalized gradient functor is illustrated in Section 1, Figure 1.

Algorithm 1 describes the distributed scheme for multitask learning derived by applying $\mathsf{GD}_{\mathcal{C}}$ to the individual subproblems and composing the resulting dynamical systems. Note that there is value in distributing computation this way as computing the gradients of weights will typically involve non-trivial backpropogation over potentially large networks.
With this, we have successfully recovered hard parameter sharing for multitask learning and shown how to make a gradient descent optimizer that can be implemented in a distributed fashion.

## 4 Discussion and Future Work

Taking stock, it is fair to wonder what has been gained by specifying parameter sharing via decorated spans. One classic argument for the utility of applied category theory is the tight link between graphical problem depiction and its realization in a mathematical or computational data structure. Figure 1 contains two string diagrams which visually capture the structure of multitask learning objectives and optimizers respectively, enabling better communication of models between machine learning practitioners. Moreover, the relative ease with which one can describe compositional objectives supports more complex problem structures. Objectives with traditional regularization terms can be trivially described as compositional objectives, and soft parameter sharing should be fairly easily accommodated via the addition of penalization terms. More exotic hierarchical parameter sharing (see e.g. []) should be readily modeled using our approach to MTL. There is a computational benefit too; distributed optimization schemes become easier to implement with a rigorous definition of problem structure. Finally the generality of the techniques discussed allows the parameter sharing pattern of the given morphism to be applied to settings beyond smooth real-valued functions and traditional gradient descent.

As for the general framework defined in Sec. 2, we believe there are interesting connections with previous work on CRDCs. Perhaps the most obvious direction for future work would be to understand the connections with the $\operatorname{Para}$ and $\mathrm{Lense}$ constructions of []. In particular, we hope that a unified treatment of both the compositional structure of models and the compositional structure of objectives can be achieved, perhaps using tools from double category theory.
Naturally another direction of future work is to find more examples of CRDCs and study not only learning problems in those settings, but other optimization problems such as those arising in operations research.

## 5 Conclusion

In summary, we have presented hypergraph categories $\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}$ and $\mathrm{\textnormal{{Dynam}}}_{\mathcal{C}}$ of open objectives and dynamical systems respectively, and shown that there exists a hypergraph functor between them that maps a generalized objective to its corresponding gradient descent optimizer. Both $\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}$ and $\mathrm{\textnormal{{Dynam}}}_{\mathcal{C}}$ are constructed with respect to an underlying Cartesian reverse derivative category, allowing for composition of optimization problems that are defined on a general optimization domain. Such compositionality induces a template for distributed optimization, and provides an accompanying graphical syntax that facilitates communication and sharing. Lastly, to provide evidence that existing machine learning paradigms can be described by our compositional framework, we described parameter sharing models for multitask learning using the hypergraph structure of $\mathrm{\textnormal{{Opt}}}_{\mathcal{C}}^{R}$, and highlighted the distributed gradient descent algorithm induced by this structure.

## Appendix A Generalized Optimization

In this section, we recall the necessary definitions and concepts from Cartesian reverse derivative categories [] and generalized optimization [].

###### Definition A.1 (Definition 3.1 in [ 18 ] ) .

A Cartesian left-additive category is a Cartesian category $\mathcal{C}$ with terminal object $*$ in which for any pair of objects $X,Y$, the hom-set $\mathcal{C}(X,Y)$ is a commutative monoid with addition operation $+$ and additive identities $0_{X,Y}\colon X\to Y$. In addition, the following axioms must be satisfied.

- •
LA.1 For any morphisms $h\colon X\to Y$ and $f,g\colon Y\to Z$, we have
$(f+g)\circ h=(f\circ h)+(g\circ h)\colon X\to Z\text{ and }0_{Y,Z}\circ h=0_{X
,Z}\colon X\to Z.$
(32)
- •
LA.2 For any projection map $\pi_{i}\colon Y\to Z$ and morphisms $f,g\colon X\to Y$ we have
$\pi_{i}\circ(f+g)=(\pi_{i}\circ f)+(\pi_{i}\circ g)\colon X\to Z\text{ and }
\pi_{i}\circ 0_{X,Y}=0_{X,Z}\colon X\to Z.$
(33)

The additive identity of $\mathcal{C}(*,X)$ is denoted $0_{X}$.

###### Rough Definition A.2 (Definition 3.3 in [ 18 ] ) .

A Cartesian differential category is a Cartesian left-additive category $\mathcal{C}$ equipped with a Cartesian derivative combinator D taking morphisms $f\colon X\to Y$ to morphisms $\mathrm{\textnormal{{D}}}[f]\colon X\times X\to Y$. This combinator must satisfy several axioms requiring that it behave like a standard forward derivative.

###### Rough Definition A.3 (Definition 3.2 in [ 18 ] ) .

A Cartesian reverse derivative category is a Cartesian left-additive category $\mathcal{C}$ equipped with a Cartesian reverse derivative combinator R taking morphisms $f\colon X\to Y$ to morphisms $\mathrm{\textnormal{{R}}}[f]\colon X\times Y\to X$. This combinator must satisfy several axioms requiring that it behave like a standard reverse derivative. In particular, we will make use of the following axioms in this paper.

- •
RD.1 $\mathrm{\textnormal{{R}}}[f+g]=\mathrm{\textnormal{{R}}}[f]+\mathrm{
\textnormal{{R}}}[g]$ and $\mathrm{\textnormal{{R}}}[0]=0$.
- •
RD.5 $\mathrm{\textnormal{{R}}}[g\circ f]=\mathrm{\textnormal{{R}}}[f]\circ(\mathrm{
id}\times\mathrm{\textnormal{{R}}}[g])\circ\langle\pi_{0},\langle f\circ\pi_{0
},\pi_{1}\rangle\rangle$.

Intuitively, RD.1 says that the reverse derivative combinator must be additive while RD.5 is the chain rule cast in the abstract setting of CRDCs. Every Cartesian reverse derivative category (CRDC) is also a Cartesian differential category (Theorem 16 in []) by defining the derivative combinator of a morphism $f\colon X\to Y$ as

$$ $\mathrm{\textnormal{{D}}}[f]=\pi_{1}\circ\mathrm{\textnormal{{R}}}[\mathrm{ \textnormal{{R}}}[f]]\circ(\langle\mathrm{id}_{X},0_{X,Y}\rangle\times\mathrm{ id}_{X})\colon X\times X\to Y.$ (34) $$

This allows us to refer to the Cartesian derivative combinator of a CRDC. Of crucial importance in the sequel is the notion of a linear map in CRDC.

###### Definition A.4 .

A linear map in a CRDC $\mathcal{C}$ is a morphism $f$ such that $\mathrm{\textnormal{{D}}}[f]=f\circ\pi_{1}$. The linear maps in $\mathcal{C}$ form a subcategory of $\mathcal{C}$ which we denote $\mathsf{Lin}(\mathcal{C})$. By Proposition 24 in [], $\mathsf{Lin}(\mathcal{C})$ is a ${\dagger}$-category with finite ${\dagger}$-biproducts. Moreover, $\mathrm{\textnormal{{R}}}[f]=f^{\dagger}\circ\pi_{1}$ and $(f+g)^{\dagger}=f^{\dagger}+g^{\dagger}$.

There are two canonical examples of CRDCs in the literature.

- 1.
The CRDC Euc has Euclidean spaces as objects and smooth functions as morphisms. The terminal object is $\mathbb{R}^{0}$ and the reverse derivative of a smooth function $f\colon\mathbb{R}^{m}\to\mathbb{R}^{n}$ is
$\mathrm{\textnormal{{R}}}[f](x,x^{\prime})\coloneqq J_{f}(x)^{T}x^{\prime},$
where $J_{f}(x)$ is the Jacobian of $f$ evaluated at $x$. The linear maps in Euc are the usual linear maps between vector spaces, with the dagger structure corresponding to dual maps.
- 2.
Given a commutative ring $r$, the CRDC $\textbf{Poly}_{r}$ has natural numbers as objects. A morphism from $n$ to $m$ is an $m$-tuple of polynomials over $r$, each in $n$ variables. The terminal object is 0 0 and the reverse derivative of a polynomial $P(x)=(p_{1}(x),\dots,p_{m}(x))$ is
$\mathrm{\textnormal{{R}}}[P](x,x^{\prime})\coloneqq\Bigl{(}\sum_{i=1}^{m}\frac
{\partial p_{i}}{\partial x_{1}}(x)x_{i}^{\prime},\dots,\sum_{i=1}^{m}\frac{
\partial p_{i}}{\partial x_{n}}(x)x_{i}^{\prime}\Bigr{)},$
where $\frac{\partial p_{i}}{\partial x_{j}}(x)$ denotes the formal derivative of $p_{i}$ in $x_{j}$ at $x$. The linear maps are polynomials $P(x)=(p_{1}(x),\dots,p_{m}(x))$ where each $p_{i}(x)$ is of the form $p_{i}(x)=\sum_{i=1}^{n}r_{i}x_{i}$ for $r_{i}\in r$ [].

###### Definition A.5 (Definition 3.4 in [ 18 ] ) .

An optimization domain is a pair $(\mathcal{C},R)$ where $\mathcal{C}$ is a CRDC and $R$ is an object of $\mathcal{C}$ equipped with the following additional structures.

- •
Each morphism $f\colon X\to Y$ in $\mathcal{C}$ has an additive inverse $-f$.
- •
Each hom-set $\mathcal{C}(*,X)$ out of the terminal object is equipped with a multiplication operation $fg$ and a multiplicative identity $1_{X}\colon*\to X$ to form a commutative ring with the left additive structure.
- •
The hom-set $\mathcal{C}(*,R)$ is totally ordered to form an ordered commutative ring.

Given a unique map $!_{X}\colon X\to*$ into the terminal object, the map $1_{Y}\circ!_{X}\colon X\to Y$ is denoted $1_{X,Y}$.
An objective in $(\mathcal{C},R)$ is a morphism $\ell\colon X\to R$ in $\mathcal{C}$.

###### Definition A.6 (Definition 3.6 in [ 18 ] ) .

Given an optimization domain $(\mathcal{C},R)$, the generalized gradient of an objective $\ell\colon X\to R$ is $\mathrm{\textnormal{{R}}}[\ell]_{1}\colon X\to X$ defined by

$$ $\mathrm{\textnormal{{R}}}[\ell]_{1}\coloneqq\mathrm{\textnormal{{R}}}[\ell] \circ\langle\mathrm{id}_{X},1_{X,R}\rangle.$ (35) $$

The following are examples given in [] of optimization domains and generalized gradients.

- 1.
The standard domain is $(\mathrm{\textnormal{{Euc}}},\mathbb{R})$. Objectives are smooth functions $\ell\colon\mathbb{R}^{n}\to\mathbb{R}$ and the gradient of $\ell\colon\mathbb{R}^{n}\to\mathbb{R}$ is $\nabla\ell\colon\mathbb{R}^{n}\to\mathbb{R}^{n}$.
- 2.
The $r$-polynomial domain is $(\textbf{Poly}_{r},1)$. Objectives are polynomials $\ell\colon n\to 1$ and the gradient of $\ell\colon n\to 1$ is $\Bigl{(}\frac{\partial\ell}{\partial x_{1}}(x),\dots,\frac{\partial\ell}{
\partial x_{n}}(x)\Bigr{)}$, where these are again formal derivatives.

###### Lemma A.7 .

Let $\phi\colon X\to Y$ and $\psi\colon X^{\prime}\to Y^{\prime}$ be arbitrary morphisms in $\mathrm{Lin}(\mathcal{C})$. Then the following equations hold:

- 1.
$(\phi\times\psi)^{\dagger}=\phi^{\dagger}\times\psi^{\dagger}$,
- 2.
$\pi_{0}\circ(\phi\times\psi)=\phi\circ\pi_{0}$,
- 3.
$\pi_{1}\circ(\phi\times\psi)=\psi\circ\pi_{1}$.

###### Proof.

For (1), we have

$$ $\begin{array}[]{ll@{}ll}(\phi\times\psi)^{\dagger}&=\langle\phi\circ\pi_{0}, \psi\circ\pi_{1}\rangle^{\dagger}&\\ &=(\iota_{0}\circ\phi\circ\pi_{0}+\iota_{1}\circ\psi\circ\pi_{1})^{\dagger}& \text{ Lemma 3 in \cite[cite]{[\@@bibref{}{cockett2019reverse}{}{}]}}\\ &=(\iota_{0}\circ\phi\circ\pi_{0})^{\dagger}+(\iota_{1}\circ\psi\circ\pi_{1})^ {\dagger}&\text{ Property of ${\dagger}$-biproducts}\\ &=\pi_{0}^{\dagger}\circ\phi^{\dagger}\circ\iota_{0}^{\dagger}+\pi_{1}^{ \dagger}\circ\psi^{\dagger}\circ\iota_{1}^{\dagger}&\text{ Contravariant functoriality}\\ &=\iota_{0}\circ\phi^{\dagger}\circ\pi_{0}+\iota_{1}\circ\psi^{\dagger}\circ \pi_{1}&\text{ Lemma 21 in \cite[cite]{[\@@bibref{}{cockett2019reverse}{}{}]}} \\ &=\langle\phi^{\dagger}\circ\pi_{0},\psi^{\dagger}\circ\pi_{1}\rangle=\phi^{ \dagger}\times\psi^{\dagger}.\end{array}$ (36) $$

For (2), let $(x,x^{\prime})\in X\times X^{\prime}$ be generalized elements. Then

$$ $\pi_{0}\circ(\phi\times\psi)(x,x^{\prime})=\phi(x)=\phi(\pi_{0}(x,x^{\prime})).$ (37) $$

A symmetric argument holds for (3).
∎

###### Definition A.1 (Definition 3.1 in [ 18 ] ) .

###### Rough Definition A.2 (Definition 3.3 in [ 18 ] ) .

###### Rough Definition A.3 (Definition 3.2 in [ 18 ] ) .

###### Definition A.4 .

###### Definition A.5 (Definition 3.4 in [ 18 ] ) .

###### Definition A.6 (Definition 3.6 in [ 18 ] ) .

###### Lemma A.7 .

###### Proof.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. You introduced the concept of the generalization ratio (GR) in this paper. Could you clarify the functional difference between V_KL for different e_i sets and GR?
2. As a novel and important criterion, could you elaborate on the properties of GR and why it serves as a strong evaluation metric?
3. Could you design some experiments that demonstrate the advantages of using GR instead of V_KL as Ye et al. (2021) also provided a model selection criterion?
soundness: 3
strengths: 1. They designed an algorithm focused on enhancing the generalization ability of gradient descent, leveraging the concept of a generalization ratio and insights from reinforcement learning.
2. They experimentally demonstrate that the algorithm outperforms traditional gradient descent.
summary: This paper proposes a novel framework for analyzing out-of-distribution (OOD) generalization, integrating a generalization ratio to quantify generalization capacity and a generalization decision process (GDP) to link losses across seen and unseen domains. This framework enables backpropagation-based training without model selection, using a combination of informativeness and variation in generalization to derive a generalization inequality, with experiments showcasing its effectiveness in improving generalization ability.
weaknesses: 1. The title is intriguing. Would it be clearer to use 'Gradient Descent with the Generalization Ratio: Enhancing.......' instead?
2. In my view, GGD is proposed based on reinforcement learning and the generalization ratio. However, the motivation behind the algorithm design could be clarified further.
3. There is considerable overlap with the paper by Ye et al. (2021). It might help to elaborate more on the ideas behind the algorithm's design.
4. The writing could be improved, as there are several grammatical errors throughout.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: None.
soundness: 3
strengths: The paper discusses the generalization problem mathematically rigorously. The ideas of generalization ratio and generalization decision process are novel and the introduction of reinforcement learning framework is unique.
summary: The paper tackles the out-of-distribution (OOD) generalization problem. The main ideas are to consider the distribution of features invariant to domain shifts and to formulate them as the ideal feature matrix function, generalized model, and non-generalized model.  The paper also introduces the variation, the informativeness, and the generalization decision process (GDP) to propose a new algorithm called the generalization gradient descent (GGD) based on the framework of reinforcement learning. The GGD performs better than the traditional gradient descent in their experiment.
weaknesses: It is not clear whether the framework is practical or not since the definition of the ideal feature matrix function that the posterior distributions of y are the same seems too strong to exist in real problems. The reviewer recommends showing some examples. The experiment is rather small and specific, that is, it compares only GGD and TGD. The effectiveness of the proposed algorithm should have been confirmed with a larger number of problems. In addition, this work strongly depends on the previous work (Ye+ 2021) and its difference seems not large since many of the concepts were introduced there.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 1
questions: Please see the weaknesses section.
soundness: 2
strengths: This is an ambitious paper.
summary: The paper proposes a variation of the framework of (Ye et al., Neurips 2021) to analyze out-of-distribution generalization, then propose a training algorithm that uses these concepts to offer cross-domain generalization performance.

Ye et al. (2021) studies how to use a small set of domains Eavail to construct a classifier that works across a larger set of domains Eall that we obviously are not expected to know in advance. They provide bounds that depends on an essentially unknowable "expansion function" that measures how features qualities expand from Eavail to Eall.  This paper instead assumes that all domains are available in a set Eavail that is partitioned into training domains Etrain and validation domains Eval. There is no Eall anymore, and one wonder why one should not train using all the known domains of Eavail (which are all the domains that matter) rather than a subset.  Other differences relative to (Ye et al. 2021) include an attempt to talk about features in all network layers (instead of just the penultimate one).
weaknesses: * This paper instead assumes that all domains are available in a set Eavail that is partitioned into training domains Etrain and validation domains Eval. There is no Eall anymore, and one wonder why one should not train using all the known domains of Eavail (which are all the domains that matter) rather than a subset. 

* There is something disturbing in equation (3) [Thm 2.1].   Remember that W(X) is a dxN matrix.  An expression of the form P(W(X^{e_i})|...) should be normalized over the values taken by this matrix, not over the i indice representing the domains.  Either the notation is inadequate (meaning that the authors mean something else) or the theorem is wrong.

* Through definition 3.8, algorithm GGD depends on all the domain distributions in Eavail and not just the domain distributions in Etra.  Since one wishes to be robust in Eavail  (unlike Ye et al. (2021), there is no Eall here], why not use Eavail directly.  The only justification that comes to my mind would be when Etra is small and Eval is much larger.  But evaluating GR_{KL}(W, Eavail) does not seem cheap. 

* I must recognize that I find all this very unclear.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 2
questions: Theorem 4.2 provides an upper bound of the generalization error. Is the generalization upper bound tight? Additionally, I noticed that the upper bound of the generalization error is composed of both the generalization ratio (GR) and the training error. The optimization in the RL section is performed with respect to GR. However, as the sample size n increases, which part—GR or training error—will dominate the generalization error? If there are some theoretical results regarding this, I believe the results could be more complete.

In the RL section, the paper uses Q-learning to adjust the training set. The state space is the set of G, and the action space is the selection of training samples. However, I’m a bit confused: for a highly complex model, can its state really be represented by a one-dimensional scalar G (Generalization Ratio)? The theory here seems a bit not solid enough.

Overall, I think the idea of this paper is interesting, but I still have these concerns. I would be happy if the authors could address my questions. Thank you.
soundness: 2
strengths: This paper presents a novel approach to gradient descent. It proposes an innovative method that selects the training samples for each step using a RLstrategy, which enhances the overall gradient descent process. This approach demonstrates a creative combination of RL and optimization techniques, potentially leading to improvements in training efficiency and model performance. The integration of RL for sample selection is a unique aspect that sets this work apart from traditional gradient descent methods, which is quite interesting.
summary: This paper introduces a new way of performing gradient descent.  Firstly, this paper defines a certain kind of generalization capability. Then, this paper proposes a method  which selects training sample of each step in a reinforcement learning (RL) manner, thereby improving the gradient descent process.
weaknesses: Line 149: eature -> feature

Line 176: The definition of $W^*$ is a little confusing. Do you mean $\forall W^* \in \mathcal{W}^*$ or $\exists W^* \in \mathcal{W}^*$?

Line 505: The font in the image is too small to be readable.

  

In Definition 3.1, this paper introduced the definition of the Ideal feature matrix function. However, I still don’t have an intuitive sense of it, and it feels like it’s missing some explanation here. I think it would be better if the reasoning behind this definition or references to relevant literature could be provided.