is defined for $p \ge 1$. The 2-Wasserstein distance is particularly significant due to its geometric interpretation and its applications in optimal transport theory, statistics, and machine learning.

Given two probability measures $\mu$ and $\nu$ on $X$, the 2-Wasserstein distance $W_2(\mu, \nu)$ is defined as the infimum of the expected squared distance between random variables $X$ and $Y$, where $X$ is distributed according to $\mu$ and $Y$ is distributed according to $\nu$, over all possible joint distributions $\pi$ of $(X, Y)$ with marginals $\mu$ and $\nu$. Mathematically, this is expressed as:

$$W_2(\mu, \nu) = \left( \inf_{\pi \in \Pi(\mu, \nu)} \int_{X \times X} d(x, y)^2 d\pi(x, y) \right)^{1/2}, \qquad (6.57)$$

where $\Pi(\mu, \nu)$ denotes the set of all joint probability measures $\pi$ on $X \times X$ with marginals $\mu$ and $\nu$, i.e., for any measurable sets $A, B \subset X$, $\pi(A \times X) = \mu(A)$ and $\pi(X \times B) = \nu(B)$. The distance $d(x, y)$ is the metric on the space $X$, and the integral $\int_{X \times X} d(x, y)^2 d\pi(x, y)$ represents the expected squared distance under the coupling $\pi$. The 2-Wasserstein distance can also be expressed in terms of the cumulative distribution functions (CDFs) of $\mu$ and $\nu$ when $X = \mathbb{R}$. Let $F_\mu$ and $F_\nu$ be the CDFs of $\mu$ and $\nu$, respectively. Then, the 2-Wasserstein distance is given by:

$$W_2(\mu, \nu) = \left( \int_0^1 |F_\mu^{-1}(q) - F_\nu^{-1}(q)|^2 dq \right)^{1/2}, \qquad (6.58)$$

where $F_\mu^{-1}$ and $F_\nu^{-1}$ are the quantile functions (inverse CDFs) of $\mu$ and $\nu$, respectively. This formulation is particularly useful in one-dimensional settings, as it reduces the problem to computing the integral of the squared difference between the quantile functions. In the case where $\mu$ and $\nu$ are absolutely continuous with respect to the Lebesgue measure, with probability density functions $f_\mu$ and $f_\nu$, the 2-Wasserstein distance can be related to the optimal transport map $T$ that pushes $\mu$ forward to $\nu$, i.e., $T_\#\mu = \nu$. The map $T$ minimizes the transport cost:

$$W_2(\mu, \nu) = \left( \int_X |x - T(x)|^2 d\mu(x) \right)^{1/2}. \qquad (6.59)$$

The optimal transport map $T$ is often characterized by the Monge-Ampère equation, which relates the densities $f_\mu$ and $f_\nu$ through the Jacobian determinant of $T$:

$$f_\mu(x) = f_\nu(T(x)) \cdot |\det(\nabla T(x))|. \qquad (6.60)$$

In higher dimensions, the 2-Wasserstein distance is more challenging to compute, but it can be approximated using numerical methods such as linear programming, entropic regularization, or Sinkhorn iterations. The dual formulation of the 2-Wasserstein distance, derived from Kantorovich duality, provides another perspective:

$$W_2(\mu, \nu)^2 = \sup_{\phi, \psi} \left( \int_X \phi(x) d\mu(x) + \int_X \psi(y) d\nu(y) \right), \qquad (6.61)$$

where the supremum is taken over all pairs of continuous functions $\phi$ and $\psi$ on $X$ satisfying the constraint:

$$\phi(x) + \psi(y) \le d(x, y)^2 \quad \forall x, y \in X. \qquad (6.62)$$

This dual formulation is particularly useful in theoretical analyses and provides a connection to convex optimization. The 2-Wasserstein distance has several important properties, including symmetry, non-negativity, and the triangle inequality, making it a true metric on the space of probability measures with finite second moments. It is also sensitive to the geometry of the underlying space $X$, as it incorporates the metric $d(x, y)$ directly into its definition. This geometric sensitivity makes it particularly useful in applications such as image processing, where the spatial arrangement of
pixels is crucial.

In summary, the 2-Wasserstein distance is a powerful tool for comparing probability distributions, with deep connections to optimal transport theory, convex analysis, and geometry. Its rigorous mathematical formulation and rich theoretical properties make it a cornerstone of modern probability and statistics.

## 6.5 Optimal Transport (OT)-based priors in Bayesian Deep Learning

The incorporation of Optimal Transport (OT)-based priors into Bayesian Deep Learning frameworks ensures that posterior distributions maintain smooth and stable representations by leveraging the rigorous mathematical properties of the Wasserstein distance and its associated geometric structure. This approach is grounded in measure theory, functional analysis, and convex optimization, providing a robust foundation for uncertainty estimation tasks. Below, we provide a deeper and more mathematically rigorous exposition.

Let $(X, d)$ be a complete separable metric space, and let $P_p(X)$ denote the space of probability measures on $X$ with finite $p$-th moments. For $P, Q \in P_p(X)$, the $p$-Wasserstein distance $W_p(P, Q)$ is defined as:

$$W_p(P, Q) = \left( \inf_{\pi \in \Pi(P, Q)} \int_{X \times X} d(x, y)^p d\pi(x, y) \right)^{1/p} \qquad (6.63)$$

where $\Pi(P, Q)$ is the set of all couplings (joint distributions) with marginals $P$ and $Q$. The Wasserstein distance induces a metric on $P_p(X)$, endowing it with a geometric structure that is sensitive to the underlying topology of $X$. This sensitivity ensures that the Wasserstein distance captures not only global but also local differences between distributions, making it particularly suitable for regularizing posterior distributions in Bayesian inference. In Bayesian Deep Learning, the prior $p(\theta)$ and the posterior $q(\theta)$ are probability measures over the parameter space $\Theta$. To ensure smoothness and stability of the posterior, we introduce an OT-based regularization term into the variational inference objective. Specifically, we augment the evidence lower bound (ELBO) with a Wasserstein penalty:

$$L_{\text{ELBO}}(q) = \mathbb{E}_{q}(\theta) [\log p(D | \theta)] - \text{KL}(q(\theta) \parallel p(\theta)) + \lambda W_p(p(\theta), q(\theta)) \qquad (6.64)$$

where $\lambda > 0$ is a regularization hyperparameter, and $\text{KL}(q(\theta) \parallel p(\theta))$ is the Kullback-Leibler divergence. The Wasserstein term $W_p(p(\theta), q(\theta))$ enforces geometric consistency between the prior and the posterior, ensuring that the posterior does not deviate excessively from the prior in a manner that respects the underlying metric structure of $\Theta$. The smoothness of the posterior is guaranteed by the properties of the Wasserstein distance. Specifically, the Wasserstein distance is Lipschitz-continuous with respect to perturbations in the distributions. For any $P, Q, R \in P_p(X)$, the triangle inequality holds:

$$W_p(P, Q) \le W_p(P, R) + W_p(R, Q) \qquad (6.65)$$

This inequality ensures that small changes in the data distribution $D$ or the prior $p(\theta)$ lead to proportionally small changes in the posterior $q(\theta)$, thereby promoting stability. Furthermore, the Wasserstein distance is convex in its arguments, which facilitates efficient optimization and guarantees the existence of a unique minimizer under appropriate conditions. To compute the Wasserstein distance in practice, we often employ entropic regularization, which transforms the OT problem into a strictly convex optimization problem. The entropic regularized Wasserstein distance $W_{p,\epsilon}(P, Q)$ is defined as:

$$W_{p,\epsilon}(P, Q) = \inf_{\pi \in \Pi(P, Q)} \int_{X \times X} d(x, y)^p d\pi(x, y) + \epsilon H(\pi) \qquad (6.66)$$
where $H(\pi) = \int_{X \times X} \pi(x, y) \log \pi(x, y) dx dy$ is the entropy of the coupling $\pi$, and $\epsilon > 0$ is the regularization parameter. The entropic regularization ensures that the optimal coupling $\pi^*$ is unique and has the form:

$$ \pi^*(x, y) = \exp \left( \frac{f(x) + g(y) - d(x, y)^p}{\epsilon} \right) \qquad (6.67) $$

where $f$ and $g$ are dual potentials that satisfy the Schrödinger system of equations. This formulation allows for efficient computation using the Sinkhorn-Knopp algorithm, which iteratively updates the dual potentials $f$ and $g$. The stability and smoothness of the posterior are further reinforced by the dual formulation of the Wasserstein distance. By the Kantorovich-Rubinstein duality, the 1-Wasserstein distance can be expressed as:

$$ W_1(P, Q) = \sup_{f \in \text{Lip}_1(X)} \left| \int_X f(x) dP(x) - \int_X f(x) dQ(x) \right| \qquad (6.68) $$

where $\text{Lip}_1(X)$ is the set of 1-Lipschitz functions on $X$. This dual formulation highlights the role of Lipschitz continuity in controlling the smoothness of the posterior, as the Wasserstein distance penalizes functions that vary too rapidly.

In summary, OT-based priors ensure that posterior distributions maintain smooth and stable representations by leveraging the geometric and analytical properties of the Wasserstein distance. The Wasserstein distance provides a rigorous framework for regularizing Bayesian inference, ensuring that the posterior remains close to the prior in a geometrically meaningful way. This approach is computationally tractable due to entropic regularization and the Sinkhorn-Knopp algorithm, and it is theoretically grounded in measure theory, convex analysis, and functional analysis. The resulting Bayesian framework is robust to noise and distributional shifts, leading to improved generalization in uncertainty estimation tasks.

## 6.6 Application of Optimal Transport in learning energy-based models

The foundational principle of Optimal Transport lies in the minimization of the cost required to transport mass from one probability distribution to another. Let $(X, d)$ be a metric space, where $X$ is the sample space and $d: X \times X \to \mathbb{R}^+$ is a metric. Given two probability measures $P_{\text{data}}$ and $P_{\theta}$ defined on $X$, the OT problem seeks to find the optimal coupling $\gamma$ that minimizes the total transportation cost:

$$ W_c(P_{\text{data}}, P_{\theta}) = \inf_{\gamma \in \Gamma(P_{\text{data}}, P_{\theta})} \int_{X \times X} c(x, y) d\gamma(x, y), \qquad (6.69) $$

where $\Gamma(P_{\text{data}}, P_{\theta})$ denotes the set of all joint probability measures $\gamma$ on $X \times X$ with marginals $P_{\text{data}}$ and $P_{\theta}$, and $c(x, y)$ is a cost function, typically chosen as $c(x, y) = d(x, y)^p$ for $p \ge 1$. For $p = 2$, this yields the squared Euclidean cost, and the corresponding OT distance is the 2-Wasserstein distance:

$$ W_2(P_{\text{data}}, P_{\theta}) = \left( \inf_{\gamma \in \Gamma(P_{\text{data}}, P_{\theta})} \int_{X \times X} \|x - y\|_2^2 d\gamma(x, y) \right)^{1/2} . \qquad (6.70) $$

In the context of energy-based models, the model distribution $P_{\theta}(x)$ is defined via an energy function $E_{\theta}(x)$ and the Boltzmann-Gibbs distribution:

$$ P_{\theta}(x) = \frac{1}{Z(\theta)} \exp(-E_{\theta}(x)), \qquad (6.71) $$

$$ Z(\theta) = \int_X \exp(-E_{\theta}(x)) dx, \qquad (6.72) $$
where $Z(\theta)$ is the partition function. The goal is to minimize the discrepancy between $P_{\text{data}}$ and $P_{\theta}$, which can be achieved by minimizing the Wasserstein distance $W_2(P_{\text{data}}, P_{\theta})$. The Kantorovich duality theorem provides a dual formulation of the OT problem:

$$W_2^2(P_{\text{data}}, P_{\theta}) = \sup_{\varphi, \psi} \left\{ \int_X \varphi(x) dP_{\text{data}}(x) + \int_X \psi(y) dP_{\theta}(y) \right\}, \qquad (6.73)$$

subject to the constraint $\varphi(x) + \psi(y) \le \|x-y\|_2$ for all $x, y \in X$. Here, $\varphi$ and $\psi$ are the Kantorovich potentials, which are related via the c-transform:

$$\varphi_c(y) = \inf_{x \in X} \{\|x - y\|_2 - \varphi(x)\}, \qquad (6.74)$$

$$\psi_c(x) = \inf_{y \in X} \{\|x - y\|_2 - \psi(y)\}. \qquad (6.75)$$

The Kantorovich potentials $\varphi$ and $\psi$ can be parameterized using neural networks, leading to the Wasserstein Generative Adversarial Network (WGAN) framework. In this setting, the discriminator learns $\varphi$, while the generator minimizes the Wasserstein distance. The gradient of the Wasserstein distance with respect to the model parameters $\theta$ is given by:

$$\nabla_{\theta} W_2^2(P_{\text{data}}, P_{\theta}) = \mathbb{E}_{x \sim P_{\text{data}}} [\nabla_{\theta} \varphi_{\theta}(x)] - \mathbb{E}_{y \sim P_{\theta}} [\nabla_{\theta} \varphi_{\theta}(y)], \qquad (6.76)$$

where $\varphi_{\theta}$ is the Kantorovich potential parameterized by $\theta$. This gradient can be estimated using Monte Carlo sampling from $P_{\text{data}}$ and $P_{\theta}$, enabling stochastic optimization. To enhance computational efficiency and stability, entropy regularization is often introduced, leading to the Sinkhorn divergence:

$$W_{2, \epsilon}^2(P_{\text{data}}, P_{\theta}) = \inf_{\gamma \in \Gamma(P_{\text{data}}, P_{\theta})} \left\{ \int_{X \times X} \|x - y\|_2^2 d\gamma(x, y) + \epsilon H(\gamma) \right\}, \qquad (6.77)$$

where

$$H(\gamma) = \int_{X \times X} \gamma(x, y) \log \gamma(x, y) dx dy \qquad (6.78)$$

is the entropy of the coupling $\gamma$, and $\epsilon > 0$ is the regularization parameter. The Sinkhorn divergence provides a smoothed approximation of the Wasserstein distance, which is particularly advantageous in high-dimensional settings. The optimization problem for learning EBMs using OT can be formulated as:

$$\min_{\theta} W_{2, \epsilon}^2 (P_{\text{data}}, P_{\theta}) + \lambda R(\theta), \qquad (6.79)$$

where $R(\theta)$ is a regularization term, such as $\|\theta\|_2^2$, and $\lambda$ is a hyperparameter. The gradient of the Sinkhorn divergence with respect to $\theta$ is:

$$\nabla_{\theta} W_{2, \epsilon}^2(P_{\text{data}}, P_{\theta}) = \mathbb{E}_{x \sim P_{\text{data}}} [\nabla_{\theta} \varphi_{\theta}(x)] - \mathbb{E}_{y \sim P_{\theta}} [\nabla_{\theta} \varphi_{\theta}(y)] + \epsilon \nabla_{\theta} H(\gamma_{\theta}), \qquad (6.80)$$

where $\gamma_{\theta}$ is the optimal coupling under entropy regularization. This gradient can be computed efficiently using the Sinkhorn-Knopp algorithm, which iteratively updates the coupling $\gamma$ and the potentials $\varphi$.

In summary, the application of Optimal Transport to learning energy-based models provides a mathematically rigorous framework for minimizing the discrepancy between data and model distributions. The interplay between Kantorovich duality, entropy regularization, and stochastic optimization enables efficient and scalable learning of complex, high-dimensional distributions. The resulting OT-based EBMs are both theoretically sound and practically effective, with connections to generative models such as WGANs and variational inference.
## 6.7 Sinkhorn-Knopp Algorithm

The Sinkhorn-Knopp algorithm is a numerical method for solving entropy-regularized optimal transport problems, which are central to many applications in deep neural networks, including generative modeling, domain adaptation, and representation learning. The optimal transport problem seeks to find a coupling matrix $P \in \mathbb{R}_+^{n \times m}$ that minimizes the transportation cost between two probability distributions $\mathbf{a} \in \mathbb{R}_+^n$ and $\mathbf{b} \in \mathbb{R}_+^m$, subject to marginal constraints. The cost is quantified by a cost matrix $C \in \mathbb{R}_+^{n \times m}$, where $C_{ij}$ represents the cost of transporting mass from point $i$ in the source distribution to point $j$ in the target distribution. The entropy-regularized optimal transport problem introduces a regularization term to the objective function, resulting in the following convex optimization problem:

$$ \min_{P \in \mathbb{R}_+^{n \times m}} \langle P, C \rangle - \epsilon H(P), \qquad (6.81) $$

subject to the constraints:

$$ P1_m = \mathbf{a}, \quad P^\top 1_n = \mathbf{b}, \qquad (6.82) $$

where $\langle P, C \rangle = \sum_{i=1}^n \sum_{j=1}^m P_{ij}C_{ij}$ is the Frobenius inner product, $\epsilon > 0$ is the regularization parameter, and $H(P) = -\sum_{i=1}^n \sum_{j=1}^m P_{ij}(\log P_{ij} - 1)$ is the entropy of the coupling matrix $P$. The regularization term $\epsilon H(P)$ ensures that the solution $P$ is smooth and computationally tractable, as it penalizes deviations from uniformity. The Sinkhorn-Knopp algorithm solves this problem by leveraging the dual formulation of the entropy-regularized optimal transport problem. The dual problem can be derived using the method of Lagrange multipliers, introducing dual variables $\mathbf{f} \in \mathbb{R}^n$ and $\mathbf{g} \in \mathbb{R}^m$ corresponding to the marginal constraints. The Lagrangian of the problem is given by:

$$ \mathcal{L}(P, \mathbf{f}, \mathbf{g}) = \langle P, C \rangle - \epsilon H(P) - \mathbf{f}^\top (P 1_m - \mathbf{a}) - \mathbf{g}^\top (P^\top 1_n - \mathbf{b}). \qquad (6.83) $$

Taking the gradient of the Lagrangian with respect to $P$ and setting it to zero yields the optimality condition:

$$ P_{ij} = \exp \left( \frac{f_i + g_j - C_{ij}}{\epsilon} \right). \qquad (6.84) $$

This expression can be rewritten in terms of the Gibbs kernel $K \in \mathbb{R}_+^{n \times m}$, defined as $K_{ij} = \exp(-C_{ij}/\epsilon)$, and the scaling vectors $\mathbf{u} = \exp(\mathbf{f}/\epsilon)$ and $\mathbf{v} = \exp(\mathbf{g}/\epsilon)$. The coupling matrix $P$ is then given by:

$$ P_{ij} = u_i K_{ij} v_j. \qquad (6.85) $$

The Sinkhorn-Knopp algorithm iteratively updates the scaling vectors $\mathbf{u}$ and $\mathbf{v}$ to satisfy the marginal constraints $P1_m = \mathbf{a}$ and $P^\top 1_n = \mathbf{b}$. The update rules are derived from the optimality conditions and are given by:

$$ u_i^{(k+1)} = \frac{a_i}{\sum_{j=1}^{m} K_{ij} v_j^{(k)}}, \quad v_j^{(k+1)} = \frac{b_j}{\sum_{i=1}^{n} K_{ij} u_i^{(k+1)}} \qquad (6.86) $$

These updates ensure that the coupling matrix $P$ converges to the unique solution of the entropy-regularized optimal transport problem. The convergence of the algorithm is guaranteed by the contractive nature of the updates, and the rate of convergence is linear in the number of iterations. In the context of deep neural networks, the Sinkhorn-Knopp algorithm is often used to compute approximate Wasserstein distances, which are essential for training generative models such as Wasserstein GANs. The Wasserstein distance between two probability distributions $\mathbf{a}$ and $\mathbf{b}$ is defined as:

$$ W(\mathbf{a}, \mathbf{b}) = \min_{P \in \Pi(\mathbf{a}, \mathbf{b})} \langle P, C \rangle, \qquad (6.87) $$

where $\Pi(\mathbf{a}, \mathbf{b})$ is the set of all coupling matrices satisfying the marginal constraints. The entropy-regularized version of the Wasserstein distance, known as the Sinkhorn distance, is given by:

$$ W_{\epsilon}(\mathbf{a}, \mathbf{b}) = \langle P^{*}, C \rangle - \epsilon H(P^{*}), \qquad (6.88) $$
where $P^*$ is the solution to the entropy-regularized optimal transport problem. The Sinkhorn-Knopp algorithm provides an efficient way to compute $P^*$ and, consequently, the Sinkhorn distance. The algorithm can also be extended to handle unbalanced optimal transport problems, where the source and target distributions do not necessarily have the same total mass. In this case, the optimization problem is modified to include additional penalty terms for deviations from the marginal constraints:

$$ \min_{P \in \mathbb{R}_+^{n \times m}} \langle P, C \rangle - \epsilon H(P) + \lambda_1 \mathrm{KL}(P \mathbf{1}_m \| \mathbf{a}) + \lambda_2 \mathrm{KL}(P^\top \mathbf{1}_n \| \mathbf{b}), \quad (6.89) $$

where $\mathrm{KL}(\cdot \| \cdot)$ is the Kullback-Leibler divergence, and $\lambda_1, \lambda_2 > 0$ are regularization parameters. The Sinkhorn-Knopp algorithm can be adapted to solve this problem by modifying the update rules for $\mathbf{u}$ and $\mathbf{v}$:

$$ u_i^{(k+1)} = \frac{a_i}{\sum_{j=1}^{m} K_{ij} v_j^{(k)}} + \lambda_1, \quad v_j^{(k+1)} = \frac{b_j}{\sum_{i=1}^{n} K_{ij} u_i^{(k+1)}} + \lambda_2 \quad (6.90) $$

This extension allows the algorithm to handle a broader range of applications in deep learning, such as semi-supervised learning and data augmentation, where the source and target distributions may not be perfectly aligned.

In summary, the Sinkhorn-Knopp algorithm is a highly efficient and mathematically rigorous method for solving entropy-regularized optimal transport problems. Its iterative nature, linear convergence rate, and ability to handle both balanced and unbalanced transport problems make it a versatile tool in deep learning. The algorithm's reliance on matrix scaling operations ensures that it can be seamlessly integrated into existing deep learning frameworks, making it a cornerstone of modern optimal transport theory and its applications.

### 6.7.1 Sinkhorn Distance

The Sinkhorn distance is a mathematically rigorous and computationally efficient approximation of the optimal transport problem, which seeks to minimize the cost of transporting mass from one probability distribution to another. Let $\mathbf{a} \in \mathbb{R}^n$ and $\mathbf{b} \in \mathbb{R}^m$ be two discrete probability distributions defined over finite spaces, and let $\mathbf{C} \in \mathbb{R}^{n \times m}$ be a cost matrix where $C_{ij}$ represents the cost of transporting a unit of mass from point $i$ in the support of $\mathbf{a}$ to point $j$ in the support of $\mathbf{b}$. The optimal transport problem is formulated as:

$$ \mathrm{OT}(\mathbf{a}, \mathbf{b}) = \min_{\mathbf{P} \in \mathcal{U}(\mathbf{a}, \mathbf{b})} \langle \mathbf{P}, \mathbf{C} \rangle, \quad (6.91) $$

where $\mathcal{U}(\mathbf{a}, \mathbf{b}) = \{\mathbf{P} \in \mathbb{R}_+^{n \times m} | \mathbf{P} \mathbf{1}_m = \mathbf{a}, \mathbf{P}^\top \mathbf{1}_n = \mathbf{b}\}$ is the set of all valid transport plans, $\mathbf{1}_n$ and $\mathbf{1}_m$ are vectors of ones, and $\langle \mathbf{P}, \mathbf{C} \rangle = \sum_{i=1}^n \sum_{j=1}^m P_{ij} C_{ij}$ is the Frobenius inner product. The Sinkhorn distance introduces an entropic regularization term to this problem, resulting in the following optimization problem:

$$ \mathrm{OT}_{\lambda}(\mathbf{a}, \mathbf{b}) = \min_{\mathbf{P} \in \mathcal{U}(\mathbf{a}, \mathbf{b})} \langle \mathbf{P}, \mathbf{C} \rangle - \lambda H(\mathbf{P}), \quad (6.92) $$

where $\lambda > 0$ is the regularization parameter, and $H(\mathbf{P}) = - \sum_{i=1}^n \sum_{j=1}^m P_{ij}(\log P_{ij} - 1)$ is the entropy of the transport plan $\mathbf{P}$. The entropic regularization ensures that the optimal transport plan $\mathbf{P}^*$ is unique and has the form:

$$ \mathbf{P}^* = \mathrm{diag}(\mathbf{u})\mathrm{diag}(\mathbf{v}), \quad (6.93) $$

where $\mathbf{K} = \exp(-\mathbf{C}/\lambda)$ is the Gibbs kernel, and $\mathbf{u} \in \mathbb{R}^n$ and $\mathbf{v} \in \mathbb{R}^m$ are positive scaling vectors that satisfy the marginal constraints:

$$ \mathbf{u} \odot (\mathbf{K}\mathbf{v}) = \mathbf{a} \quad \text{and} \quad \mathbf{v} \odot (\mathbf{K}^\top\mathbf{u}) = \mathbf{b}. \quad (6.94) $$
Here, $\odot$ denotes element-wise multiplication. The Sinkhorn-Knopp algorithm is used to iteratively compute these scaling vectors:

$$ \mathbf{u}^{(k+1)} = \frac{\mathbf{a}}{K\mathbf{v}^{(k)}} \quad \text{and} \quad \mathbf{v}^{(k+1)} = \frac{\mathbf{b}}{K^{\top}\mathbf{u}^{(k+1)}}, \qquad (6.95) $$

where $\mathbf{u}^{(0)} = \mathbf{1}_n$ and $\mathbf{v}^{(0)} = \mathbf{1}_m$ are initialized as vectors of ones. The algorithm converges to the optimal scaling vectors $\mathbf{u}^*$ and $\mathbf{v}^*$, which define the optimal transport plan $\mathbf{P}^*$. The Sinkhorn distance is then given by:

$$ \text{Sinkhorn}_{\lambda}(\mathbf{a}, \mathbf{b}) = \langle \mathbf{P}^*, \mathbf{C} \rangle - \lambda H(\mathbf{P}^*). \qquad (6.96) $$

The regularization parameter $\lambda$ controls the trade-off between the transportation cost and the entropy of the transport plan. As $\lambda \to \infty$, the Sinkhorn distance approaches the original optimal transport distance, while as $\lambda \to 0$, the transport plan becomes more diffuse and the distance approaches the entropy-regularized Wasserstein distance. The Sinkhorn distance is computationally efficient due to the iterative nature of the Sinkhorn-Knopp algorithm, making it suitable for large-scale applications. The dual formulation of the regularized optimal transport problem provides additional insights into the structure of the Sinkhorn distance. The dual problem is given by:

$$ \text{OT}_{\lambda}(\mathbf{a}, \mathbf{b}) = \max_{\mathbf{f}, \mathbf{g}} \mathbf{f}^{\top}\mathbf{a} + \mathbf{g}^{\top}\mathbf{b} - \lambda \sum_{i=1}^{n} \sum_{j=1}^{m} \exp \left( \frac{f_i + g_j - C_{ij}}{\lambda} \right), \qquad (6.97) $$

where $\mathbf{f} \in \mathbb{R}^n$ and $\mathbf{g} \in \mathbb{R}^m$ are dual variables. The optimal dual variables $\mathbf{f}^*$ and $\mathbf{g}^*$ are related to the scaling vectors $\mathbf{u}^*$ and $\mathbf{v}^*$ by:

$$ \mathbf{u}^* = \exp\left(\frac{\mathbf{f}^*}{\lambda}\right) \quad \text{and} \quad \mathbf{v}^* = \exp\left(\frac{\mathbf{g}^*}{\lambda}\right). \qquad (6.98) $$

The Sinkhorn distance can thus be computed using either the primal or dual formulation, depending on the specific application. The dual formulation is particularly useful for deriving theoretical properties of the Sinkhorn distance, such as its convexity and smoothness with respect to the input distributions $\mathbf{a}$ and $\mathbf{b}$. The Sinkhorn distance satisfies several key mathematical properties. First, it is symmetric, meaning that $\text{Sinkhorn}_{\lambda}(\mathbf{a}, \mathbf{b}) = \text{Sinkhorn}_{\lambda}(\mathbf{b}, \mathbf{a})$. Second, it satisfies the triangle inequality, making it a valid metric on the space of probability distributions. Third, it is differentiable with respect to the input distributions $\mathbf{a}$ and $\mathbf{b}$, which is useful for gradient-based optimization in machine learning applications. Finally, the Sinkhorn distance is computationally efficient, with a time complexity of $O(nm)$ per iteration of the Sinkhorn-Knopp algorithm, making it scalable to high-dimensional problems. The Sinkhorn distance can also be generalized to continuous probability measures using the Kantorovich formulation of optimal transport. Let $\mu$ and $\nu$ be probability measures defined on metric spaces $\mathcal{X}$ and $\mathcal{Y}$, respectively, and let $c: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}_+$ be a continuous cost function. The entropically regularized optimal transport problem in the continuous setting is given by:

$$ \text{OT}_{\lambda}(\mu, \nu) = \min_{\pi \in \Pi(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) d\pi(x, y) - \lambda H(\pi), \qquad (6.99) $$

where $\Pi(\mu, \nu)$ is the set of all joint probability measures $\pi$ with marginals $\mu$ and $\nu$, and $H(\pi) = - \int_{\mathcal{X} \times \mathcal{Y}} \log \left( \frac{d\pi}{d\mu \otimes d\nu} \right) d\pi$ is the entropy of the coupling $\pi$. The optimal coupling $\pi^*$ has the form:

$$ d\pi^*(x, y) = \exp\left(\frac{f(x) + g(y) - c(x, y)}{\lambda}\right) d\mu(x) d\nu(y), \qquad (6.100) $$

where $f$ and $g$ are continuous functions satisfying the Schrödinger system:

$$ \int_{\mathcal{Y}} \exp \left( \frac{f(x) + g(y) - c(x, y)}{\lambda} \right) d\nu(y) = 1 \quad \text{and} \quad \int_{\mathcal{X}} \exp \left( \frac{f(x) + g(y) - c(x, y)}{\lambda} \right) d\mu(x) = 1. \qquad (6.101) $$
The Sinkhorn distance in the continuous setting is then given by:

$$ \text{Sinkhorn}_{\lambda}(\mu, \nu) = \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) d\pi^*(x, y) - \lambda H(\pi^*). \qquad (6.102) $$

This formulation extends the Sinkhorn distance to continuous probability measures, enabling its application to a broader class of problems in probability theory, statistics, and machine learning. The Sinkhorn distance thus provides a rigorous and computationally efficient framework for measuring the dissimilarity between probability distributions, combining the theoretical foundations of optimal transport with the practical advantages of entropic regularization.

## 6.7.2 Wasserstein GANs

Wasserstein Generative Adversarial Networks (WGANs) are a class of generative models that utilize the Wasserstein-1 distance, also known as the Earth Mover's Distance (EMD), to measure the discrepancy between the real data distribution $P_r$ and the generated data distribution $P_g$. The Wasserstein-1 distance is defined as:

$$ W_1(\mathbb{P}_r, \mathbb{P}_g) = \inf_{\gamma \in \Pi(\mathbb{P}_r, \mathbb{P}_g)} \mathbb{E}_{(x,y) \sim \gamma} [\|x - y\|] \qquad (6.103) $$

Here, $\Pi(\mathbb{P}_r, \mathbb{P}_g)$ represents the set of all joint distributions $\gamma(x, y)$ with marginals $\mathbb{P}_r$ and $\mathbb{P}_g$, and the infimum is taken over all such joint distributions. The Wasserstein-1 distance quantifies the minimal cost required to transport mass from $\mathbb{P}_r$ to $\mathbb{P}_g$, where the cost is proportional to the Euclidean distance $\|x - y\|$. In WGANs, the generator $G$ and the critic $D$ are optimized to minimize and maximize the Wasserstein-1 distance, respectively. The critic $D$ is trained to approximate the Wasserstein-1 distance between $\mathbb{P}_r$ and $\mathbb{P}_g$. The objective function for the critic is given by:

$$ L_D = \mathbb{E}_{x \sim \mathbb{P}_r} [D(x)] - \mathbb{E}_{z \sim p_z(z)} [D(G(z))] \qquad (6.104) $$

Here, $z$ is a random noise vector sampled from a prior distribution $p_z(z)$, and $G(z)$ is the generated sample. To ensure that the critic $D$ is a 1-Lipschitz function, a constraint is imposed on its gradients. This constraint is typically enforced using a gradient penalty term, leading to the modified critic loss function:

$$ L_D^{GP} = \mathbb{E}_{x \sim \mathbb{P}_r} [D(x)] - \mathbb{E}_{z \sim p_z(z)} [D(G(z))] + \lambda \mathbb{E}_{\hat{x} \sim \mathbb{P}_{\hat{x}}} [(\|\nabla_{\hat{x}} D(\hat{x})\|_2 - 1)^2] \qquad (6.105) $$

In this equation, $\lambda$ is a hyperparameter controlling the strength of the gradient penalty, and $\hat{x}$ is sampled uniformly along straight lines connecting pairs of points sampled from $\mathbb{P}_r$ and $\mathbb{P}_g$. The term $(\|\nabla_{\hat{x}} D(\hat{x})\|_2 - 1)^2$ penalizes deviations from the 1-Lipschitz constraint, ensuring that the critic's gradients have a norm of at most 1. The generator $G$ is trained to minimize the Wasserstein-1 distance by maximizing the critic's output on generated samples. The objective function for the generator is given by:

$$ L_G = -\mathbb{E}_{z \sim p_z(z)} [D(G(z))] \qquad (6.106) $$

The optimization process alternates between updating the critic and the generator. The critic is updated multiple times for each update of the generator to ensure that it remains close to the optimal 1-Lipschitz function. The parameter updates for the critic and generator are given by:

$$ \theta_D \leftarrow \theta_D - \eta_D \nabla_{\theta_D} L_D^{GP} \qquad (6.107) $$

$$ \theta_G \leftarrow \theta_G - \eta_G \nabla_{\theta_G} L_G \qquad (6.108) $$

Here, $\theta_D$ and $\theta_G$ are the parameters of the critic and generator, respectively, and $\eta_D$ and $\eta_G$ are the learning rates for the critic and generator. The training continues until the generator produces samples that are indistinguishable from real data according to the critic. The Wasserstein-1 distance provides several theoretical advantages over traditional GAN objectives, such as the
Jensen-Shannon divergence used in the original GAN formulation. The Wasserstein-1 distance is continuous and differentiable almost everywhere, which leads to more stable training dynamics. Additionally, the Wasserstein-1 distance correlates better with sample quality, providing a meaningful metric for evaluating the performance of the generator. The mathematical properties of the Wasserstein-1 distance, including its sensitivity to the geometry of the underlying space and its ability to provide meaningful gradients even when the supports of $P_r$ and $P_g$ are disjoint, make WGANs a powerful tool for generative modeling.

In summary, WGANs optimize the following minimax objective:

$$ \min_{G} \max_{D \in \mathcal{D}} \mathbb{E}_{x \sim P_r}[D(x)] - \mathbb{E}_{z \sim p_z(z)}[D(G(z))] \quad (6.109) $$

where $\mathcal{D}$ is the set of 1-Lipschitz functions. The Wasserstein-1 distance $W_1(P_r, P_g)$ is approximated by the critic $D$, and the generator $G$ is trained to minimize this distance. The use of the Wasserstein-1 distance in WGANs provides a theoretically grounded framework for generative modeling, leading to improved stability and performance compared to traditional GANs. The mathematical rigor of the Wasserstein-1 distance and its properties ensure that WGANs are well-suited for a wide range of generative tasks.

## 6.8 Kantorovich Duality

The Kantorovich duality theorem is a fundamental result in optimal transport theory that establishes the equivalence between the primal formulation of the optimal transport problem and a dual formulation based on potential functions. Given two probability spaces $(\mathcal{X}, \mu)$ and $(\mathcal{Y}, \nu)$, with a cost function $c: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}$, the primal Kantorovich formulation of optimal transport seeks to minimize the total transport cost over all possible couplings $\gamma$ of $\mu$ and $\nu$, where a coupling is a joint probability measure on $\mathcal{X} \times \mathcal{Y}$ whose marginals are $\mu$ and $\nu$, respectively. Mathematically, this is expressed as

$$ \inf_{\gamma \in \Pi(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) d\gamma(x, y), \quad (6.110) $$

where $\Pi(\mu, \nu)$ is the set of all couplings satisfying

$$ \int_{\mathcal{X}} d\gamma(x, y) = d\nu(y), \quad \int_{\mathcal{Y}} d\gamma(x, y) = d\mu(x). \quad (6.111) $$

The Kantorovich duality states that the optimal transport cost can equivalently be expressed in terms of potential functions $\varphi: \mathcal{X} \to \mathbb{R}$ and $\psi: \mathcal{Y} \to \mathbb{R}$ satisfying the inequality

$$ \varphi(x) + \psi(y) \le c(x, y), \quad \forall (x, y) \in \mathcal{X} \times \mathcal{Y}. \quad (6.112) $$

Under appropriate conditions, the optimal transport cost is given by the dual formulation

$$ \sup_{\varphi, \psi} \left\{ \int_{\mathcal{X}} \varphi(x) d\mu(x) + \int_{\mathcal{Y}} \psi(y) d\nu(y) \right\} \quad (6.113) $$

subject to the constraint

$$ \varphi(x) + \psi(y) \le c(x, y), \quad \forall (x, y). \quad (6.114) $$

A key result in this formulation is that when the cost function $c(x, y)$ is lower semi-continuous, the supremum is attained by functions $\varphi$ and $\psi$ that satisfy

$$ \psi(y) = \inf_{x \in \mathcal{X}} [c(x, y) - \varphi(x)]. \quad (6.115) $$
For the specific case of the Wasserstein-1 distance, where $c(x, y) = \|x - y\|$, the dual problem simplifies to

$$W_1(\mu, \nu) = \sup_{\|\varphi\|_L \le 1} \left\{ \int_X \varphi(x) d\mu(x) - \int_Y \varphi(y) d\nu(y) \right\}. \qquad (6.116)$$

Here, the supremum is taken over all 1-Lipschitz functions $\varphi : \mathcal{X} \to \mathbb{R}$, satisfying

$$|\varphi(x) - \varphi(y)| \le \|x - y\|, \quad \forall x, y \in \mathcal{X}. \qquad (6.117)$$

This dual formulation is of fundamental importance in deep learning applications, particularly in the training of Wasserstein Generative Adversarial Networks (WGANs). In this setting, the generator $G: \mathcal{Z} \to \mathcal{X}$ transforms samples $z \sim p_z$ from a latent space $\mathcal{Z}$ into synthetic data points in $\mathcal{X}$, aiming to approximate a target data distribution $\mu$. The discriminator network approximates the Kantorovich potential $\varphi$, leading to the objective

$$\min_G \max_{\|\varphi\|_L \le 1} \left\{ \mathbb{E}_{x \sim \mu}[\varphi(x)] - \mathbb{E}_{y \sim \nu}[\varphi(G(z))] \right\}. \qquad (6.118)$$

This adversarial framework provides a more stable training procedure than standard GANs, which rely on the Jensen-Shannon divergence. The Kantorovich duality ensures that the discriminator function learns a transport potential that characterizes the optimal transport map, guiding the generator toward minimizing the Wasserstein distance between the real and generated distributions. Beyond generative modeling, optimal transport theory is also applied in supervised learning, where the Wasserstein distance serves as a loss function for comparing distributions in high-dimensional spaces. When computational efficiency is required, entropic regularization techniques introduce an additional entropy penalty term $H(\gamma)$, yielding the regularized transport problem

$$\inf_{\gamma \in \Pi(\mu, \nu)} \left\{ \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) d\gamma(x, y) + \epsilon H(\gamma) \right\}, \qquad (6.119)$$

where

$$H(\gamma) = \int_{\mathcal{X} \times \mathcal{Y}} \gamma(x, y) \log \gamma(x, y) dx dy. \qquad (6.120)$$

This modification allows for efficient computation via Sinkhorn iterations, leveraging the dual formulation

$$\sup_{\varphi, \psi} \left\{ \int_{\mathcal{X}} \varphi(x) d\mu(x) + \int_{\mathcal{Y}} \psi(y) d\nu(y) - \epsilon \int_{\mathcal{X} \times \mathcal{Y}} e^{(\varphi(x) + \psi(y) - c(x, y))/\epsilon} dx dy \right\}. \qquad (6.121)$$

Such approaches are particularly useful in deep learning, where high-dimensional transport problems arise in applications such as domain adaptation, clustering, and metric learning. The theoretical underpinnings of Kantorovich duality thus play a critical role in designing robust algorithms for optimizing probability distributions in high-dimensional spaces, ensuring efficient convergence and improved generalization performance in deep neural networks.

## 6.9 Entropy Regularization

The entropy regularization of optimal transport arises from the classical Monge-Kantorovich formulation, where the goal is to find a transport plan $\gamma \in \Pi(\mu, \nu)$ that minimizes the expected transportation cost between two probability measures $\mu$ and $\nu$. Mathematically, this can be expressed as the minimization problem

$$W_c(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) d\gamma(x, y), \qquad (6.122)$$
where $\Pi(\mu, \nu)$ is the space of joint probability distributions on $\mathcal{X} \times \mathcal{Y}$ with marginals $\mu$ and $\nu$, respectively, satisfying the marginalization conditions

$$ \int_{\mathcal{Y}} d\gamma(x, y) = d\mu(x), \quad \int_{\mathcal{X}} d\gamma(x, y) = d\nu(y). \qquad (6.123) $$

While this optimization problem is well-posed, it is computationally intractable in high-dimensional settings due to the combinatorial explosion of feasible transport plans. To mitigate this, an entropy regularization term is introduced, leading to the entropy-regularized optimal transport problem

$$ W_c^{\lambda}(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \left( \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) d\gamma(x, y) + \lambda H(\gamma) \right), \qquad (6.124) $$

where $H(\gamma)$ is the Shannon entropy of the transport plan $\gamma$, given by

$$ H(\gamma) = - \int_{\mathcal{X} \times \mathcal{Y}} \gamma(x, y) \log \gamma(x, y) \, dx \, dy. \qquad (6.125) $$

The introduction of entropy regularization has fundamental mathematical consequences. First, it ensures the strict convexity of the optimization problem, which allows for efficient computation using iterative methods. Second, it guarantees the uniqueness of the optimal transport plan, preventing the degeneracy issues that arise in the classical unregularized setting. The optimal transport plan in the entropy-regularized setting takes the form

$$ \gamma^{\lambda}(x, y) = \exp\left(-\frac{c(x, y)}{\lambda}\right) \mu(x)\nu(y), \qquad (6.126) $$

which is known as the Gibbs kernel representation. This form follows from the variational principle associated with entropy maximization under marginal constraints, leading to the closed-form solution

$$ \gamma^{\lambda}(x, y) = \mu(x)\nu(y)e^{\frac{\phi(x)+\psi(y)-c(x,y)}{\lambda}}, \qquad (6.127) $$

where $\phi(x)$ and $\psi(y)$ are dual potentials arising from the Lagrange multipliers enforcing the marginal constraints. The dual formulation of entropy-regularized optimal transport is derived from the Fenchel-Rockafellar duality theorem and takes the form

$$ \sup_{\phi, \psi} \left( \int_{\mathcal{X}} \phi(x) d\mu(x) + \int_{\mathcal{Y}} \psi(y) d\nu(y) - \lambda \int_{\mathcal{X} \times \mathcal{Y}} e^{\frac{\phi(x)+\psi(y)-c(x,y)}{\lambda}} dxdy \right). \qquad (6.128) $$

The dual potentials $\phi(x)$ and $\psi(y)$ satisfy the iterative update equations

$$ \phi^{(k+1)}(x) = -\lambda \log \int_{\mathcal{Y}} e^{\frac{\psi^{(k)}(y)-c(x,y)}{\lambda}} d\nu(y), \qquad (6.129) $$

$$ \psi^{(k+1)}(y) = -\lambda \log \int_{\mathcal{X}} e^{\frac{\phi^{(k)}(x)-c(x,y)}{\lambda}} d\mu(x). \qquad (6.130) $$

These updates converge exponentially to the optimal dual potentials, enabling the efficient computation of entropy-regularized transport distances. In the discrete setting, where the probability measures $\mu$ and $\nu$ are represented as vectors $u \in \mathbb{R}^n$ and $v \in \mathbb{R}^m$, the entropy-regularized optimal transport problem reduces to a matrix scaling problem. Given a cost matrix $C \in \mathbb{R}^{n \times m}$, the optimal transport matrix $T$ satisfies

$$ T_{ij} = \exp\left(-\frac{C_{ij}}{\lambda}\right) u_i v_j. \qquad (6.131) $$

This leads to the Sinkhorn-Knopp iteration

$$ u^{(k+1)} = \frac{\mu}{K v^{(k)}}, \quad v^{(k+1)} = \frac{\nu}{K^T u^{(k)}}, \qquad (6.132) $$
where $K_{ij} = e^{-C_{ij}/\lambda}$ and division is element-wise. These updates ensure fast convergence and numerical stability, making entropy-regularized optimal transport suitable for large-scale applications in deep learning. In the context of deep neural networks, entropy-regularized optimal transport plays a crucial role in generative modeling, domain adaptation, and adversarial training. Given a parameterized generator $G_{\theta}(z)$ with latent variable $z \sim P_Z$, the entropy-regularized Wasserstein loss for training a generative model is given by

$$ \mathcal{L}(\theta) = W_c^{\lambda}(P_{G_{\theta}}, P_X) = \inf_{\gamma \in \Pi(P_{G_{\theta}}, P_X)} \left( \int c(G_{\theta}(z), x) d\gamma(z, x) + \lambda H(\gamma) \right). \quad (6.133) $$

This formulation enables stable gradient-based training by ensuring smoothness in the transport plan. Furthermore, in domain adaptation, the entropy-regularized Wasserstein distance is minimized between the source and target feature distributions $P_S$ and $P_T$, leading to the optimization problem

$$ \mathcal{L}_{\text{OT}} = W_c^{\lambda}(P_S, P_T) = \inf_{\gamma \in \Pi(P_S, P_T)} \left( \int c(x, y) d\gamma(x, y) + \lambda H(\gamma) \right). \quad (6.134) $$

This objective ensures that the learned representation preserves geometric consistency between domains while mitigating domain shift. Thus, entropy regularization in optimal transport provides a mathematically rigorous, computationally efficient, and theoretically well-grounded framework for deep learning applications, ensuring convexity, uniqueness, and numerical stability while enabling large-scale optimization in high-dimensional probability spaces.

### 6.9.1 Classical Monge-Kantorovich formulation

The **Monge-Kantorovich optimal transport problem** seeks to determine the most efficient way to transport mass from one distribution to another while minimizing a given cost function. Let $X$ and $Y$ be two measurable spaces, each endowed with probability measures $\mu$ and $\nu$, respectively, which are absolutely continuous with respect to the Lebesgue measure, having corresponding density functions $f(x)$ and $g(y)$. The fundamental problem first considered by Monge is to find a measurable transport map $T: X \to Y$ satisfying the **pushforward condition**,

$$ T_{\#}\mu = \nu, \quad \text{or equivalently,} \quad \int_{A} g(y)dy = \int_{T^{-1}(A)} f(x)dx, \quad \forall A \subset Y, \quad (6.135) $$

which ensures that mass is conserved. Given a cost function $c: X \times Y \to \mathbb{R}$, the total transport cost associated with the map $T$ is given by the functional

$$ C[T] = \int_X c(x, T(x))f(x)dx. \quad (6.136) $$

The **Monge problem** is then formulated as the minimization problem

$$ \inf_{T} C[T] = \inf_{T} \int_{X} c(x, T(x))f(x)dx, \quad (6.137) $$

where the infimum is taken over all measurable transport maps satisfying the pushforward condition. The primary difficulty in solving Monge's problem lies in the fact that a transport map may not exist in general, particularly when the measure $\mu$ is more **spread out** than $\nu$, or when mass-splitting is required, which motivates the **relaxation** of Monge's problem to the Kantorovich formulation. Instead of restricting attention to deterministic transport maps, Kantorovich introduced the notion of **transport plans**, which are probability measures $\gamma$ on the product space $X \times Y$ satisfying the **marginal constraints**

$$ \int_{Y} d\gamma(x, y) = d\mu(x), \quad \int_{X} d\gamma(x, y) = d\nu(y). \quad (6.138) $$
The set of all such **couplings** is denoted as $\Pi(\mu, \nu)$, which consists of all probability measures $\gamma(x, y)$ whose projections onto $X$ and $Y$ coincide with $\mu$ and $\nu$, respectively. The **Kantorovich optimal transport problem** is then formulated as the minimization problem

$$ \inf_{\gamma \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\gamma(x, y). \qquad (6.139) $$

Under mild conditions, such as compactness of $X$ and $Y$ and continuity of $c(x, y)$, existence of an **optimal transport plan** $\gamma^*$ is guaranteed. If $\gamma^*$ is supported on the graph of a function $y = T(x)$, then the optimal plan corresponds to a Monge-type transport. The problem of characterizing such solutions leads naturally to **Kantorovich duality**, which provides a variational formulation of optimal transport in terms of potentials $\varphi: X \to \mathbb{R}$ and $\psi: Y \to \mathbb{R}$. The dual problem is expressed as

$$ \sup_{\varphi, \psi} \left\{ \int_X \varphi(x) d\mu(x) + \int_Y \psi(y) d\nu(y) \right\} \qquad (6.140) $$

subject to the inequality constraint

$$ \varphi(x) + \psi(y) \leq c(x, y), \quad \forall (x, y) \in X \times Y. \qquad (6.141) $$

The optimality conditions for the Kantorovich dual problem imply the **complementary slackness condition**

$$ \varphi(x) + \psi(y) = c(x, y) \quad \text{for} \quad (x, y) \in \text{supp}(\gamma^*), \qquad (6.142) $$

which characterizes the support of the optimal transport plan. When the cost function is the squared Euclidean distance, $c(x, y) = \|x - y\|^2$, the transport problem reduces to finding a convex function $u(x)$ whose gradient defines the optimal transport map. This leads to the **Monge-Ampère equation**

$$ \det D^2 u(x) = \frac{f(x)}{g(\nabla u(x))} \qquad (6.143) $$

which provides a **geometric interpretation** of optimal transport: the transport map is given by the gradient of a convex function, ensuring volume preservation under mass transport. In this setting, the transport map satisfies

$$ T(x) = \nabla u(x), \qquad (6.144) $$

which is the Brenier map for the quadratic cost function. The induced transport cost is given by the **Wasserstein-2 distance**

$$ W_2(\mu, \nu) = \left( \inf_{\gamma \in \Pi(\mu, \nu)} \int_{X \times Y} \|x - y\|^2 d\gamma(x, y) \right)^{1/2}. \qquad (6.145) $$

A fundamental computational approach to solving optimal transport problems is **entropic regularization**, which introduces an entropy term

$$ \epsilon H(\gamma) = \epsilon \int_{X \times Y} \gamma(x, y) \log \gamma(x, y) dx dy, \qquad (6.146) $$

leading to the **regularized optimal transport problem**

$$ \inf_{\gamma \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\gamma(x, y) + \epsilon H(\gamma). \qquad (6.147) $$

These equations reveal deep connections between optimal transport, geometric analysis, and functional inequalities, providing a unifying framework for problems in probability, physics, and machine learning.
## 6.9.2 Fenchel-Rockafellar duality theorem

The **Fenchel-Rockafellar duality theorem** is a cornerstone result in convex analysis, establishing a deep connection between a primal convex optimization problem and its associated dual problem via the framework of convex conjugates and subdifferential calculus. It generalizes classical duality principles by employing functional analytic techniques and provides a powerful formulation for the optimality conditions in convex minimization problems. The theorem holds in the setting of a Banach space $X$ paired with its continuous dual space $X^*$, and its implications extend to a broad range of optimization and variational problems.

Consider two proper, convex, and lower semicontinuous functions $f: X \to \mathbb{R} \cup \{+\infty\}$ and $g: X \to \mathbb{R} \cup \{+\infty\}$. The **primal optimization problem** is given by

$$ \inf_{x \in X} \{f(x) + g(x)\}, \qquad (6.148) $$

which seeks to minimize the sum of these two convex functions over the domain $X$. The associated **Fenchel conjugate function** (also referred to as the **Legendre-Fenchel transform**) is defined as

$$ f^*(x^*) = \sup_{x \in X} \{\langle x^*, x \rangle - f(x)\}, \quad x^* \in X^*, \qquad (6.149) $$

which captures the maximum possible affine lower bound of $f(x)$. Similarly, the conjugate of $g$ is given by

$$ g^*(x^*) = \sup_{x \in X} \{\langle x^*, x \rangle - g(x)\}, \quad x^* \in X^*. \qquad (6.150) $$

The **Fenchel dual problem** is then formulated as

$$ \sup_{x^* \in X^*} \{-(f^* + g^*)(x^*)\}. \qquad (6.151) $$

The Fenchel-Rockafellar duality theorem states that under appropriate constraint qualifications, strong duality holds, meaning that the optimal values of the primal and dual problems coincide:

$$ \inf_{x \in X} \{f(x) + g(x)\} = \sup_{x^* \in X^*} \{-(f^* + g^*)(x^*)\}. \qquad (6.152) $$

To ensure strong duality, a sufficient condition is the existence of a point $x_0 \in X$ such that

$$ x_0 \in \text{dom}(f) \cap \text{dom}(g), \quad \text{and} \quad 0 \in \text{core}(\text{dom}(g) - x_0), \qquad (6.153) $$

where the **effective domain** of a function is defined as

$$ \text{dom}(f) = \{x \in X \mid f(x) < +\infty\}, \qquad (6.154) $$

and the **core** (or algebraic interior) of a set $S$ is given by

$$ \text{core}(S) = \{x \in S \mid \forall y \in X, \exists t > 0 \text{ such that } x + t(y - x) \in S\}. \qquad (6.155) $$

This constraint qualification ensures that the subdifferential intersection necessary for optimality is nonempty. In particular, the subdifferential of a convex function $f$ at a point $x$ is defined as

$$ \partial f(x) = \{x^* \in X^* \mid f(y) \ge f(x) + \langle x^*, y - x \rangle, \forall y \in X\}. \qquad (6.156) $$

For optimality, we require that the subdifferentials of $f$ and $g$ satisfy

$$ 0 \in \partial f(x) + \partial g(x), \qquad (6.157) $$
which guarantees the existence of a Lagrange multiplier $x^*$ such that

$$x^* \in \partial f(x), \quad -x^* \in \partial g(x). \tag{6.158}$$

That is, the dual optimizer $x^*$ must belong to the intersection

$$\partial f(x) \cap (-\partial g(x)) \neq \emptyset. \tag{6.159}$$

This condition provides the necessary and sufficient optimality criterion in terms of subdifferential calculus. The theorem extends naturally to infinite-dimensional settings, where it serves as a foundation for convex variational problems in Hilbert and Banach spaces. By employing the Fenchel conjugates, it allows for a reformulation of many classical variational problems in dual terms, enabling powerful analytical techniques such as saddle-point theory, strong convexity conditions, and dual space representations. The theorem's implications in optimization, functional analysis, and mathematical economics highlight its fundamental role in modern analysis.

## 6.10 Optimal transport theory Based Training Dynamics and Sinkhorn Divergences

Optimal transport theory provides a mathematical framework for comparing probability distributions by minimizing the cost of transporting mass from one distribution to another. The theory is rooted in the Monge-Kantorovich problem, which seeks a transport plan that minimizes the total cost of moving mass between two probability measures. Given two probability measures $\mu$ and $\nu$ on spaces $X$ and $Y$, respectively, and a cost function $c: X \times Y \to \mathbb{R}_+$, the optimal transport problem is formulated as:

$$\inf_{\pi \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\pi(x, y),$$

where $\Pi(\mu, \nu)$ denotes the set of all joint probability measures on $X \times Y$ with marginals $\mu$ and $\nu$. The solution to this problem defines the Wasserstein distance, which quantifies the minimal cost of transforming one distribution into another.

Training dynamics in machine learning can be analyzed through the lens of optimal transport, particularly when studying the evolution of model parameters or the distribution of data. Optimal transport-based training dynamics focus on minimizing the Wasserstein distance between the model's output distribution and the target distribution. This approach is especially useful in generative modeling, where the goal is to align the generated samples with the true data distribution. The gradient flow of the Wasserstein distance induces a dynamics that pushes the model parameters toward configurations that minimize the transport cost. For a parameterized model $P_\theta$, the optimization problem becomes:

$$\min_{\theta} W_c(P_{\theta}, Q),$$

where $Q$ is the target distribution and $W_c$ is the Wasserstein distance with cost function $c$. The gradient descent update rule for $\theta$ involves computing the gradient of the Wasserstein distance, which can be approximated using the dual formulation of the optimal transport problem.

Sinkhorn divergences provide a computationally efficient alternative to the Wasserstein distance by introducing entropy regularization into the optimal transport problem. The entropy-regularized optimal transport problem is given by:

$$\inf_{\pi \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\pi(x, y) + \epsilon H(\pi),$$
where $H(\pi)$ is the entropy of the transport plan $\pi$ and $\epsilon > 0$ is the regularization parameter. The Sinkhorn divergence is then defined as the difference between the regularized optimal transport cost and its self-transport terms:

$$S_{\epsilon}(\mu, \nu) = W_{c,\epsilon}(\mu, \nu) - \frac{1}{2}W_{c,\epsilon}(\mu, \mu) - \frac{1}{2}W_{c,\epsilon}(\nu, \nu),$$

where $W_{c,\epsilon}$ denotes the entropy-regularized Wasserstein distance. The Sinkhorn divergence retains many desirable properties of the Wasserstein distance, such as metricity for certain cost functions, while being amenable to efficient computation using iterative scaling algorithms like the Sinkhorn-Knopp algorithm. This makes it particularly suitable for large-scale machine learning applications where exact computation of the Wasserstein distance is infeasible.

In summary, optimal transport theory provides a rigorous foundation for analyzing training dynamics through the minimization of transport costs between distributions. The Wasserstein distance serves as a natural objective for aligning model outputs with target distributions, while Sinkhorn divergences offer a practical and scalable alternative through entropy regularization. These tools enable the development of efficient and theoretically grounded algorithms for generative modeling, domain adaptation, and other machine learning tasks involving distributional alignment.

## 6.11 Optimal transport theory in Neural Architecture and Gradient Flows

Optimal transport theory provides a mathematical framework for comparing probability distributions by minimizing the cost of transporting mass from one distribution to another. The theory is rooted in the Monge-Kantorovich problem, which seeks a transport plan that minimizes the total cost of moving mass between two probability measures. Given two probability measures $\mu$ and $\nu$ on spaces $X$ and $Y$, respectively, and a cost function $c: X \times Y \to \mathbb{R}_+$, the optimal transport problem is formulated as:

$$\inf_{\pi \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\pi(x, y),$$

where $\Pi(\mu, \nu)$ denotes the set of all joint probability measures on $X \times Y$ with marginals $\mu$ and $\nu$. The solution to this problem defines the Wasserstein distance, which quantifies the minimal cost of transforming one distribution into another.

Training dynamics in machine learning can be analyzed through the lens of optimal transport, particularly when studying the evolution of model parameters or the distribution of data. Optimal transport-based training dynamics focus on minimizing the Wasserstein distance between the model's output distribution and the target distribution. This approach is especially useful in generative modeling, where the goal is to align the generated samples with the true data distribution. The gradient flow of the Wasserstein distance induces a dynamics that pushes the model parameters toward configurations that minimize the transport cost. For a parameterized model $P_{\theta}$, the optimization problem becomes:

$$\min_{\theta} W_c(P_{\theta}, Q),$$

where $Q$ is the target distribution and $W_c$ is the Wasserstein distance with cost function $c$. The gradient descent update rule for $\theta$ involves computing the gradient of the Wasserstein distance, which can be approximated using the dual formulation of the optimal transport problem.

Sinkhorn divergences provide a computationally efficient alternative to the Wasserstein distance by introducing entropy regularization into the optimal transport problem. The entropy-regularized
optimal transport problem is given by:

$$ \inf_{\pi \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\pi(x, y) + \epsilon H(\pi), $$

where $H(\pi)$ is the entropy of the transport plan $\pi$ and $\epsilon > 0$ is the regularization parameter. The Sinkhorn divergence is then defined as the difference between the regularized optimal transport cost and its self-transport terms:

$$ S_{\epsilon}(\mu, \nu) = W_{c, \epsilon}(\mu, \nu) - \frac{1}{2}W_{c, \epsilon}(\mu, \mu) - \frac{1}{2}W_{c, \epsilon}(\nu, \nu), $$

where $W_{c, \epsilon}$ denotes the entropy-regularized Wasserstein distance. The Sinkhorn divergence retains many desirable properties of the Wasserstein distance, such as metricity for certain cost functions, while being amenable to efficient computation using iterative scaling algorithms like the Sinkhorn-Knopp algorithm. This makes it particularly suitable for large-scale machine learning applications where exact computation of the Wasserstein distance is infeasible.

In summary, optimal transport theory provides a rigorous foundation for analyzing training dynamics through the minimization of transport costs between distributions. The Wasserstein distance serves as a natural objective for aligning model outputs with target distributions, while Sinkhorn divergences offer a practical and scalable alternative through entropy regularization. These tools enable the development of efficient and theoretically grounded algorithms for generative modeling, domain adaptation, and other machine learning tasks involving distributional alignment.

## 6.12 Geometric Regularization and Barycenters

Geometric regularization in optimal transport theory, particularly in the context of deep neural networks, refers to the imposition of structural constraints on the learned representations or model parameters to ensure they adhere to desirable geometric properties. This regularization leverages the Wasserstein distance, which inherently captures the underlying geometry of the data space, to penalize deviations from a prescribed geometric structure. Given a cost function $c: \mathcal{X} \times \mathcal{X} \to \mathbb{R}_+$ and two probability measures $\mu$ and $\nu$, the Wasserstein-$p$ distance is defined as:

$$ W_p(\mu, \nu) = \left( \inf_{\pi \in \Pi(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{X}} c(x, y)^p d\pi(x, y) \right)^{1/p}, $$

where $\Pi(\mu, \nu)$ is the set of all joint distributions with marginals $\mu$ and $\nu$. In deep learning, geometric regularization can be incorporated by adding a Wasserstein-based penalty term to the loss function, ensuring that the learned distributions or feature maps respect the metric structure of the input space. For instance, if $P_{\theta}$ is the model's output distribution and $Q$ is the target distribution, the regularized objective may take the form:

$$ \mathcal{L}(\theta) = \mathcal{L}_{\text{task}}(\theta) + \lambda W_p(P_{\theta}, Q), $$

where $\mathcal{L}_{\text{task}}$ is the task-specific loss and $\lambda$ controls the strength of the regularization. This approach is particularly useful in scenarios where preserving the geometric relationships between data points is critical, such as in manifold learning or domain adaptation.

Barycenters in optimal transport theory generalize the notion of averaging to probability measures on metric spaces, providing a way to compute a "mean" distribution that minimizes the total Wasserstein distance to a set of input distributions. Given a collection of probability measures
$\{\mu_i\}_{i=1}^N$ and weights $\{w_i\}_{i=1}^N$ such that $\sum_{i=1}^N w_i = 1$, the Wasserstein barycenter $\bar{\mu}$ is defined as the solution to:

$$ \bar{\mu} = \min_{\nu \in \mathcal{P}(\mathcal{X})} \sum_{i=1}^{N} w_i W_p^p(\mu_i, \nu), $$

where $\mathcal{P}(\mathcal{X})$ is the space of probability measures on $\mathcal{X}$. In deep neural networks, barycenters can be employed to aggregate information from multiple sources or to interpolate between distributions in a geometrically meaningful way. For example, in multi-task learning, the barycenter can represent a consensus distribution that balances the objectives of different tasks, while in generative modeling, it can be used to blend styles or features from diverse datasets. The computation of barycenters often relies on iterative algorithms, such as the iterative Bregman projection method or entropy-regularized approximations, which scale efficiently to high-dimensional spaces.

The integration of geometric regularization and barycenters into deep neural networks underscores the power of optimal transport theory in providing principled tools for incorporating geometric constraints into learning algorithms. By leveraging the Wasserstein distance, these methods ensure that the learned representations respect the intrinsic geometry of the data, leading to more robust and interpretable models. The barycenter framework, in particular, offers a flexible mechanism for combining distributions in a way that preserves their structural relationships, making it invaluable for applications ranging from collaborative learning to generative modeling. Together, these concepts highlight the deep interplay between geometry and probability in modern machine learning.
# 7 Categorical Foundations of Deep Learning
The categorical foundations of deep learning provide a unifying mathematical framework for modeling neural networks as compositional structures within category theory. At its core, this approach formalizes neural architectures as morphisms in a monoidal category, where objects represent data types (e.g., vector spaces $\mathbb{R}^n$) and morphisms correspond to parameterized functions (e.g., layers of a network). Given a symmetric monoidal category $(\mathcal{C}, \otimes, I)$, a neural network $f: X \to Y$ is constructed via sequential and parallel compositions of simpler morphisms, governed by the categorical axioms of associativity and unitality. For instance, the composition of layers $f: \mathbb{R}^n \to \mathbb{R}^m$ and $g: \mathbb{R}^m \to \mathbb{R}^k$ is represented by the morphism $g \circ f: \mathbb{R}^n \to \mathbb{R}^k$, satisfying $(h \circ g) \circ f = h \circ (g \circ f)$ for all compatible $h$. The monoidal product $\otimes$ captures parallel processing, as in convolutional or attention layers, where the tensor product of morphisms $f \otimes g: X_1 \otimes X_2 \to Y_1 \otimes Y_2$ represents independent transformations of parallel data streams. This framework generalizes backpropagation via functoriality: the derivative of a composite morphism $f \circ g$ is given by the chain rule, expressed categorically as the composition of the derivative functors $D(f \circ g) = D(f) \circ D(g)$, where $D: \mathcal{C} \to \mathcal{C}$ maps objects to their tangent bundles and morphisms to their Jacobians.

The optimization of neural networks is naturally expressed through universal constructions in category theory. The backpropagation algorithm arises as a right Kan extension in the category of parameterized functions. Let $\mathcal{P}$ be a category where objects are parameters $\theta \in \Theta$ and morphisms are updates $\theta \to \theta'$ induced by gradient descent. The loss functor $\mathcal{L}: \mathcal{P} \to \mathbb{R}$ sends each parameter configuration to its associated loss value, and the gradient $\nabla_{\theta\mathcal{L}}$ is the universal arrow completing the diagram:

<figure><img src="image_66.png" /></figure>

This universal property ensures that backpropagation optimally approximates the directional derivative of $\mathcal{L}$ with respect to the parameter update strategy. Furthermore, the universal approximation theorem is reformulated categorically as a density result: the hom-set $\text{Hom}(X, Y)$ in the category of neural networks is dense in the hom-set of continuous functions $C(X, Y)$ under the topology of uniform convergence on compacta.

The categorical lens formalism captures stateful computations in recurrent networks and memory-augmented architectures. A lens $\ell: (X, S) \to (Y, T)$ consists of a forward map $f: X \times S \to Y \times T$ and a backward map $b: X \times S \times Y \to T$, modeling both inference and gradient propagation. In a differentiable programming setting, lenses generalize automatic differentiation by encoding the full computation graph as a morphism in the category **Lens**, where composition is given by the chain rule for stateful functions:

$$ (f, b) \circ (g, d) = ((f \circ g), (x, s, z) \mapsto d(x, s, y) \cdot b(g(x), t, z)), $$

where $y, t = g(x, s)$. This perspective unifies architectures like LSTMs and neural ODEs, where the lens structure explicitly tracks hidden state dynamics. The hierarchical organization of deep networks is modeled using higher categories, where 2-morphisms represent transformations between layers (e.g., residual connections). For example, a ResNet block $f + \text{id}: X \to X$ is a 2-morphism in the bicategory **NN**, where horizontal composition corresponds to sequential layers and vertical composition to skip connections. Symmetries, such as permutation equivariance in graph neural networks, are captured by actions of symmetric monoidal categories **Sym**, ensuring that morphisms satisfy $f \circ \sigma = \sigma' \circ f$ for permutations $\sigma, \sigma'$.

By grounding deep learning in category theory, these foundations reveal the intrinsic algebraic and geometric structures underlying neural architectures, enabling rigorous analysis of their expressivity, optimization, and generalization.
# 7.1 Introduction

## 7.1.1 Motivation: Why Category Theory for Deep Learning?

Deep Learning has achieved remarkable success across domains—from computer vision to natural language processing—yet its theoretical foundations remain fragmented. Traditional mathematical tools (linear algebra, calculus, probability) provide only a local understanding of neural networks, failing to capture their compositional, structural, and relational nature.

Category Theory offers a unifying framework for Deep Learning by:

1. **Abstraction over computation:** Neural networks are fundamentally morphisms (maps) between data representations, and Category Theory provides tools to study compositions of such maps.

2. **Generalization across architectures:** Whether working with CNNs, RNNs, or Transformers, Category Theory allows us to describe them as instances of categorical constructs (e.g., monoidal categories, functors).

3. **Handling complex data flows:** Many Deep Learning constructs (e.g., attention mechanisms, residual connections) are naturally expressed via universal properties (limits, adjunctions).

4. **Bridging discrete and continuous reasoning:** Gradient-based optimization (backpropagation) can be formulated categorically via lenses or reverse derivative categories.

Thus, Category Theory provides a meta-language for Deep Learning, enabling:

1. Formal reasoning about model behavior.

2. Systematic architecture design via compositionality.

3. Novel generalizations (e.g., neural networks over graphs, higher-order interactions).

## 7.1.2 Limits of Traditional Mathematical Formalism in Modern Machine Learning

While traditional mathematics has been indispensable for Deep Learning, key limitations arise when scaling to modern systems:

1. **Lack of Compositionality:** Deep Learning models are built via hierarchical composition of layers, but linear algebra treats them as sequences of matrix operations without inherent structure. The Category Theory Solution is to represent networks as morphisms in a category, where composition is associative by definition.

2. **Poor Handling of Complex Data Types:** Tensors, graphs, and sequences require ad-hoc notation. Backpropagation, for instance, is often derived via index manipulation rather than structurally. The Category Theory Solution is to use monoidal categories to model tensors and string diagrams for graphical reasoning.

3. **Opaque Optimization Dynamics:** Gradient descent is typically analyzed via local approximations (Taylor expansions), obscuring global behavior. The Category Theory Solution is to use reverse derivative categories (RDCs) to give an abstract formulation of backpropagation.

4. **Difficulty in Generalizing Architectures:** Novel components (e.g., attention, memory) are introduced empirically without a unifying framework. The Category Theory Solution is to use universal properties (e.g., adjunctions) to guide principled extensions.
5. **Weak Theoretical Guarantees**: Traditional theory (e.g., VC dimension) fails to explain the success of overparametrized Deep Learning models. The Category Theory Solution is to use functorial semantics to connect learning to algebraic invariants.

### 7.1.3 Overview of the Chapter and Learning Goals

Category Theory provides a powerful framework for understanding and designing Deep Learning models by abstracting their compositional structure and behavior. This chapter develops CT for DL systematically, with the following structure:

1. **Categories and Functors**: At its core, a category formalizes the notion of objects (such as data spaces) and morphisms (such as neural network layers) that can be composed as- sociatively. Functors extend this by mapping one category to another, preserving struc- ture—allowing us to describe how neural architectures transform data systematically. We shall discuss the basic definitions, examples (e.g., FeedForwardNet, the category of neural networks).

2. **Universal Constructions**: Universal constructions, including products and limits, char- acterize optimal ways to combine or relate objects. In Deep Learning, these appear in ar- chitectures like residual networks, where skip connections arise as coproducts, or attention mechanisms, which aggregate information universally. The ability to reason about such con- structions enables principled architecture design rather than relying on intuition alone. We shall discuss the products, coproducts, and how they relate to branching architectures (e.g., ResNet skip connections).

3. **Monoidal Categories**: Monoidal categories introduce tensor products, capturing parallel computation and multi-dimensional data flow. This framework elegantly models operations in convolutional networks or transformers, where tensor contractions and parallel processing are fundamental. String diagrams, a graphical notation from monoidal categories, offer an intuitive way to visualize and manipulate complex network structures. We shall discuss the modeling of tensors and parallel computation (e.g., CNNs as functors).

4. **Lenses and Backpropagation**: Differentiation and optimization, central to Deep Learning, are naturally expressed using lenses and reverse derivative categories. A lens pairs a forward pass (evaluation) with a backward pass (gradient propagation), abstracting backpropagation as a compositional process. This shifts gradient-based learning from an algorithmic procedure to a categorical construct, clarifying its mathematical essence. We shall discuss the categorical differentiation via reverse derivative categories.

5. **Higher-Order Abstractions**: Higher-order abstractions like adjunctions and monads fur- ther extend the expressive power of the framework. Adjunctions formalize relationships between paired functors, such as encoder-decoder networks in autoencoders, while monads model computational effects like stochasticity in dropout or reinforcement learning. These tools allow us to describe complex behaviors—such as memory, recursion, or attention—in a unified way. We shall discuss the adjunctions, monads, and their role in attention/ memory mechanisms.

Category theory does not replace traditional ML theory but extends it with a higher-level language for compositionality, abstraction, and generalization. By adopting CT, we gain tools to:

1. Design models more systematically.

2. Analyze them more rigorously.

3. Discover new architectures more confidently.
The next sections will make these ideas precise, starting with the definition of a category and its applications to DL. Together, these categorical concepts provide a rigorous foundation for Deep Learning, transforming empirical practices into structured mathematical principles. By viewing neural networks as morphisms in appropriate categories, we gain a language for compositionality, optimization, and generalization—bridging the gap between abstract theory and practical implementation. This perspective not only clarifies existing architectures but also guides the invention of new ones, offering a pathway toward more systematic and interpretable machine learning.

## 7.2 Category Theory Primer for Machine Learners

We shall now discuss the foundational mathematical framework to describe the structure and compositionality inherent in machine learning models through the language of categories. At its core, a category $\mathcal{C}$ consists of a collection of objects (e.g., data spaces, model architectures) and morphisms (e.g., parameterized functions, layers) between these objects, equipped with an associative composition operation and identity morphisms. For machine learning, the category **Learn** is often informally considered, where objects are spaces of data representations (e.g., $\mathbb{R}^n$ for $n$-dimensional feature vectors) and morphisms $f_\theta: X \to Y$ are learnable maps (e.g., neural network layers) parameterized by $\theta \in \Theta$. Composition of morphisms $f_\theta: X \to Y$ and $g_\phi: Y \to Z$ yields a new morphism $g_\phi \circ f_\theta: X \to Z$, corresponding to the forward pass of a neural network:

$$ (g_\phi \circ f_\theta)(x) = g_\phi(f_\theta(x)), \quad \text{with parameters } (\theta, \phi). $$

The identity morphism $\text{id}_X: X \to X$ represents the trivial transformation $\text{id}_X(x) = x$, though in practice, exact identities may be approximated due to architectural constraints. Universal constructions in category theory provide templates for common machine learning paradigms. The product of objects $X$ and $Y$, denoted $X \times Y$, models parallel data processing (e.g., multi-modal inputs), with projection morphisms $\pi_X: X \times Y \to X$ and $\pi_Y: X \times Y \to Y$ extracting components. Dually, the coproduct $X \sqcup Y$ represents disjoint unions (e.g., batch concatenation). Limits and colimits generalize these notions to arbitrary diagrams, capturing architectures like attention mechanisms, where the output is a weighted limit of input features. For instance, the attention operation for queries $Q$, keys $K$, and values $V$ can be viewed as a colimit construction:

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V, $$

where the softmax computes a convex combination of values, analogous to a categorical colimit over the alignment scores.

Functors map one category to another while preserving structure, offering a formalism for model transformations and embeddings. A loss functor $L: \mathbf{Learn} \to \mathbf{Met}$ might map learnable morphisms to their associated loss landscapes, where $\mathbf{Met}$ is a category of metric spaces. Natural transformations between functors describe parameter updates: for a model $F_\theta: \mathbf{Data} \to \mathbf{Pred}$, gradient descent induces a natural transformation $\eta: F_\theta \Rightarrow F_{\theta-\alpha\nabla_\theta\mathcal{L}}$, where $\mathcal{L}$ is the loss function and $\alpha$ the learning rate. Adjunctions formalize optimization paradigms, such as the free-forgetful adjunction between unstructured data (e.g., raw images) and structured representations (e.g., latent embeddings), where the left adjoint $F$ constructs a representation and the right adjoint $U$ forgets it:

$$ \text{Hom}_{\text{Learn}}(F(X), Y) \cong \text{Hom}_{\text{Data}}(X, U(Y)). $$

Monoidal categories capture parallel computation in neural networks. The tensor product $X \otimes Y$ represents parallelized data (e.g., channels in a convolutional layer), with morphisms $f \otimes g: X \otimes Y \to Z \otimes W$ acting independently on subsystems. The interchange law $(f \otimes g) \circ (h \otimes k) = (f \circ h) \otimes (g \circ k)$ reflects the locality of operations like convolutions. String diagrams provide a graphical calculus
for such settings, where wires represent data flow and boxes represent layers, simplifying proofs of network properties. Abstraction and generalization in machine learning are formalized via universal properties. For example, an universal approximator (e.g., an MLP) satisfies a universal property: for any continuous $f: X \to Y$ and $\epsilon > 0$, there exists a morphism (neural network) $f_\theta: X \to Y$ such that $\|f(x) - f_\theta(x)\| < \epsilon$ on compact $K \subset X$. This mirrors the universal property of polynomial approximation in Stone-Weierstrass theorem. Sheaf theory further refines this by modeling local-to-global learning (e.g., graph neural networks), where data is structured as a presheaf and predictions are glued from local patches.

In summary, category theory provides a unifying language for machine learning, where models are morphisms, training is functorial optimization, and architectures are universal constructions. This framework bridges abstract mathematical principles with concrete implementations, revealing the compositional essence of learning systems while accommodating their statistical and approximate nature.

## 7.2.1 Basic Definitions: Categories, Functors, Natural Transformations, Adjunctions, Monoidal categories

### 7.2.1.1 Categories

A category $\mathcal{C}$ consists of:

* A collection of **objects** $\text{Ob}(\mathcal{C})$ (e.g., sets, vector spaces).

* For every pair $X, Y \in \text{Ob}(\mathcal{C})$, a set $\text{Hom}_\mathcal{C}(X, Y)$ of **morphisms** (arrows) $f: X \to Y$.

* A **composition rule** $\circ$ such that for $f: X \to Y$ and $g: Y \to Z$, there exists $g \circ f: X \to Z$.

* An **identity morphism** $\text{id}_X: X \to X$ for each object, satisfying:

$$ \forall f: X \to Y, \quad f \circ \text{id}_X = f = \text{id}_Y \circ f. $$

* **Associativity**: For $f: X \to Y, g: Y \to Z, h: Z \to W$,

$$ h \circ (g \circ f) = (h \circ g) \circ f. $$

A category is a fundamental mathematical structure that formalizes the notion of mathematical objects and the relationships between them through a precise algebraic framework. Formally, a category $\mathcal{C}$ consists of two primary components: a collection of objects and a collection of morphisms satisfying specific axiomatic conditions. The objects, denoted by $\text{Ob}(\mathcal{C})$, represent the entities of interest within the category, which could be sets, vector spaces, topological spaces, or any other mathematically definable constructs. The morphisms, also called arrows, represent structure-preserving maps between these objects. For every pair of objects $X, Y \in \text{Ob}(\mathcal{C})$, there exists a set $\text{Hom}_\mathcal{C}(X, Y)$ whose elements are the morphisms from $X$ to $Y$. Each morphism $f \in \text{Hom}_\mathcal{C}(X, Y)$ is denoted as $f: X \to Y$, where $X$ is the domain and $Y$ is the codomain.

The structure of a category is governed by two fundamental operations: composition and identity. Composition is a binary operation that takes two morphisms $f: X \to Y$ and $g: Y \to Z$ and produces a composite morphism $g \circ f: X \to Z$. This operation must satisfy the associativity axiom, which states that for any three morphisms $f: X \to Y, g: Y \to Z$, and $h: Z \to W$, the following equality holds:

$$ h \circ (g \circ f) = (h \circ g) \circ f. $$

This axiom ensures that the order in which compositions are evaluated does not affect the result, making the algebraic structure coherent. The second operation is the existence of an identity
morphism for every object. Specifically, for each object $X \in \text{Ob}(\mathcal{C})$, there exists a distinguished morphism $\text{id}_X : X \to X$ such that for any morphism $f : X \to Y$, the identities act as neutral elements for composition:

$$f \circ \text{id}_X = f \quad \text{and} \quad \text{id}_Y \circ f = f.$$

These axioms collectively ensure that the morphisms behave consistently under composition and that each object has a trivial self-map.

The power of category theory lies in its ability to abstract and generalize mathematical concepts across diverse fields. For instance, in the category **Set**, the objects are sets, and the morphisms are ordinary functions between sets. In $\text{Vect}_K$, the objects are vector spaces over a field $K$, and the morphisms are linear transformations. The categorical framework allows these seemingly distinct mathematical domains to be studied under a unified language, where universal properties, functors, and natural transformations become the tools for analysis. The abstract nature of categories enables mathematicians to distill the essence of mathematical structures, focusing on their relational properties rather than their internal constitution. This perspective has profound implications in areas ranging from algebraic topology to theoretical computer science and, more recently, in the formalization of machine learning architectures.

The notion of commutativity is central to category theory and is often visualized through commutative diagrams. A diagram in a category is said to commute if all paths between any two objects yield the same morphism under composition. For example, given morphisms $f : X \to Y$ and $g : Y \to Z$, the commutativity of the diagram:

$$\begin{array}{ccc}
X & \xrightarrow{f} & Y \\
\bigcirc \circled{g} & & \bigcirc \circled{g} \\
Z & & Z
\end{array}$$

expresses the equality $g \circ f = g \circ f$, which trivially holds but illustrates the diagrammatic reasoning pervasive in category theory. More complex diagrams encode richer algebraic conditions, such as universal properties, which characterize objects like products, coproducts, and limits up to unique isomorphism. These properties are foundational in constructing and understanding categorical structures across mathematics.

In summary, a category is an algebraic structure that abstracts the relationships between mathematical objects through morphisms, composition, and identities, governed by axioms ensuring associativity and the existence of neutral elements. This abstraction provides a powerful language for unifying diverse mathematical disciplines and has become indispensable in modern theoretical frameworks, including those in machine learning. The rigorous formulation of categories enables precise reasoning about structure-preserving mappings, universal properties, and hierarchical constructions, making it a cornerstone of contemporary mathematics.

### 7.2.1.2 Functors

A (covariant) functor $F : \mathcal{C} \to \mathcal{D}$ between categories consists of:

* A mapping $X \mapsto F(X)$ on objects.

* For each morphism $f : X \to Y$ in $\mathcal{C}$, a morphism $F(f) : F(X) \to F(Y)$ in $\mathcal{D}$, preserving:

- **Identities**: $F(\text{id}_X) = \text{id}_F(X)$.

- **Composition**: $F(g \circ f) = F(g) \circ F(f)$.
A contravariant functor reverses arrows: $F(f) : F(Y) \to F(X)$.

A functor is a structure-preserving mapping between categories that serves as the fundamental notion of morphism in the realm of category theory. Formally, given two categories $\mathcal{C}$ and $\mathcal{D}$, a (covariant) functor $F : \mathcal{C} \to \mathcal{D}$ consists of two components that respect the categorical structure. First, it associates to each object $X \in \text{Ob}(\mathcal{C})$ an object $F(X) \in \text{Ob}(\mathcal{D})$. Second, to each morphism $f : X \to Y$ in $\mathcal{C}$, it assigns a morphism $F(f) : F(X) \to F(Y)$ in $\mathcal{D}$. This assignment must satisfy the twin conditions of preserving identity morphisms and composition of morphisms. For every object $X \in \mathcal{C}$, the functor preserves the identity morphism as follows:

$$F(\text{id}_X) = \text{id}_{F(X)},$$

ensuring that the identity structure is maintained under the mapping. For composable morphisms $f : X \to Y$ and $g : Y \to Z$ in $\mathcal{C}$, the functor preserves composition through the equation:

$$F(g \circ f) = F(g) \circ F(f),$$

which guarantees that the image of a composite morphism equals the composite of the images, thereby maintaining the compositional structure of the original category.

The contravariant functor represents a dual notion where the direction of morphisms is reversed. Specifically, a contravariant functor $F : \mathcal{C} \to \mathcal{D}$ maps each morphism $f : X \to Y$ in $\mathcal{C}$ to a morphism $F(f) : F(Y) \to F(X)$ in $\mathcal{D}$, reversing the arrow direction while still preserving identities and composition. The composition preservation for contravariant functors takes the form:

$$F(g \circ f) = F(f) \circ F(g),$$

where the order of composition is inverted compared to the covariant case. This reversal captures important mathematical phenomena such as duality in vector spaces, where the dual space functor $(-)^* : \text{Vect}_k \to \text{Vect}_k$ sends each linear map $T : V \to W$ to its transpose $T^* : W^* \to V^*$.

Functors serve as the primary vehicles for transporting mathematical structures between different categorical contexts. The power of functoriality lies in its ability to systematically relate disparate mathematical domains while preserving their essential structural relationships. For instance, the fundamental group functor $\pi_1 : \text{Top}_* \to \text{Grp}$ maps pointed topological spaces to their fundamental groups and continuous maps to group homomorphisms, translating topological problems into algebraic ones. Similarly, in algebraic geometry, the spectrum functor $\text{Spec} : \text{CRing}^{\text{op}} \to \text{Top}$ establishes a bridge between commutative rings and topological spaces through the contravariant correspondence of prime ideals and Zariski topology.

The naturality condition for functors becomes particularly significant when considering their behavior with respect to transformations between categories. A natural transformation $\eta : F \Rightarrow G$ between two functors $F, G : \mathcal{C} \to \mathcal{D}$ consists of a family of morphisms $\eta_X : F(X) \to G(X)$ for each $X \in \mathcal{C}$ that commute with the images of morphisms under the functors. This means that for every morphism $f : X \to Y$ in $\mathcal{C}$, the following diagram commutes:

$$\begin{tikzcd}
F(X) \arrow[r, "\eta_X"] & G(X) \\
F(f) \arrow[u, "\eta_Y"] \arrow[d, "G(f)"'] & \\
F(Y) & G(Y)
\end{tikzcd}$$

expressing the coherence condition $G(f) \circ \eta_X = \eta_Y \circ F(f)$. Natural transformations formalize the notion of consistent translation between different functorial perspectives on the same categorical data.
The ubiquity of functors in modern mathematics stems from their capacity to systematically capture structural parallels across diverse mathematical domains. In representation theory, group representations are understood through functors from the delooping of a group to vector spaces. In algebraic topology, homology and cohomology theories are realized as functors from topological categories to graded algebraic structures. The categorical language of functors provides a unifying framework that reveals deep connections between apparently unrelated mathematical constructions, while the rigorous preservation of composition and identities ensures the structural integrity of these connections. This makes functors indispensable tools not only in pure mathematics but also in theoretical computer science, where they model type constructors and computational effects, and in mathematical physics, where they describe symmetries and state space transformations.

### 7.2.1.3 Natural Transformations

Given functors $F, G : \mathcal{C} \to \mathcal{D}$, a **natural transformation** $\eta : F \Rightarrow G$ assigns to each $X \in \mathcal{C}$ a morphism $\eta_X : F(X) \to G(X)$ such that for every $f : X \to Y$ in $\mathcal{C}$, the following **naturality condition** holds:

$$G(f) \circ \eta_X = \eta_Y \circ F(f).$$

This ensures that $\eta$ "transformats $F$ into $G$ coherently."

A natural transformation represents a systematic way to transform one functor into another while respecting the underlying categorical structure. Formally, given two parallel functors $F, G : \mathcal{C} \to \mathcal{D}$ between the same categories, a natural transformation $\eta : F \Rightarrow G$ consists of a family of morphisms indexed by the objects of $\mathcal{C}$, where for each object $X \in \mathcal{C}$, there exists a component morphism $\eta_X : F(X) \to G(X)$ in $\mathcal{D}$. These components must satisfy the naturality condition, which ensures that for every morphism $f : X \to Y$ in $\mathcal{C}$, the following diagram commutes:

$$\begin{tikzcd}
F(X) \arrow[r, "\eta_X"] & G(X) \\
F(f) \arrow[u, "\eta_Y"] \arrow[d, "G(f)"'] & \\
F(Y) & G(Y)
\end{tikzcd}$$

This diagrammatic condition translates algebraically to the equation $G(f) \circ \eta_X = \eta_Y \circ F(f)$, expressing the coherence of the transformation across all morphisms in the category. The naturality condition guarantees that the transformation $\eta$ behaves uniformly with respect to the action of the functors on both objects and morphisms, making it a proper morphism between functors in the functor category $\mathcal{D}^c$.

The concept of naturality captures the mathematical intuition of "canonical" or "universally defined" transformations that depend solely on the categorical structure, without arbitrary choices. For instance, in the category of vector spaces over a field $k$, the double dual embedding $\eta_V : V \to V^{**}$ defined by $\eta_V(v)(\phi) = \phi(v)$ for $v \in V$ and $\phi \in V^*$ forms a natural transformation from the identity functor to the double dual functor. The naturality of this transformation means that for any linear map $T : V \to W$, the following equality holds:

$$T^{**} \circ \eta_V = \eta_W \circ T,$$

where $T^{**}$ denotes the double dual map induced by $T$. This equation verifies that the embedding commutes with the application of linear maps, demonstrating that the transformation respects the linear structure across all vector spaces and their morphisms.

Natural transformations play a pivotal role in category theory by enabling the comparison of functors and the formulation of universal properties. They serve as the 2-morphisms in the 2-category
of categories, where objects are categories, 1-morphisms are functors, and 2-morphisms are natural transformations. This hierarchical structure allows for the expression of sophisticated mathematical concepts, such as adjunctions and equivalences between categories, which can be characterized through natural isomorphisms—natural transformations where every component $\eta_X$ is an isomorphism. For example, an adjunction between functors $F: C \to D$ and $G: D \to C$ consists of a pair of natural transformations $\eta: \mathrm{Id}_C \Rightarrow GF$ and $\epsilon: FG \Rightarrow \mathrm{Id}_D$ satisfying the triangle identities:

$$ \epsilon F \circ F \eta = \mathrm{id}_F \quad \text{and} \quad G \epsilon \circ \eta G = \mathrm{id}_G. $$

These identities ensure the coherence of the adjunction, providing a rigorous framework for understanding dualities and universal constructions across mathematics. In algebraic topology, for instance, the Hurewicz transformation from homotopy groups to homology groups constitutes a natural transformation that preserves the topological structure, facilitating the translation between homotopy-theoretic and homological data. Similarly, in algebraic geometry, the sheafification functor comes equipped with a natural transformation from the identity functor to the sheafification functor, reflecting the universal property of sheafification.

The power of natural transformations lies in their ability to formalize the notion of structure-preserving mappings between different representations of the same categorical data. This capability is crucial in areas such as functional programming, where natural transformations between type constructors ensure the correctness of polymorphic functions, and in theoretical physics, where they model symmetries and dualities between physical theories. By providing a rigorous language for comparing and transforming functorial constructions, natural transformations unify diverse mathematical phenomena under a common conceptual framework, revealing deep connections between seemingly unrelated areas of mathematics and science.

### 7.2.1.4 Adjunctions

Adjunctions in Category Theory formalize profound relationships between pairs of functors, capturing a notion of "approximate equivalence" or "duality" between categories. Given two categories $C$ and $D$, an adjunction consists of functors $F: C \to D$ (the left adjoint) and $G: D \to C$ (the right adjoint), together with a natural isomorphism between hom-sets:

$$ \mathrm{Hom}_D(F(X), Y) \cong \mathrm{Hom}_C(X, G(Y)) $$

for all objects $X \in C$ and $Y \in D$. This isomorphism must be natural in both $X$ and $Y$, meaning that for any morphisms $f: X' \to X$ in $C$ and $g: Y \to Y'$ in $D$, the following diagrams commute:

$$ \begin{array}{ccc} \mathrm{Hom}_D(F(X), Y) & \xrightarrow{\sim} & \mathrm{Hom}_C(X, G(Y)) \\ F(f)^* \downarrow & & \downarrow f^* \\ \mathrm{Hom}_D(F(X'), Y) & \xrightarrow{\sim} & \mathrm{Hom}_C(X', G(Y)) \end{array} \qquad \begin{array}{ccc} \mathrm{Hom}_D(F(X), Y) & \xrightarrow{\sim} & \mathrm{Hom}_C(X, G(Y)) \\ g_* \downarrow & & \downarrow G(g)_* \\ \mathrm{Hom}_D(F(X), Y') & \xrightarrow{\sim} & \mathrm{Hom}_C(X, G(Y')) \end{array} $$

Here, $F(f)^*$ denotes precomposition with $F(f)$, and $g_*$ denotes postcomposition with $g$. An equivalent formulation uses the unit $\eta: \mathrm{Id}_C \Rightarrow G \circ F$ and $\epsilon: F \circ G \Rightarrow \mathrm{Id}_D$, which are natural transformations satisfying the triangle identities:

$$ \epsilon F \circ F \eta = \mathrm{Id}_F \quad \text{and} \quad G \epsilon \circ \eta G = \mathrm{Id}_G. $$

The unit $\eta_X: X \to G(F(X))$ universalizes the "embedding" of $X$ into its image under the adjunction, while the counit $\epsilon_Y: F(G(Y)) \to Y$ evaluates or "realizes" the abstractly constructed object $F(G(Y))$. For example, in the adjunction between the free group functor $F: \mathbf{Set} \to \mathbf{Grp}$ and the forgetful functor $U: \mathbf{Grp} \to \mathbf{Set}$, the unit injects a set $X$ into its free group, and the counit maps a group back to its underlying set.
Adjunctions generalize universal constructions: limits/colimits arise as adjoints to diagonal functors, and tensor-hom adjunctions in closed categories ($\operatorname{Hom}(X \otimes Y, Z) \cong \operatorname{Hom}(X, \operatorname{Hom}(Y, Z))$) encode currying. In machine learning, adjunctions model optimization (e.g., latent representation learning), where $F$ constructs features and $G$ decodes them, with the adjunction ensuring coherence between encoding and reconstruction. The naturality of the hom-set isomorphism reflects the stability of these processes under transformations of inputs or outputs. Adjunctions also underlie Kan extensions, where left/right adjoints to precomposition functors extend diagrams universally. For a functor $F: C \to D$ and category $\mathcal{E}$, the left Kan extension $\operatorname{Lan}_F: [C, \mathcal{E}] \to [D, \mathcal{E}]$ is the left adjoint to precomposition $F^*: [D, \mathcal{E}] \to [C, \mathcal{E}]$, providing a way to "complete" partial diagrams. This abstract machinery clarifies phenomena like gradient descent (a Kan extension of loss functions along optimization paths) and attention mechanisms (extensions of local feature interactions).

In essence, adjunctions are the "optimal" relationships between categories, where one functor approximates the inverse of another, governed by universal properties and natural coherence conditions. They unify seemingly disparate constructions across mathematics and machine learning, from algebraic duality to neural network training dynamics, through a precise language of compositional equivalence.

### 7.2.1.5 Monoidal categories

Monoidal categories provide a rigorous mathematical framework for capturing parallel computation in neural networks by formalizing the algebraic structure of simultaneous, independent operations across multiple data streams. At the core of this formalism lies the tensor product $\otimes$, which represents the parallel combination of objects (data spaces) and morphisms (network layers). In the category $\mathbf{NN}$ of neural networks, objects are vector spaces $\mathbb{R}^n$ (or more generally, manifolds of activations), and the tensor product $\mathbb{R}^n \otimes \mathbb{R}^m \cong \mathbb{R}^{n+m}$ corresponds to the concatenation of parallel data streams—such as the channels in a convolutional layer or the heads in a multi-head attention mechanism. Morphisms $f \otimes g: \mathbb{R}^n \otimes \mathbb{R}^k \to \mathbb{R}^m \otimes \mathbb{R}^l$ between these combined spaces act independently on their respective subspaces, defined for $x \in \mathbb{R}^n$ and $y \in \mathbb{R}^k$ by:

$$ (f \otimes g)(x \otimes y) = f(x) \otimes g(y), $$

where $f(x) = \sigma(W_f x + b_f)$ and $g(y) = \sigma(W_g y + b_g)$ are individual layer operations with weights $W_f, W_g$ and biases $b_f, b_g$. This operation preserves the separability of parallel computations until explicitly combined by subsequent layers.

The monoidal unit $I = \mathbb{R}^0$ (a trivial space) serves as the identity for $\otimes$, with isomorphisms $\mathbb{R}^n \otimes I \cong \mathbb{R}^n \cong I \otimes \mathbb{R}^n$ reflecting that augmenting a network with empty data leaves it unchanged. The associativity of $\otimes$ is captured by the natural isomorphism $\alpha_{A,B,C}: (A \otimes B) \otimes C \to A \otimes (B \otimes C)$, which permits the re-parenthesization of parallel streams—a property critical for modular network architectures like ResNets, where skip connections and branches must coherently merge. For neural networks, this translates to the equivalence of processing pipelines like $(f \otimes g) \otimes h$ and $f \otimes (g \otimes h)$, ensuring that the order of combining parallel paths does not affect the final output. The interchange law governs the interaction between sequential and parallel composition, stating that for morphisms $f, h: \mathbb{R}^n \to \mathbb{R}^m$ and $g, k: \mathbb{R}^k \to \mathbb{R}^l$, the following equality holds:

$$ (h \circ f) \otimes (k \circ g) = (h \otimes k) \circ (f \otimes g). $$

This law ensures that parallel and sequential operations commute in a well-defined manner, mirroring the locality of operations in convolutional networks or attention layers, where filters or attention heads process inputs independently before aggregation. Violations of this law due to non-linearities (e.g., activation functions) are mitigated by the approximate linearity of learned representations in
practice.

Symmetric monoidal categories further introduce a braiding isomorphism $\beta_{A,B}: A \otimes B \to B \otimes A$, which models the permutation of parallel data channels. In neural networks, this captures operations like channel shuffling in ShuffleNet or the reordering of attention heads, where $\beta$ swaps subspaces without altering the underlying computation. The coherence conditions for $\beta$—such as $\beta_{B,A} \circ \beta_{A,B} = \mathrm{id}_{A \otimes B}$—guarantee that permuting and reverting returns the original configuration, a property leveraged in equivariant networks designed for symmetric inputs. The closure property of monoidal categories, when present, internalizes the space of parallelizable operations. In NN, the hom-object $[A, B]$ represents the space of all possible layers from $A$ to $B$, and the isomorphism $\mathrm{Hom}(A \otimes B, C) \cong \mathrm{Hom}(A, [B, C])$ encodes the currying of parallel operations—a abstraction foundational to attention mechanisms, where keys, queries, and values are processed in parallel before interaction. The evaluation morphism $\mathrm{eval}: [A, B] \otimes A \to B$ then corresponds to the application of an attention head to its inputs.

In summary, monoidal categories rigorously axiomatize the parallel computation in neural networks through the tensor product $\otimes$, its associativity, unit, and interchange laws, while symmetry and closure properties model advanced architectural features. This framework abstracts the algebraic essence of parallelism, providing a unified language to reason about complex network architectures and their compositional invariants.

## 7.2.2 Use of String diagrams and Abstraction and Generalization in Machine Learning

### 7.2.2.1 String diagrams as Graphical Calculus for Neural Networks

String diagrams furnish a rigorous graphical calculus for reasoning about monoidal categories, offering an intuitive yet mathematically precise representation of neural network architectures where wires denote data flow (objects) and boxes (nodes) represent layers (morphisms). The formalism is grounded in the Penrose graphical notation, which translates categorical axioms into topological equivalences of diagrams, enabling syntactic manipulations that preserve semantic meaning. In the context of neural networks, a layer $f: \mathbb{R}^n \to \mathbb{R}^m$ is depicted as a box with an input wire (of type $\mathbb{R}^n$) entering from below and an output wire (of type $\mathbb{R}^m$) exiting above:

$$ \bullet \xrightarrow{\mathbb{R}^n} [f] \xrightarrow{\mathbb{R}^m} \bullet $$

Composition of layers $g \circ f$ corresponds to vertical stacking of boxes, while the tensor product $f \otimes g$ for parallel layers is represented by horizontal juxtaposition:

Sequential composition:

$$ \bullet \xrightarrow{\mathbb{R}^n} [f] \xrightarrow{\mathbb{R}^k} [g] \xrightarrow{\mathbb{R}^m} \bullet $$

Parallel composition:

$$ \bullet \xrightarrow{\mathbb{R}^n} [f] \xrightarrow{\mathbb{R}^m} \bullet $$

$$ \bullet \xrightarrow{\mathbb{R}^k} [g] \xrightarrow{\mathbb{R}^l} \bullet $$

The interchange law $(h \circ f) \otimes (k \circ g) = (h \otimes k) \circ (f \otimes g)$ becomes a geometric triviality in string diagrams, as the equality manifests as isotopic deformation of the graph (sliding boxes past each other without cutting wires):
<figure><img src="image_67.png" /></figure>

**Coherence conditions** of monoidal categories (associativity, unit laws) are implicitly enforced by the diagrammatic syntax. For instance, associativity of $\otimes$ $(\alpha_{A,B,C}: (A \otimes B) \otimes C \rightarrow A \otimes (B \otimes C))$ corresponds to the absence of parentheses in diagrams-wires may freely cross or bend, as the underlying topological equivalence classes abstract over syntactic re-bracketing. The unit object $I$ is represented by the absence of a wire, and the unit isomorphisms $\lambda_A: I \otimes A \rightarrow A$ and $\rho_A: A \otimes I \rightarrow A$ are rendered as "yanking" wires straight:

<figure><img src="image_68.png" /></figure>

**Symmetry** in braided monoidal categories introduces crossings of wires, represented by the natural isomorphism $\sigma_{A,B}: A \otimes B \rightarrow B \otimes A$. In neural networks, this captures operations like channel shuffling or attention head permutation, where $\sigma$ swaps two parallel data streams:

<figure><img src="image_69.png" /></figure>

**Evaluation and currying** in closed monoidal categories are depicted as caps and cups. For a layer $f: A \otimes B \rightarrow C$, its curried form $\tilde{f}: A \rightarrow [B, C]$ is represented by bending the input wire $B$ into an output:

<figure><img src="image_70.png" /></figure>

**Neural network properties** reduce to diagrammatic simplifications. For example, the universality of a multilayer perceptron (MLP) corresponds to the fact that any diagram of wires and boxes can be approximated (up to $\epsilon$) by a sufficiently large subdiagram of MLP layers. Residual connections $f + \text{id}$ are represented as a fork where the identity wire bypasses the box for $f$:

<figure><img src="image_71.png" /></figure>

The backpropagation algorithm arises as the dual (mirror image) of the forward pass diagram, with gradients flowing backward along reversed wires-a manifestation of the **time-reversal symmetry** in string diagrams for autonomous categories.

In essence, string diagrams distill the algebraic structure of neural networks into topological invariants, where equivalence of diagrams under deformation encodes equality of morphisms. This graphical calculus not only simplifies proofs of network properties (e.g., universality, equivariance) but also guides architectural design through visual intuition, bridging categorical abstraction and engineering practice.
### 7.2.2.2 Abstraction and Generalization in Machine Learning

Abstraction and generalization in machine learning are rigorously formalized through universal properties in category theory, which characterize mathematical objects by their relationships to all other objects in a category rather than by their internal structure. A universal property specifies that a particular construction is the "most efficient" solution to a given problem, unique up to isomorphism. In machine learning, this manifests when a model architecture or learning algorithm satisfies conditions that make it optimally representative or predictive across a class of problems.

For a learning task where data $X$ is mapped to targets $Y$ through some ground truth function $f: X \to Y$, a universal approximator $F_\theta: X \to Y$ (e.g., a neural network) satisfies the following universal property: for any $\epsilon > 0$ and compact $K \subset X$, there exists a parameterization $\theta$ such that:

$$ \sup_{x \in K} \|F_{\theta}(x) - f(x)\| < \epsilon. $$

This universality is analogous to the universal property of free objects in algebra, where the free group $F(S)$ on a set $S$ satisfies: for any group $G$ and map $f: S \to G$, there exists a unique homomorphism $\tilde{f}: F(S) \to G$ extending $f$. Here, $F(S)$ is the "most general" group containing $S$, just as a universal approximator is the "most general" function approximator for $f$. The universal property of optimization formalizes gradient descent as a universal arrow in the category of loss landscapes. Given a loss functor $\mathcal{L}: \mathbf{Params} \to \mathbf{Met}$, mapping parameters $\theta$ to metric spaces of their loss values, gradient descent constructs a universal morphism to the global minimum $\theta^*$: for any initial $\theta_0$, there exists a sequence of updates $\theta_{k+1} = \theta_k - \alpha \nabla_{\theta} \mathcal{L}$ converging to $\theta^*$ under appropriate conditions. This universality is expressed via the commutativity of:

$$ \begin{gathered} \mathcal{L}(\theta_k) \xrightarrow{\text{gradient step}} \mathcal{L}(\theta_{k+1}) \\ \text{projection} \downarrow \quad \downarrow \text{projection} \\ \|\theta_k - \theta^*\| \xrightarrow{\text{decreasing}} \|\theta_{k+1} - \theta^*\|, \end{gathered} $$

where the diagram commutes as the loss decreases monotonically.

Generalization is formalized through universal properties of Kan extensions, which extend partial models to complete ones. Given a training dataset $D: \mathcal{I} \to \mathbf{Data}$ (a functor from a small index category $\mathcal{I}$), the left Kan extension $\text{Lan}_D F$ of a model $F$ along $D$ computes the "best approximation" of $F$ on unseen data. This universal property is expressed by the natural isomorphism:

$$ \text{Hom}_{\text{Model}}(\text{Lan}_D F, G) \cong \text{Hom}_{\text{Data}^\mathcal{I}}(F, G \circ D), $$

where $G$ is any candidate model. The right-hand side represents training-set performance, while the left-hand side captures generalization, making the Kan extension the optimal balance between the two. In representation learning, the universal property of quotients formalizes abstraction. For a data space $X$ with equivalence relation $\sim$ (e.g., invariance to nuisance factors), the quotient map $q: X \to X/\sim$ satisfies: any $\sim$-invariant map $f: X \to Y$ factors uniquely through $q$ as $f = \tilde{f} \circ q$. In deep learning, this corresponds to the encoder-decoder paradigm:

$$ X \xrightarrow{\text{encoder}} Z \xrightarrow{\text{decoder}} Y, $$

where $Z$ is the latent space of abstract representations, and the encoder realizes the quotient by discarding $\sim$-irrelevant information. The universality ensures $Z$ is the "most efficient" such space, justifying architectures like VAEs. Attention mechanisms are characterized by universal properties in monoidal categories. The attention operation $\text{Attn}: (Q, K, V) \mapsto Y$ is universal
among all bilinear maps from $Q \times K \times V$ that respect the query-key-value structure, expressed by the commuting diagram:

$$ \begin{array}{ccc} Q \otimes K \otimes V & \xrightarrow{\text{Attn}} & Y \\ \text{bilinear} & \searrow & \nearrow \exists! \phi, \end{array} $$

where $\phi$ is any other bilinear operation. This universality explains why attention can approximate arbitrary sequence-to-sequence operations given sufficient heads.

The universal property of neural networks thus lies at the intersection of approximation theory (Stone-Weierstrass), optimization (gradient descent as universal arrow), and representation (quotients and Kan extensions). These properties collectively ensure that well-designed ML systems are both maximally expressive (abstracting patterns) and minimally biased (generalizing beyond training data), with the category-theoretic framework providing the language to prove such claims rigorously.

### 7.2.3 Examples from Familiar Settings

#### 7.2.3.1 Category of Sets (Set)

* **Objects:** Sets.

* **Morphisms:** Functions $f: A \to B$.

* **Functors:**

- The **powerset functor $\mathcal{P}: \mathbf{Set} \to \mathbf{Set}$** maps $A$ to its power set $\mathcal{P}(A)$ and $f: A \to B$ to the direct image $f_*: \mathcal{P}(A) \to \mathcal{P}(B)$.

* **Natural Transformation:**

- The **singleton map $\eta_A: A \to \mathcal{P}(A)$**, where $\eta_A(x) = \{x\}$, is natural because for any $f: A \to B$,

$$ f_*(\{x\}) = \{f(x)\} = \eta_B(f(x)). $$

The category of sets, denoted **Set**, constitutes a foundational object in category theory where objects are all possible sets and morphisms are all possible functions between these sets. Formally, **Set** is defined as follows: for any two objects A and B in **Set**, the hom-set $\text{Hom}_{\text{Set}}(A, B)$ consists of all functions $f: A \to B$. Composition of morphisms is given by ordinary function composition: for $f: A \to B$ and $g: B \to C$, the composite $g \circ f: A \to C$ is defined by $(g \circ f)(a) = g(f(a))$ for all $a \in A$. This composition operation is strictly associative, meaning $h \circ (g \circ f) = (h \circ g) \circ f$ for all composable functions $f, g, h$, and the identity morphism $\text{id}_A: A \to A$ for each set A is the function defined by $\text{id}_A(a) = a$ for all $a \in A$, satisfying $f \circ \text{id}_A = f$ and $\text{id}_B \circ f = f$ for any $f: A \to B$.

The category **Set** is Cartesian closed, meaning it possesses all finite products and exponentials. The product of two sets A and B is their Cartesian product $A \times B$, equipped with projection morphisms $\pi_1: A \times B \to A$ and $\pi_2: A \times B \to B$ defined by $\pi_1(a, b) = a$ and $\pi_2(a, b) = b$. The universal property of the product states that for any set X and morphisms $f: X \to A$ and $g: X \to B$, there exists a unique morphism $\langle f, g \rangle: X \to A \times B$ such that $\pi_1 \circ \langle f, g \rangle = f$ and $\pi_2 \circ \langle f, g \rangle = g$. The exponential object $B^A$ represents the set of all functions from A to B, with evaluation morphism eval: $B^A \times A \to B$ defined by eval($f, a) = f(a)$. This structure ensures that for any morphism $f: C \times A \to B$, there exists a unique morphism $\lambda f: C \to B^A$ such that eval $\circ$ ($\lambda f \times \text{id}_A$) = $f$.

**Set** also admits all limits and colimits, making it complete and cocomplete. The limit of a diagram $D: \mathcal{J} \to \mathbf{Set}$ is given by the set of cones over $D$, while the colimit is the quotient of the
disjoint union of sets in $D$ under an appropriate equivalence relation. For instance, the equalizer of two functions $f, g: A \to B$ is the subset $E = \{a \in A | f(a) = g(a)\}$ with inclusion $e: E \hookrightarrow A$, satisfying $f \circ e = g \circ e$. Dually, the coequalizer is the quotient set $Q = B/\sim$, where $\sim$ is the smallest equivalence relation containing all pairs $(f(a), g(a))$ for $a \in A$, with projection $q: B \to Q$ satisfying $q \circ f = q \circ g$.

The category **Set** serves as a canonical example of a topos, a category with Cartesian closure, finite limits, and a subobject classifier. The subobject classifier in **Set** is the two-element set $\Omega = \{\text{true, false}\}$, equipped with the characteristic function $\chi: A \to \Omega$ for any subset $S \subseteq A$, defined by $\chi(a) = \text{true}$ if $a \in S$ and false otherwise. This structure allows **Set** to internalize logic via its subobject lattice, where logical operations correspond to set-theoretic operations on subsets.

Monomorphisms in **Set** are precisely the injective functions, characterized by the left-cancellation property: $f: A \to B$ is monic if for all functions $g, h: X \to A$, $f \circ g = f \circ h$ implies $g = h$. Epimorphisms are the surjective functions, satisfying the dual right-cancellation property. Isomorphisms are bijections, i.e., functions that are both injective and surjective. The Yoneda embedding $y: \mathbf{Set} \to [\mathbf{Set}^{\text{op}}, \mathbf{Set}]$ maps each set $A$ to its representable functor $\text{Hom}_{\mathbf{Set}}(-, A)$, demonstrating that **Set** is a locally small category where every object is determined by its relationships to other objects via morphisms.

The functor category $[\mathcal{C}, \mathbf{Set}]$ for any small category $\mathcal{C}$ inherits many properties from **Set**, including completeness and cocompleteness, making **Set** a foundational tool for studying presheaves and sheaves in algebraic geometry and logic. The interplay between **Set** and other categories is mediated by adjunctions, such as the free-forgetful adjunction between **Set** and categories of algebraic structures (e.g., groups, rings), where the free functor $F: \mathbf{Set} \to \mathbf{Grp}$ assigns to each set the free group generated by its elements, and the forgetful functor $U: \mathbf{Grp} \to \mathbf{Set}$ returns the underlying set of a group. This adjunction satisfies the natural isomorphism $\text{Hom}_{\mathbf{Grp}}(F(A), G) \cong \text{Hom}_{\mathbf{Set}}(A, U(G))$ for any set $A$ and group $G$.

In summary, **Set** provides a universal framework for mathematical constructions through its rich categorical structure, serving as the archetypal model for concepts like limits, adjunctions, and topos theory. Its role extends beyond foundations, offering concrete intuition for abstract categorical notions while maintaining rigorous formal power.

### 7.2.3.2 Category of Vector Spaces ($\text{Vect}_K$)

* **Objects**: Vector spaces over $K$.

* **Morphisms**: Linear maps $T: V \to W$.

* **Functors**:

- The **dual space functor** $(-)^* : \text{Vect}_K \to \text{Vect}_K$ maps $V$ to $V^* = \text{Hom}(V, K)$ and $T: V \to W$ to its transpose $T^*: W^* \to V^*$.

* **Natural Transformation**:

- The double dual embedding $\eta_V: V \to V^{**}$, where $\eta_V(v)(\phi) = \phi(v)$, is natural because for any linear $T: V \to W$,

$$T^{**}(\eta_V(v)) = \eta_W(T(v)).$$

The category of vector spaces over a field $K$, denoted $\text{Vect}_K$, constitutes a fundamental example of an abelian category in which objects are vector spaces over $K$ and morphisms are $K$-linear transformations between them. Formally, for any two objects $V$ and $W$ in $\text{Vect}_K$, the hom-set
$\text{Hom}_{\mathbf{Vect}_K}(V, W)$ consists of all linear maps $T: V \to W$ satisfying $T(\lambda v + \mu w) = \lambda T(v) + \mu T(w)$ for all $\lambda, \mu \in \mathbb{K}$ and $v, w \in V$. Composition of morphisms is given by the standard composition of linear maps, which preserves linearity: for $T: V \to W$ and $S: W \to U$, the composite $S \circ T: V \to U$ is defined by $(S \circ T)(v) = S(T(v))$. This composition is strictly associative, and the identity morphism $\text{id}_V: V \to V$ is the linear operator $\text{id}_V(v) = v$, satisfying $\text{id}_W \circ T = T = T \circ \text{id}_V$ for any $T: V \to W$.

The category $\mathbf{Vect}_K$ is enriched over itself, meaning that each hom-set $\text{Hom}_{\mathbf{Vect}_K}(V, W)$ naturally carries the structure of a $\mathbb{K}$-vector space. The vector space operations on $\text{Hom}_{\mathbf{Vect}_K}(V, W)$ are defined pointwise: for $T, S \in \text{Hom}_{\mathbf{Vect}_K}(V, W)$ and $\lambda \in \mathbb{K}$, the sum $T + S$ and scalar multiple $\lambda T$ are given by $(T + S)(v) = T(v) + S(v)$ and $(\lambda T)(v) = \lambda T(v)$. This enrichment is compatible with composition, which is bilinear: $(S_1 + S_2) \circ T = S_1 \circ T + S_2 \circ T$ and $S \circ (T_1 + T_2) = S \circ T_1 + S \circ T_2$ for all appropriately composable linear maps. The zero morphism $0: V \to W$ is the linear map sending every $v \in V$ to the zero vector in $W$.

$\mathbf{Vect}_K$ is a monoidal category with respect to the tensor product $\otimes_K$. The tensor product $V \otimes_K W$ of two vector spaces is characterized by the universal property that any bilinear map $f: V \times W \to U$ factors uniquely through a linear map $\tilde{f}: V \otimes_K W \to U$. This gives the adjunction $\text{Hom}_{\mathbf{Vect}_K}(V \otimes_K W, U) \cong \text{Bilin}(V \times W, U)$, where $\text{Bilin}(V \times W, U)$ denotes the set of bilinear maps. The monoidal unit is the field $\mathbb{K}$ itself, viewed as a one-dimensional vector space over itself, with natural isomorphisms $V \otimes_K \mathbb{K} \cong V \cong \mathbb{K} \otimes_K V$ for any $V \in \mathbf{Vect}_K$.

Finite products and coproducts coincide in $\mathbf{Vect}_K$, both being given by the direct sum $V \oplus W$. The projections $\pi_V: V \oplus W \to V$ and $\pi_W: V \oplus W \to W$ are linear maps defined by $\pi_V(v, w) = v$ and $\pi_W(v, w) = w$, while the inclusions $\iota_V: V \to V \oplus W$ and $\iota_W: W \to V \oplus W$ are given by $\iota_V(v) = (v, 0)$ and $\iota_W(w) = (0, w)$. The universal properties of product and coproduct are satisfied: for any linear maps $f: U \to V$ and $g: U \to W$, there exists a unique linear map $\langle f, g \rangle: U \to V \oplus W$ such that $\pi_V \circ \langle f, g \rangle = f$ and $\pi_W \circ \langle f, g \rangle = g$; dually, for any $f: V \to U$ and $g: W \to U$, there exists a unique linear map $[f, g]: V \oplus W \to U$ such that $[f, g] \circ \iota_V = f$ and $[f, g] \circ \iota_W = g$.

The category $\mathbf{Vect}_K$ is abelian, meaning it has all kernels and cokernels, every monomorphism is a kernel, and every epimorphism is a cokernel. The kernel of a linear map $T: V \to W$ is the subspace $\ker T = \{v \in V \mid T(v) = 0\}$ equipped with the inclusion $\iota: \ker T \hookrightarrow V$, while the cokernel is the quotient space $\text{coker}T = W/\text{im}T$ with the projection $\pi: W \to W/\text{im}T$. A sequence $V \xrightarrow{T} W \xrightarrow{S} U$ is exact at $W$ if and only if $\text{im}T = \ker S$, and short exact sequences $0 \to V \xrightarrow{T} W \xrightarrow{S} U \to 0$ correspond to situations where $T$ is injective, $S$ is surjective, and $\text{im}T = \ker S$.

The duality between vector spaces and their duals is captured by the contravariant functor $(-)^* = \text{Hom}_{\mathbf{Vect}_K}(-, \mathbb{K})$: $\mathbf{Vect}_K^{\text{op}} \to \mathbf{Vect}_K$. For finite-dimensional vector spaces, this duality is perfect: the canonical map $V \to V^{**}$ given by $v \mapsto (\phi \mapsto \phi(v))$ is a natural isomorphism. This reflects the fact that $\mathbf{Vect}_K$ is a compact closed category, where the dual space $V^*$ serves as the left and right dual of $V$ with evaluation and coevaluation maps $\text{ev}_V: V^* \otimes_K V \to \mathbb{K}$ and $\text{coev}_V: \mathbb{K} \to V \otimes_K V^*$ defined by $\text{ev}_V(\phi \otimes v) = \phi(v)$ and $\text{coev}_V(1) = \sum_{i=1}^n v_i \otimes v_i^*$ for a basis $\{v_i\}$ of $V$ with dual basis $\{v_i^*\}$ of $V^*$.

The functor category $[\mathcal{C}, \mathbf{Vect}_K]$ for any small category $\mathcal{C}$ inherits the abelian structure from $\mathbf{Vect}_K$, with all limits and colimits computed pointwise. This makes $\mathbf{Vect}_K$ a central object in representation theory, where representations of groups, algebras, or other structures are studied as functors into $\mathbf{Vect}_K$. The tensor-hom adjunction $\text{Hom}_{\mathbf{Vect}_K}(V \otimes_K W, U) \cong \text{Hom}_{\mathbf{Vect}_K}(V, \text{Hom}_{\mathbf{Vect}_K}(W, U))$ further enriches the categorical structure, making $\mathbf{Vect}_K$ a closed monoidal category.
In summary, $Vect_K$ provides a rich framework for linear algebra through its categorical properties, serving as a paradigmatic example of an abelian, monoidal, and compact closed category. Its structure facilitates the study of linear transformations, duality, and higher-dimensional algebraic constructs while maintaining deep connections to other branches of mathematics through representation theory and homological algebra.

### 7.2.3.3 Category of Neural Networks (Informal Example)

* **Objects**: Data spaces (e.g., $\mathbb{R}^n$).

* **Morphisms**: Neural network layers (e.g., affine maps $f(x) = Wx + b$).

* **Functors**:

- A **training functor** could map a network architecture to its trained parameter space.

The category of neural networks can be informally conceptualized as a categorical framework where objects represent data spaces (typically Euclidean spaces $\mathbb{R}^n$) and morphisms correspond to parameterized neural network layers. For a fixed activation function $\sigma: \mathbb{R} \to \mathbb{R}$, each morphism $f_\theta: \mathbb{R}^n \to \mathbb{R}^m$ is defined by a parameter set $\theta = (W, b)$ where $W \in \mathbb{R}^{m \times n}$ is a weight matrix and $b \in \mathbb{R}^m$ is a bias vector, with the layer operation given by $f_\theta(x) = \sigma(Wx + b)$. Here, $\sigma$ is applied component-wise, and the parameter space $\Theta$ for each morphism is the manifold $\mathbb{R}^{m \times n} \times \mathbb{R}^m$ of possible weights and biases. Composition of morphisms $f_{\theta_1}: \mathbb{R}^n \to \mathbb{R}^k$ and $f_{\theta_2}: \mathbb{R}^k \to \mathbb{R}^m$ yields a new morphism $f_{\theta_2 \circ f_{\theta_1}}: \mathbb{R}^n \to \mathbb{R}^m$ with parameters $\theta = (\theta_1, \theta_2)$, where $(f_{\theta_2 \circ f_{\theta_1}})(x) = \sigma(W_2(\sigma(W_1x + b_1)) + b_2)$. This composition is associative up to the non-linearities introduced by $\sigma$, and the identity morphism $\text{id}_{\mathbb{R}^n}: \mathbb{R}^n \to \mathbb{R}^n$ is theoretically given by $\text{id}_{\mathbb{R}^n}(x) = x$, though exact implementation requires $\sigma$ to be the identity function and $(W, b) = (I_n, 0)$, which is often violated in practice due to architectural constraints.

The categorical structure becomes non-trivial when considering parallel compositions via monoidal products. Given two layers $f_{\theta_1}: \mathbb{R}^n \to \mathbb{R}^m$ and $f_{\theta_2}: \mathbb{R}^k \to \mathbb{R}^l$, their tensor product $f_{\theta_1} \otimes f_{\theta_2}: \mathbb{R}^{n+k} \to \mathbb{R}^{m+l}$ is defined by block-diagonal weight matrices and concatenated biases: $(f_{\theta_1} \otimes f_{\theta_2})(x_1, x_2) = (\sigma(W_1x_1 + b_1), \sigma(W_2x_2 + b_2))$. This models architectural features like multi-head attention or residual connections, where parallel processing occurs. The interchange law $(f_{\theta_2} \otimes g_{\phi_2}) \circ (f_{\theta_1} \otimes g_{\phi_1}) = (f_{\theta_2} \circ f_{\theta_1}) \otimes (g_{\phi_2} \circ g_{\phi_1})$ holds only approximately due to activation function non-linearities, reflecting the empirical behavior of neural networks where parallel and sequential compositions interact non-ideally.

Universal properties in this category are relaxed to approximate versions. For instance, the "universal approximation" property of neural networks corresponds to the existence of a morphism $f_\theta: \mathbb{R}^n \to \mathbb{R}^m$ that $\epsilon$-approximates any continuous function $g: \mathbb{R}^n \to \mathbb{R}^m$ on compact subsets, where $\epsilon > 0$ is an error tolerance. This is expressed as $\|f_\theta(x) - g(x)\| < \epsilon$ for all $x \in K \subset \mathbb{R}^n$, with $K$ compact. The categorical product is replaced by a weakened notion where the diagram commutes only up to $\epsilon$, reflecting the optimization-driven nature of neural networks. For a network $N$ composed of layers $f_{\theta_1}, \dots, f_{\theta_k}$, the actual computation $N(x)$ approximates the idealized composition $f_{\theta_k} \circ \dots \circ f_{\theta_1}(x)$ due to numerical errors and training imperfections.

This informal category differs from classical mathematical categories in its tolerance for approximation errors and its reliance on optimization to instantiate morphisms. The "functoriality" of neural architectures is observed empirically: a computational graph $\mathcal{G}$ maps to a network $N_\mathcal{G}$ in the category, but the functor laws $N_{G_1 \circ G_2} \approx N_{G_1} \circ N_{G_2}$ and $N_{\text{id}} \approx \text{id}$ hold only approximately after training. The category thus captures the pragmatic essence of neural networks as composable, parameterized functions whose categorical structure emerges statistically through gradient-based optimization rather than through exact algebraic satisfaction of axioms.
### 7.2.4 Diagrams as Proofs and Reasoning Tools

In category theory, diagrams serve as rigorous proof tools through the formal machinery of commutative diagrams, where the equality of morphisms along different paths is guaranteed by the underlying categorical structure. A diagram commutes when for any two objects $X$ and $Y$ and all paths of morphisms between them, the composition of morphisms along each path yields identical results. For example, given morphisms $f: A \to B$, $g: B \to C$, and $h: A \to C$, the triangular diagram:

<figure><img src="image_72.png" /></figure>

commutes if and only if $g \circ f = h$. This condition imposes an algebraic constraint on the morphisms, ensuring coherence in the categorical structure. Diagrams generalize equational reasoning to higher-dimensional compositions. The pasting lemma formalizes how smaller commutative diagrams can be glued along common morphisms to form larger commutative diagrams. Given two commuting squares sharing a morphism $f: X \to Y$:

<figure><img src="image_73.png" /></figure>

the commutativity of the left square ($f \circ p = q \circ u$) and right square ($g \circ q = r \circ v$) implies the commutativity of the outer rectangle (($g \circ f) \circ p = r \circ (v \circ u)$). This compositional property mirrors the associativity of morphisms and extends to diagrams of arbitrary complexity. Universal constructions, such as limits and colimits, are defined via universal properties expressed through commutative diagrams. A product $A \times B$ of objects $A$ and $B$ is characterized by a universal cone:

<figure><img src="image_74.png" /></figure>

such that for any object $X$ with morphisms $f: X \to A$ and $g: X \to B$, there exists a **unique** morphism $h: X \to A \times B$ making the diagram commute:

$$ \pi_A \circ h = f \quad \text{and} \quad \pi_B \circ h = g. $$

This universal property ensures that the product is uniquely determined up to isomorphism, and similar diagrams define coproducts, equalizers, and other categorical limits.

In categorical logic, string diagrams provide a graphical calculus for reasoning in monoidal categories, where morphisms are represented as boxes and composition as wiring. The interchange law:

$$ (g \circ f) \otimes (k \circ h) = (g \otimes k) \circ (f \otimes h) $$

is visually evident in string diagrams, where vertical stacking represents composition and horizontal juxtaposition represents the tensor product. This graphical language rigorously captures the algebraic structure while simplifying proofs of coherence conditions in symmetric monoidal categories. Diagrams also formalize adjunctions and Kan extensions through bijective correspondences between hom-sets. An adjunction $F \dashv G$ between functors $F: \mathcal{C} \to \mathcal{D}$ and $G: \mathcal{D} \to \mathcal{C}$ is encoded by the natural isomorphism:

$$ \mathcal{D}(F(X), Y) \cong \mathcal{C}(X, G(Y)) $$
which can be represented as a family of commuting diagrams expressing the unit $\eta$: $\text{Id} \to GF$ and counit $\epsilon$: $FG \to \text{Id}$. These diagrams enforce the triangle identities:

$$ \epsilon F \circ F \eta = \text{Id}_F \quad \text{and} \quad G \epsilon \circ \eta G = \text{Id}_G, $$

ensuring the coherence of the adjunction.

Thus, diagrams in category theory are not merely illustrative but constitute formal proofs through their commutativity conditions, universal properties, and graphical equivalences. They provide a synthetic language for expressing complex algebraic relationships, reducing intricate equational reasoning to diagrammatic manipulations that preserve categorical structure.

### 7.2.4.1 Commutative Diagrams

A diagram **commutes** if all paths between two objects yield the same morphism. For example, the naturality square for $\eta: F \Rightarrow G$:

$$ \begin{tikzcd} F(X) \arrow[r, "\eta_X"] & G(X) \\ F(f) \arrow[u, "\eta_Y"] \arrow[d] & G(f) \arrow[u, "\eta_Y"] \\ F(Y) \arrow[r, "\eta_Y"] & G(Y) \end{tikzcd} $$

commutes by definition of $\eta$.

Commutative diagrams constitute the visual language of category theory, providing a rigorous graphical calculus for expressing algebraic relationships between morphisms. Formally, a commutative diagram in a category $\mathcal{C}$ is a directed graph whose vertices represent objects and whose edges represent morphisms, where the fundamental property asserts that all paths between any two vertices yield identical morphisms through composition. This means that for any pair of objects $X, Y$ in the diagram, given two parallel paths $f_n \circ \dots \circ f_1$ and $g_m \circ \dots \circ g_1$ from $X$ to $Y$, the equality $f_n \circ \dots \circ f_1 = g_m \circ \dots \circ g_1$ must hold in $\mathcal{C}$. The canonical example appears in the definition of natural transformations, where for $\eta: F \Rightarrow G$ and $f: X \to Y$, the square:

$$ \begin{tikzcd} F(X) \arrow[r, "\eta_X"] & G(X) \\ F(f) \arrow[u, "\eta_Y"] \arrow[d] & G(f) \arrow[u, "\eta_Y"] \\ F(Y) \arrow[r, "\eta_Y"] & G(Y) \end{tikzcd} $$

commutes by naturality, expressing the equation $G(f) \circ \eta_X = \eta_Y \circ F(f)$. This single diagram encapsulates the coherence condition that each component $\eta_X$ must satisfy relative to every morphism in the category. The power of this notation lies in its ability to compress complex algebraic identities into topological configurations of arrows, where the commutativity constraint ensures the consistency of the categorical structure.

The universality of commutative diagrams emerges through their capacity to characterize limits and colimits. A cone over a diagram $D: \mathcal{J} \to \mathcal{C}$ consists of an object $L$ with morphisms $\pi_X: L \to D(X)$ for each $X \in \mathcal{J}$, making all induced triangles commute. The limit object $\lim D$ satisfies the universal property that for any other cone $(L', \{\pi'_X\})$, there exists a unique morphism $u: L' \to \lim D$ rendering commutative all diagrams:

$$ \begin{tikzcd} L' \arrow[r, "u"] & \lim D \\ \pi'_X \arrow[u, "\pi_X"] \arrow[d] & D(X) \arrow[u, "\pi_X"] \\ D(X) & D(X) \end{tikzcd} $$
for every $X \in \mathcal{J}$. Dually, colimits reverse all arrows while preserving commutativity. These diagrammatic conditions distill infinite verbal descriptions into finite graphical data, where the universal morphism $u$'s uniqueness follows from the commutativity constraints on all possible paths.

In abelian categories, commutative diagrams acquire enhanced potency through exact sequences and homology. The snake lemma's proof revolves around the meticulous construction of a commutative diagram connecting kernels to cokernels:

<figure><img src="image_75.png" /></figure>

where exactness and commutativity jointly induce the connecting homomorphism $\delta : \ker h \to \operatorname{coker} f$. Each commutative subdiagram encodes specific algebraic relations, while their collective commutativity ensures the lemma's validity. The diagram serves not merely as illustration but as the actual mathematical object of study, where the geometric arrangement of arrows carries equational content.

Higher-dimensional category theory extends commutative diagrams to higher commutativity through pasting diagrams in 2-categories. For 2-morphisms $\alpha : f \Rightarrow g$ and $\beta : h \Rightarrow k$, their vertical and horizontal compositions must satisfy the interchange law, visualized through the commutativity of:

<figure><img src="image_76.png" /></figure>

where the equality of pastings expresses the coherence between different compositions of 2-cells. This elevates commutative diagrams from planar graphs to higher-dimensional structures, where commutativity becomes homotopy coherence in $(\infty, 1)$-categories. The strict commutativity of 1-categorical diagrams generalizes to coherent homotopy commutativity in higher categories, mediated by systems of higher morphisms.

### 7.2.4.2 Universal Properties via Diagrams

A **product** $X \times Y$ is defined by a universal property: for any object $Z$ with morphisms $f : Z \to X$ and $g : Z \to Y$, there exists a unique $h : Z \to X \times Y$ making the following commute:

<figure><img src="image_77.png" /></figure>

where $\pi_X, \pi_Y$ are projections.

Universal properties in category theory are most powerfully expressed through commutative diagrams that capture the essential uniqueness and existence conditions characterizing mathematical constructions. Formally, a universal property describes an object and morphisms that are either
initial or terminal with respect to some diagrammatic condition, where this extremality is expressed through the existence of a unique factorization. The product $X \times Y$ of two objects $X, Y \in \mathcal{C}$, for instance, is defined by a universal cone consisting of the object $X \times Y$ equipped with projection morphisms $\pi_X : X \times Y \to X$ and $\pi_Y : X \times Y \to Y$, such that for any other object $Z$ with morphisms $f : Z \to X$ and $g : Z \to Y$, there exists a unique morphism $u : Z \to X \times Y$ making the following diagram commute:

<figure><img src="image_78.png" /></figure>

This universal property simultaneously asserts existence through the morphism $u$ and uniqueness through its distinctiveness in satisfying the commutative conditions. The product's universal property thus reduces the infinite collection of possible cones to a single optimal solution, where the commutativity constraints enforce the precise alignment of structural relationships.

Colimits dualize this construction by reversing all arrows while preserving the universal factorization property. The coproduct $X \sqcup Y$ is characterized by injection morphisms $\iota_X : X \to X \sqcup Y$ and $\iota_Y : Y \to X \sqcup Y$ such that for any object $Z$ with morphisms $f : X \to Z$ and $g : Y \to Z$, there exists a unique $u : X \sqcup Y \to Z$ rendering commutative:

<figure><img src="image_79.png" /></figure>

The diagrammatic formulation encodes the coproduct's universal mapping property as a terminal object in the category of cocones, where the existence of $u$ mediates between the individual morphisms $f, g$ and their joint extension through the coproduct.

More sophisticated universal constructions like equalizers and pullbacks exhibit richer diagrammatic conditions. The equalizer $\mathrm{eq}(f, g)$ of parallel morphisms $f, g : X \to Y$ is an object $E$ with morphism $e : E \to X$ satisfying $f \circ e = g \circ e$, universal for this property in that any $h : Z \to X$ with $f \circ h = g \circ h$ factors uniquely through $e$:

<figure><img src="image_80.png" /></figure>

The pullback $X \times_Z Y$ of morphisms $f : X \to Z$ and $g : Y \to Z$ extends this to a two-dimensional universal property, where the pullback square commutes and is universal among all such commuting squares:

<figure><img src="image_81.png" /></figure>

Every pair of morphisms $q_X : Q \to X$ and $q_Y : Q \to Y$ with $f \circ q_X = g \circ q_Y$ induces a unique $u : Q \to P$ making all resulting triangles commute. These diagrammatic universal properties collectively form the blueprint for categorical constructions, where the diagrams themselves become the definitions rather than mere illustrations. The universality is precisely the requirement that certain diagrams commute and that solutions to specified diagrammatic conditions exist and are unique up to isomorphism. This diagrammatic approach provides a unified language for universal properties across all branches of mathematics, from algebra to topology to logic, where the same categorical patterns recur in different concrete instantiations.
### 7.2.4.3 Applications in ML

* **Backpropagation as a Functorial Diagram:** The chain rule can be expressed as a commuting diagram in the category of differentiable functions.

* **Attention Mechanisms:** The query-key-value interaction in transformers can be modeled as a limit diagram.

#### 7.2.4.3.1 Backpropagation as a Functorial Diagram

Backpropagation in neural networks admits a rigorous categorical formulation as a functor between appropriately defined categories of differentiable functions. Consider the category $\mathcal{D}\rangle\{\}$ whose objects are Euclidean spaces $\mathbb{R}^n$ and whose morphisms are differentiable functions $f : \mathbb{R}^n \to \mathbb{R}^m$, and the category $\mathcal{P}\vdash\nabla\vdash\nabla$ whose objects are parameter spaces $\Theta$ and whose morphisms are parameter updates. The backpropagation functor $\mathcal{B} : \mathcal{D}\rangle\{\to \mathcal{P}\vdash\nabla\vdash\nabla\}$ acts on morphisms by sending each differentiable function $f_\theta : \mathbb{R}^n \to \mathbb{R}^m$ (where $\theta \in \Theta$ denotes the parameters) to its parameter gradient computation:

$$ \begin{array}{ccc} (\mathbb{R}^n, f_\theta) & \xrightarrow{\mathcal{B}} & (\Theta, \nabla_\theta \mathcal{L}) \\ \big\downarrow\big\downarrow & & \big\downarrow\big\downarrow_{(\mathcal{B}(g))} \\ (\mathbb{R}^k, f'_{\theta'}) & \xrightarrow{\mathcal{B}} & (\Theta', \nabla'_{\theta'} \mathcal{L}') \end{array} $$

where the vertical morphisms represent changes in architecture and $\mathcal{L}$ is the loss function. The functoriality of $\mathcal{B}$ manifests in the chain rule, which ensures that for composable morphisms $f$ and $g$, the gradient computation satisfies:

$$ \mathcal{B}(g \circ f) = \mathcal{B}(f) \circ \mathcal{B}(g) $$

This equality expresses the fundamental property that backpropagation through a composition of layers corresponds to the composition of individual gradient computations. The diagram commutes precisely when the backward pass correctly accumulates gradients through the computational graph, with the functor $\mathcal{B}$ preserving the compositional structure of the network.

The universal property of backpropagation emerges through its formulation as a lens in the category of reverse derivatives. Each parameterized function $f_\theta$ induces a lens $(f_\theta, \nabla_\theta f_\theta)$ pairing the forward pass with its parameter gradient. The backpropagation functor $\mathcal{B}$ then becomes the realization of this universal construction, where the gradient computation is uniquely determined by the requirement to maintain coherence with the chain rule. This perspective reveals backpropagation as the canonical solution to the problem of gradient computation in compositional systems, with the functorial diagram:

$$ \begin{array}{ccc} \mathbb{R}^n & \xrightarrow{f_\theta} & \mathbb{R}^m \\ \mathcal{B}(f_\theta) \downarrow & & \mathcal{B}(id) \downarrow \\ \Theta \times \mathbb{R}^m & \xrightarrow{\pi_1} & \mathbb{R}^m \end{array} $$

capturing the essential relationship between forward computation and backward gradient flow. The commutativity of this diagram enforces the constraint that parameter updates must properly account for downstream gradient contributions, making explicit the information flow required for correct learning dynamics in neural networks.

#### 7.2.4.3.2 Attention Mechanisms
The query-key-value interaction in transformer architectures admits a rigorous categorical formulation as a limit construction in an appropriate category of attention diagrams. Consider the category $\mathcal{A} \sqcup \sqcup$ whose objects are triples $(Q, K, V)$ representing query, key, and value vector spaces, and whose morphisms are linear transformations between these spaces that preserve the attention structure. The attention mechanism defines a diagram $D : \mathcal{J} \to \mathcal{A} \sqcup \sqcup$ where $\mathcal{J}$ is an indexing category capturing the pairwise interactions between queries and keys, with objects $(q_i, k_j)$ for all $i, j$ in the sequence positions and morphisms representing their compatibility relations. The limit of this diagram $\lim D$ constitutes the universal object equipped with projection morphisms $\pi_{ij} : \lim D \to \text{Hom}(Q_i, K_j)$ making all subdiagrams commute:

$$ \begin{array}{ccc} \lim D & \xrightarrow{\pi_{ij}} & \text{Hom}(Q_i, K_j) \\ \pi_{i'j'} \downarrow & & \downarrow \alpha_{ij}^{i'j'} \\ \text{Hom}(Q_{i'}, K_{j'}) & \xrightarrow{\beta_{ij}^{i'j'}} & \mathcal{S} \end{array} $$

where $\mathcal{S}$ is the space of attention scores and $\alpha_{ij}^{i'j'}$, $\beta_{ij}^{i'j'}$ are the natural transformations encoding the softmax normalization and scaling by $\sqrt{d_k}$. The universal property ensures that for any other object $X$ with morphisms $f_{ij} : X \to \text{Hom}(Q_i, K_j)$ respecting these compatibility relations, there exists a unique morphism $u : X \to \lim D$ factoring all $f_{ij}$ through the projections $\pi_{ij}$.

The value aggregation emerges through the cone condition over the limit, where the attention weights $\text{Att}(Q, K) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})$ mediate between the limit projections and the value space $V$. This constructs a natural transformation $\eta : \Delta V \to D$ from the constant functor $\Delta V$ to the diagram $D$, whose component at each position $(i, j)$ is the linear map $V_j \to \bigoplus_i V_i$ weighted by $\text{Att}(Q_i, K_j)$. The multi-head attention mechanism appears as a product of such limits, where the universal property guarantees the independence of attention heads while maintaining their joint interaction through the value space:

$$ \begin{array}{ccc} \prod_{h=1}^{H} \lim D_h & \xrightarrow{\otimes} & \text{End}(V) \\ \otimes_h \downarrow & & \downarrow \text{merge} \\ \lim D_h & \xrightarrow{\text{Att}_h} & \text{Hom}(V, V) \end{array} $$

Here merge represents the linear combination of head outputs, and the commutativity of the diagram enforces the consistency between individual head computations and their aggregated result. The residual connection and layer normalization further arise as universal constructions, where the skip connection forms a coproduct diagram and the normalization implements a coequalizer of scale transformations. This limit diagram perspective reveals the transformer's attention mechanism as fundamentally a universal solution to the problem of context-aware sequence processing, where the architectural constraints emerge naturally from the categorical universal properties governing the interaction of queries, keys, and values.

### 7.2.5 Summary

Category Theory provides a **rigorous language** for ML by:

1. **Abstracting** computation via categories and functors.

2. **Unifying** structures (e.g., vector spaces, neural networks) under common principles.

3. **Enabling diagrammatic proofs** for complex architectures.

This formalism is not just theoretical—it guides the design of scalable, interpretable models. For instance, **monoidal categories** formalize tensor operations in deep learning, while **adjunctions** describe optimality conditions in autoencoders.
# 7.3 Neural Networks as Composable Morphisms

## 7.3.1 Layers as Morphisms; Networks as Compositions

The formalization of neural network layers as morphisms in a categorical framework provides a rigorous mathematical foundation for understanding their compositional nature. Within the category **Para**, each neural network layer constitutes a morphism $f_i: X_i \times \Theta_i \to X_{i+1}$ where $X_i$ and $X_{i+1}$ represent the input and output vector spaces respectively, and $\Theta_i$ denotes the parameter space of weights and biases. These morphisms are required to be smooth maps, ensuring differentiability for gradient-based optimization. The parameterized nature of these morphisms is captured through the Cartesian product with the parameter space, making explicit the functional dependence on both input data and learned parameters.

The composition of layers in a feedforward neural network corresponds precisely to the categorical composition operation in **Para**. For a sequence of layers $(f_1, \theta_1): X_0 \to X_1$ through $(f_n, \theta_n): X_{n-1} \to X_n$, the entire network implements the morphism:

$$ (f_n \circ \dots \circ f_1, (\theta_1, \dots, \theta_n)): X_0 \to X_n $$

with the composite function defined recursively through the composition law:

$$ (f_{i+1} \circ f_i)(x, (\theta_i, \theta_{i+1})) = f_{i+1}(f_i(x, \theta_i), \theta_{i+1}) $$

This composition operation satisfies the associativity axiom strictly, meaning that for any three consecutive layers, the equality $(f_3 \circ (f_2 \circ f_1)) = ((f_3 \circ f_2) \circ f_1)$ holds exactly at the level of parameterized functions. The identity morphism $\mathrm{id}_X: X \to X$, defined by $\mathrm{id}_X(x, \emptyset) = x$ where $\emptyset$ represents the empty parameter set, serves as the neutral element for composition and corresponds to a trivial pass-through layer in neural network terms.

The monoidal structure of **Para** captures parallel composition of layers through the tensor product operation. Given two layers $(f, \theta_f): X \to Y$ and $(g, \theta_g): A \to B$, their parallel composition is the morphism:

$$ (f \otimes g, (\theta_f, \theta_g)): X \times A \to Y \times B $$

defined by $(f \otimes g)((x, a), (\theta_f, \theta_g)) = (f(x, \theta_f), g(a, \theta_g))$. This operation models architectural features like multi-head attention mechanisms where multiple layer computations occur in parallel. The interchange law governing the relationship between sequential and parallel composition ensures the coherence of these operations:

$$ (f_2 \otimes g_2) \circ (f_1 \otimes g_1) = (f_2 \circ f_1) \otimes (g_2 \circ g_1) $$

when the domains and codomains are properly aligned.

The categorical perspective reveals that common neural network operations correspond to standard categorical constructions. Batch normalization emerges as a natural transformation between parallel layers, while dropout can be modeled as a probabilistic morphism in a suitably enriched category. The universal property of certain colimits explains why skip connections in residual networks facilitate gradient flow - they provide categorical coproducts that preserve information through the composition chain. This abstract formulation demonstrates that successful neural network architectures often implicitly implement fundamental categorical concepts, providing a mathematical explanation for their empirical effectiveness.
### 7.3.1.1 Neural Networks as Composable Functions

A neural network can be viewed as a composition of parameterized functions (layers). Let:

* Each **layer** $L_i$ be a function $f_i: X_i \times \Theta_i \to X_{i+1}$, where:

- $X_i$ is the input space (e.g., $\mathbb{R}^{n_i}$),

- $\Theta_i$ is the parameter space (e.g., weights and biases $\mathbb{R}^{p_i}$),

- $X_{i+1}$ is the output space.

* A **network** $N$ with $k$ layers is the composition:

$$N(x) = (f_k \circ f_{k-1} \circ \dots \circ f_1)(x; \theta_1, \dots, \theta_k)$$

where $\theta_i \in \Theta_i$.

Neural networks can be rigorously formalized as composable functions within a categorical framework, where each layer is treated as a morphism and the entire network is the composition of these morphisms. Let **Para** denote the category of parameterized functions, where objects are Euclidean spaces $\mathbb{R}^n$ and morphisms are smooth functions equipped with parameters. Specifically, a morphism $(f, \theta): \mathbb{R}^n \to \mathbb{R}^m$ in **Para** is defined by a function $f: \mathbb{R}^n \times \Theta \to \mathbb{R}^m$, where $\Theta$ is the parameter space associated with $f$. The parameters $\theta \in \Theta$ modulate the behavior of the function, analogous to the weights and biases in a neural network layer.

The composition of morphisms in **Para** captures the sequential application of layers in a neural network. Given two morphisms $(f, \theta): \mathbb{R}^n \to \mathbb{R}^m$ and $(g, \phi): \mathbb{R}^m \to \mathbb{R}^p$, their composite $(g \circ f, (\theta, \phi)) : \mathbb{R}^n \to \mathbb{R}^p$ is defined by the function:

$$(g \circ f)(x, (\theta, \phi)) = g(f(x, \theta), \phi).$$

This operation is associative, meaning that for three morphisms $(f, \theta)$, $(g, \phi)$, and $(h, \psi)$, the equality:

$$(h \circ (g \circ f), ((\theta, \phi), \psi)) = ((h \circ g) \circ f, (\theta, (\phi, \psi)))$$

holds, ensuring that the order of layer composition is unambiguous. The identity morphism $\mathrm{id}_{\mathbb{R}^n}: \mathbb{R}^n \to \mathbb{R}^n$ is given by the function $\mathrm{id}_{\mathbb{R}^n}(x, ( )) = x$, which acts as a trivial layer preserving the input.

The categorical perspective generalizes to parallel compositions via monoidal structure. Given two morphisms $(f, \theta): \mathbb{R}^n \to \mathbb{R}^m$ and $(g, \phi): \mathbb{R}^k \to \mathbb{R}^l$, their tensor product $(f \otimes g, (\theta, \phi)) : \mathbb{R}^{n+k} \to \mathbb{R}^{m+l}$ is defined by:

$$(f \otimes g)((x, y), (\theta, \phi)) = (f(x, \theta), g(y, \phi)).$$

This allows for the modeling of parallel layer computations, such as those in multi-head attention or residual connections. The universality of neural networks arises from their ability to approximate arbitrary continuous functions through composition. The associativity and identity laws ensure that networks can be constructed modularly, while the monoidal structure accommodates branching architectures. This formalism extends naturally to recurrent networks by considering dynamical systems as morphisms in a suitable category, further demonstrating the flexibility of the compositional approach.

In summary, neural networks as composable functions are characterized by their representation as morphisms in **Para**, where composition corresponds to layer chaining, associativity guarantees structural coherence, and identity morphisms enable skip connections. The framework provides a rigorous foundation for analyzing and constructing neural architectures while preserving their functional and algebraic properties.
### 7.3.1.2 Category-Theoretic Interpretation

We define a category **Para** (for “parameterized functions”) where:

* **Objects** are spaces $X_i$ (e.g., $\mathbb{R}^{n_i}$),

* **Morphisms** $X \to Y$ are pairs $(f, \theta)$, where:

$$f: X \times \Theta \to Y$$

is a smooth function, and $\theta \in \Theta$.

* **Composition** of morphisms $(f, \theta) : X \to Y$ and $(g, \phi) : Y \to Z$ is given by:

$$(g \circ f, (\theta, \phi)) : X \to Z$$

where:

$$(g \circ f)(x, (\theta, \phi)) = g(f(x, \theta), \phi).$$

* **Identity morphism** $\text{id}_X : X \to X$ is given by $\text{id}_X(x, (\)) = x$ (no parameters).

Under this formalism, a neural network is a composition of morphisms in **Para**.

The category-theoretic interpretation of neural networks provides a rigorous mathematical framework for understanding their compositional structure through the lens of abstract algebraic concepts. At its core, this interpretation formalizes neural networks as morphisms in a specially constructed category where objects represent data spaces and arrows represent parameterized transformations. Let **Para** denote the category of parameterized functions, whose objects are smooth manifolds (typically Euclidean spaces $\mathbb{R}^n$) representing possible input and output spaces of network layers. Each morphism $f: X \to Y$ in this category is realized as a smooth function $f: X \times \Theta \to Y$ where $\Theta$ is a manifold of parameters, making explicit the dependence on learnable weights and biases.

The fundamental operation of neural network composition corresponds to the categorical composition of morphisms. Given two parameterized functions $(f, \theta) : X \to Y$ and $(g, \phi) : Y \to Z$, their composite $(g \circ f, (\theta, \phi)) : X \to Z$ is defined through the diagrammatic composition:

$$X \times \Theta \times \Phi \to Y \times \Phi \to Z \\ (x, \theta, \phi) \mapsto (f(x, \theta), \phi) \mapsto g(f(x, \theta), \phi)$$

This composition operation satisfies the strict associativity axiom: for any triple of composable morphisms $(f, \theta)$, $(g, \phi)$, and $(h, \psi)$, the equation

$$h \circ (g \circ f) = (h \circ g) \circ f$$

holds exactly, not just up to approximation. The identity morphism $\text{id}_X : X \to X$ is given by the canonical projection $\pi_X : X \times \{\ast\} \to X$ where $\{\ast\}$ is the singleton parameter space, satisfying the identity laws $\text{id}_Y \circ f = f = f \circ \text{id}_X$ for any morphism $f: X \to Y$.

The categorical perspective reveals deeper structure through universal properties and adjunctions. The parameter space construction forms a monoidal structure where the tensor product of objects $X \otimes Y = X \times Y$ represents parallel composition of data spaces, and the tensor product of morphisms $(f \otimes g)((x, y), (\theta, \phi)) = (f(x, \theta), g(y, \phi))$ models parallel execution of network components. This gives rise to a symmetric monoidal category structure that captures both sequential and parallel composition of neural networks. Factoriality emerges when considering network architectures as functors from a computational graph category to **Para**. Each node in the computational graph
maps to an object in **Para**, and each edge to a parameterized morphism. The functor preserves composition, guaranteeing that the semantic interpretation of the network respects the graph's structure. This perspective enables rigorous treatment of network architectures as mathematical objects in their own right, rather than just particular instantiations with concrete parameters.

The categorical framework naturally accommodates advanced architectural features through standard constructions. Skip connections arise as coproducts, attention mechanisms as enriched categorical operations, and recurrent networks as fixpoints of endofunctors. The universal property of certain categorical limits explains why residual networks can overcome vanishing gradients - the identity morphism provides a canonical pathway for gradient flow that is preserved by the categorical structure. This abstract viewpoint reveals that many successful neural network design patterns correspond to fundamental categorical constructions.

### 7.3.2 The Category of Parameterized Functions

The category of parameterized functions, denoted **Para**, is a foundational construct in categorical machine learning that rigorously formalizes the structure of parameter-dependent transformations central to neural networks and other learned systems. Objects in **Para** are measurable spaces $X, Y$ (typically Euclidean spaces $\mathbb{R}^n$ or smooth manifolds), representing the domains and codomains of parameterized maps. A morphism $f_\theta: X \to Y$ is defined by a triple $(f, \Theta, \theta)$ where $\Theta$ is a measurable parameter space (often $\mathbb{R}^p$) and $f: X \times \Theta \to Y$ is a measurable function in both arguments. For fixed parameters $\theta \in \Theta$, this yields a deterministic map $f(-, \theta): X \to Y$.

Composition in **Para** is governed by the parameterized product rule. Given morphisms $f_\theta: X \to Y$ and $g_\phi: Y \to Z$, their composite $g_\phi \circ f_\theta: X \to Z$ is defined by the function:

$$ (g \circ f)(x, (\theta, \phi)) = g(f(x, \theta), \phi), $$

equipped with the product parameter space $\Theta \times \Phi$. This composition is strictly associative, with the identity morphism $\mathrm{id}_X: X \to X$ given by $\mathrm{id}_X(x, \emptyset) = x$, where the parameter space is trivial (a singleton). The associativity and unit laws hold exactly:

$$ h_\psi \circ (g_\phi \circ f_\theta) = (h_\psi \circ g_\phi) \circ f_\theta, \quad \mathrm{id}_Y \circ f_\theta = f_\theta = f_\theta \circ \mathrm{id}_X. $$

**Para** is equipped with a symmetric monoidal structure $(\otimes, I)$ for parallel computation. The tensor product $X \otimes Y$ is the Cartesian product $X \times Y$, and the tensor product of morphisms $f_\theta: X \to Y$ and $g_\phi: A \to B$ is:

$$ (f \otimes g)((x, a), (\theta, \phi)) = (f(x, \theta), g(a, \phi)). $$

The monoidal unit $I$ is the singleton space $\{\ast\}$, with unitor isomorphisms $\lambda_X: I \otimes X \to X$ and $\rho_X: X \otimes I \to X$ ensuring that parallel composition with trivial data is coherent. The braiding isomorphism $\sigma_{X,Y}: X \otimes Y \to Y \otimes X$ encodes the commutativity of parallel processes, critical for architectures like attention mechanisms. A gradient structure emerges when **Para** is restricted to differentiable functions. For $f_\theta: \mathbb{R}^n \to \mathbb{R}^m$ with $\Theta = \mathbb{R}^p$, the tangent map $\nabla_\theta f: \mathbb{R}^n \times \mathbb{R}^p \to \mathbb{R}^{m \times p}$ provides a functorial lift to the tangent category, enabling backpropagation as a categorical limit. Universal properties in **Para** characterize optimization: the learning adjunction $L \dagger \mathcal{E}$ relates parameterized functions $f_\theta$ to their empirical losses $\mathcal{E}(f_\theta)$, with the universal arrow $\eta_f: f_\theta \to \mathcal{E}^\dagger(\mathrm{argmin}_\theta \mathcal{E}(f_\theta))$ representing gradient descent.

In essence, **Para** provides a universal language for parameterized computation, where morphisms encapsulate both functional and parametric aspects, and their compositions model the hierarchical and parallel structures inherent in learned systems. The category's monoidal and differential structures rigorously unify forward computation, backpropagation, and parallelization under a single mathematical framework.
### 7.3.2.1 Definition of the Category Para

The category **Para** is defined rigorously as follows:

* **Objects**: Euclidean spaces $\mathbb{R}^n$ (or more generally, smooth manifolds).

* **Morphisms**: A morphism $(f, \theta) : \mathbb{R}^n \to \mathbb{R}^m$ is a smooth function:

$$f : \mathbb{R}^n \times \Theta \to \mathbb{R}^m,$$

where $\Theta$ is the parameter space (also Euclidean).

* **Composition**: As defined above, composition is associative due to the associativity of function composition.

* **Identity**: The identity morphism $\mathrm{id}_{\mathbb{R}^n}$ is the projection $\mathrm{id}_{\mathbb{R}^n}(x, ()) = x$.

The category of parameterized functions, denoted **Para**, is rigorously defined as a symmetric monoidal category that formalizes the structure of parameter-dependent transformations central to machine learning and optimization. Its objects are measurable spaces $X \in \mathrm{Ob}(\mathbf{Para})$, typically Euclidean spaces $\mathbb{R}^n$ or more general smooth manifolds representing data domains. A morphism $f_\theta: X \to Y$ in **Para** consists of:

1. A parameter space $\Theta$ (a measurable space, often $\mathbb{R}^p$)

2. A measurable function $f: X \times \Theta \to Y$

such that for fixed parameters $\theta \in \Theta$, the map $f(-, \theta): X \to Y$ is a deterministic transformation. Composition of morphisms $f_\theta: X \to Y$ and $g_\phi: Y \to Z$ is defined by the parameterized composite:

$$(g \circ f)_{\theta, \phi}(x) := g(f(x, \theta), \phi), \quad \text{where} \quad (\theta, \phi) \in \Theta \times \Phi,$$

equipped with the product parameter space $\Theta \times \Phi$. This composition is strictly associative, with identity morphisms $\mathrm{id}_X: X \to X$ given by the trivial parameterization $\mathrm{id}_X(x, \emptyset) = x$ where the parameter space is the singleton $\{\emptyset\}$. The monoidal structure of **Para** captures parallel computation. The tensor product of objects $X \otimes Y$ is their Cartesian product $X \times Y$, while the tensor product of morphisms $f_\theta: X \to Y$ and $g_\phi: A \to B$ is defined by:

$$(f \otimes g)_{\theta, \phi}(x, a) := (f(x, \theta), g(a, \phi)),$$

with parameter space $\Theta \times \Phi$. The monoidal unit is the singleton space $I = \{\ast\}$, and the associator/unitor isomorphisms are inherited from the Cartesian product in the base category (e.g., $\alpha_{X,Y,Z}: (X \times Y) \times Z \cong X \times (Y \times Z)$).

**Para** is equipped with a gradient structure when modeling differentiable functions. Given a morphism $f_\theta: \mathbb{R}^n \to \mathbb{R}^m$ where $f(-, \theta)$ is smooth, the **parameter derivative** $\nabla_\theta f: \mathbb{R}^n \times \Theta \to \mathbb{R}^{m \times p}$ (for $\Theta = \mathbb{R}^p$) provides the tangent map for gradient-based optimization. This makes **Para** a category internal to the category of smooth manifolds, where morphisms are smooth in both inputs and parameters. The universal property of learning in **Para** arises via adjunctions. For a fixed architecture $F: \mathbf{Para} \to [\mathcal{D}, \mathbf{Set}]$ mapping parameterized functions to their behaviors on data $\mathcal{D}$, the left adjoint $L: [\mathcal{D}, \mathbf{Set}] \to \mathbf{Para}$ constructs the "freely learned" parameterization from data observations, satisfying:

$$\mathrm{Hom}_{\mathbf{Para}}(L(D), f_\theta) \cong \mathrm{Hom}_{[\mathcal{D}, \mathbf{Set}]}(D, F(f_\theta)),$$

where $D: \mathcal{D} \to \mathbf{Set}$ is a dataset functor. This adjunction formalizes empirical risk minimization as a universal construction. **Para** also models **stochastic maps** by allowing $f: X \times \Theta \to \Delta(Y)$,
where $\Delta(Y)$ is the space of probability measures on $Y$. In this case, composition uses the Chapman-Kolmogorov equation for Markov kernels:

$$ (g \circ f)_{\theta, \phi}(x, A) := \int_Y g(y, \phi, A) f(x, \theta, dy), \quad A \subseteq Y \text{ measurable.} $$

This extends **Para** to a Markov category, where Bayesian inference becomes a morphism factorization problem.

The categorical structure of **Para** thus unifies deterministic, stochastic, and differentiable parameterized processes under a single framework, with universal properties ensuring that learning algorithms (e.g., backpropagation, variational inference) arise as canonical constructions. Its morphisms encapsulate both the functional and parameter-space aspects of learned transformations, making it the natural setting for theoretical machine learning.

### 7.3.2.2 **Functoriality of Neural Networks**

A neural network architecture defines a functor from a **computation graph category** (a directed acyclic graph) to **Para**:

* Each node maps to a space $\mathbb{R}^n$,

* Each edge maps to a parameterized morphism $(f, \theta)$,

* **Functoriality** ensures that compositions in the graph correspond to compositions in **Para**.

The functoriality of neural networks constitutes a rigorous categorical framework wherein neural architectures are formalized as functors between appropriately defined categories, preserving compositional structure while encapsulating both parameterized transformations and their learning dynamics. Let **Arch** denote the category of neural architectures, whose objects are tuples $(X, \Theta_X)$ representing layer input/output spaces $X$ and parameter spaces $\Theta_X$, and whose morphisms $f: (X, \Theta_X) \to (Y, \Theta_Y)$ are equivalence classes of parameterized functions $f_\theta: X \to Y$ for $\theta \in \Theta_X \times \Theta_Y$. A network functor $\mathcal{F}: \mathbf{Arch} \to \mathbf{Para}$ maps each architectural component to its parameterized realization, such that:

1. **Object Preservation:** $\mathcal{F}(X, \Theta_X) = X$ (the underlying data space).

2. **Morphism Action:** $\mathcal{F}(f: (X, \Theta_X) \to (Y, \Theta_Y)) = f_\theta: X \times \Theta \to Y$ where $\Theta = \Theta_X \times \Theta_Y$.

3. **Compositionality:** $\mathcal{F}(g \circ f) = \mathcal{F}(g) \circ \mathcal{F}(f)$ for composable $f, g$, with parameter spaces concatenated:

$$ \mathcal{F}(g \circ f)(x, (\theta, \phi)) = g(f(x, \theta), \phi), \quad \theta \in \Theta_X \times \Theta_Y, \phi \in \Theta_Y \times \Theta_Z. $$

4. **Identity Preservation:** $\mathcal{F}(\mathrm{id}_{(X,\Theta_X)}) = \mathrm{id}_X: X \times \{\emptyset\} \to X$.

The functoriality condition ensures that the computational graph of a neural network strictly commutes with its categorical representation: for any diagram $D: \mathcal{J} \to \mathbf{Arch}$ (where $\mathcal{J}$ is a finite indexing category representing the network's graph structure), the functor $\mathcal{F}$ induces a limit-preserving map to **Para**, guaranteeing that parallel and sequential compositions in the architecture correspond to monoidal products and morphism compositions in **Para**.

Backpropagation emerges as a natural transformation $\nabla: \mathcal{F} \Rightarrow \mathcal{T}$ between the network functor $\mathcal{F}$ and a tangent functor $\mathcal{T}$, where $\mathcal{T}(f_\theta) = (f_\theta, \nabla_\theta f_\theta)$ augments each parameterized function with its gradient. The chain rule is encoded via the naturality condition:

$$ \nabla(g \circ f) = (\nabla g \circ \nabla f) \cdot \sigma, $$
where $\sigma$ is the symmetry isomorphism in **Para** permuting gradient terms. This naturality ensures compatibility between the algebraic structure of the network and its differential optimization dynamics. Universal approximation is expressed as a functorial property: let **Top** be the category of topological spaces and continuous maps. The restriction functor $R: \mathbf{Para} \to \mathbf{Top}$ forgets parametrizations, and a network family $\{\mathcal{F}_\alpha\}$ is universally approximating if for any $f: X \to Y$ in Top and compact $K \subset X$, there exists $\alpha$ such that:

$$ \sup_{x \in K} \|R(\mathcal{F}_\alpha)(x) - f(x)\| < \epsilon. $$

Here, $\mathcal{F}_\alpha$ is a functorial neural architecture (e.g., MLPs with width $\alpha$), and the condition asserts that the image of $\mathcal{F}_\alpha$ under $R$ is dense in the space of continuous functions. Factoriality also governs parameter updates: stochastic gradient descent defines an endofunctor $SGD: \mathbf{Para} \to \mathbf{Para}$ acting on morphisms by:

$$ SGD(f_\theta)(x, \theta') = f_{\theta-\eta\nabla_\theta\mathcal{L}}(x), \quad \theta' \in \Theta \text{ (updated parameters)}, $$

where $\mathcal{L}$ is a loss functor **Para** $\to \mathbb{R}$. The factoriality of $SGD$ ensures that parameter updates commute with network composition, preserving the coherence of distributed training protocols.

In summary, the functorial perspective rigorously unifies the algebraic, differential, and approximation-theoretic aspects of neural networks, demonstrating that their empirical success stems from their inherent categorical structure as parameter-space functors preserving compositional and optimization-theoretic invariants.

### 7.3.2.3 Monoidal Structure (for Parallel Composition)

If we extend **Para** to a **monoidal category**, we can model parallel layer compositions (e.g., in multi-head attention or residual connections):

* Tensor product $\otimes$ combines spaces and parameters:

$$ (f, \theta) \otimes (g, \phi) = (f \times g, (\theta, \phi)) : \mathbb{R}^n \times \mathbb{R}^m \to \mathbb{R}^p \times \mathbb{R}^q. $$

The monoidal structure in the context of parallel composition formalizes the algebraic and categorical foundations necessary to rigorously describe simultaneous, independent computations within neural networks and other parallelizable systems. At its core, a monoidal category $(C, \otimes, I)$ consists of a category $C$ equipped with a tensor product $\otimes: C \times C \to C$ and a unit object $I$, satisfying associativity and unit constraints up to natural isomorphism. For parallel computation in neural networks, the tensor product $\otimes$ represents the concurrent application of operations to distinct data streams, while the unit object $I$ corresponds to the trivial case of no computational load.

Given two neural network layers $f: X \to Y$ and $g: A \to B$, their parallel composition is encoded as the tensor product $f \otimes g: X \otimes A \to Y \otimes B$, where $X \otimes A$ and $Y \otimes B$ denote the concatenation or Cartesian product of the respective input and output spaces. Explicitly, for $x \in X$ and $a \in A$, the action of $f \otimes g$ is defined by:

$$ (f \otimes g)(x, a) = (f(x), g(a)). $$

This operation preserves the independence of $f$ and $g$, ensuring no spurious interactions between the parallel streams. The associativity of $\otimes$ is captured by the natural isomorphism $\alpha_{X,Y,Z}: (X \otimes Y) \otimes Z \to X \otimes (Y \otimes Z)$, which allows re-parenthesization of parallel computations without altering their semantic meaning. For neural networks, this justifies the arbitrary grouping of parallel branches, as in multi-head attention or residual connections, where the order of combining channels or features
is irrelevant to the final output. The unit object $I$ satisfies the isomorphisms $\lambda_X: I \otimes X \to X$ and $\rho_X: X \otimes I \to X$, which assert that parallel composition with an empty or trivial process leaves the original computation unchanged. In practice, this corresponds to identity skip connections or zero-padding in neural architectures, where the presence of an inert branch does not affect the active computation. The monoidal coherence conditions—pentagon and triangle identities—ensure these isomorphisms interact harmoniously, guaranteeing that diagrams of parallel computations commute. The interchange law governs the interaction between sequential and parallel composition, stating that for morphisms $f, h: X \to Y$ and $g, k: A \to B$, the following equality holds:

$$ (h \circ f) \otimes (k \circ g) = (h \otimes k) \circ (f \otimes g). $$

This law enforces that independent sequences of operations can be composed either sequentially within each branch or in parallel across branches, yielding identical results. In neural networks, this underlies the modularity of layers, where parallel sub-networks (e.g., convolutional filters) process inputs independently before aggregation. For symmetric monoidal categories, the additional natural isomorphism $\sigma_{X,Y}: X \otimes Y \to Y \otimes X$ encodes the commutativity of parallel operations, permitting the reordering of data streams. This symmetry is critical for architectures like attention mechanisms, where the output is invariant to the permutation of input heads. The coherence of $\sigma$ with the associator and unitors ensures that arbitrary rearrangements of parallel computations are semantically consistent.

In summary, the monoidal structure provides the algebraic scaffolding for parallel computation, where the tensor product $\otimes$ models concurrent execution, associativity and unit isomorphisms guarantee compositional flexibility, and the interchange law enforces the coherence of sequential and parallel operations. This framework rigorously captures the essence of parallelism in neural networks, from basic layer-wise processing to complex modular architectures, while ensuring mathematical consistency through categorical axioms.

### 7.3.3 Role of Associativity and Identity in Sequential Models

The role of associativity and identity in sequential models is foundational to their mathematical coherence and operational consistency, rigorously governed by the axioms of category theory. In the context of neural networks and other compositional systems, associativity ensures that the chaining of layers is unambiguous, while identity morphisms guarantee the existence of neutral transformations that preserve information flow.

For a sequential model composed of layers $f: X \to Y$, $g: Y \to Z$, and $h: Z \to W$, the associativity axiom demands that the order of function composition does not alter the model's output:

$$ h \circ (g \circ f) = (h \circ g) \circ f. $$

This equality holds strictly when the layers are deterministic functions, as in feedforward networks, where $(h \circ g \circ f)(x) = h(g(f(x)))$ is invariant to parenthesization. In parameterized settings (e.g., the category **Para**), associativity extends to parameter spaces: for $f_\theta$, $g_\phi$, and $h_\psi$, the composite $(h_\psi \circ g_\phi) \circ f_\theta$ and $h_\psi \circ (g_\phi \circ f_\theta)$ share the same parameterization $(\theta, \phi, \psi)$ and map $x \mapsto h(g(f(x, \theta), \phi), \psi)$. This property is critical for modular architectures like ResNets or Transformers, where submodules can be freely recomposed without altering the model's semantics. The identity morphism $\mathrm{id}_X: X \to X$ acts as a trivial layer satisfying:

$$ f_{\theta} \circ \mathrm{id}_{X} = f_{\theta} = \mathrm{id}_{Y} \circ f_{\theta}, $$

for any $f_\theta: X \to Y$. In **Para**, the identity is defined by $\mathrm{id}_X(x, \emptyset) = x$, where the parameter space is empty. This axiom enforces that skip connections or residual branches behave as pure
information passthroughs, ensuring that $f_{\theta} + \text{id}_X$ (as in ResNets) is well-defined. The identity's naturality extends to gradient computations, where $\nabla_{\theta}(\text{id}_X \circ f_{\theta}) = \nabla_{\theta}f_{\theta}$, preserving the chain rule in backpropagation. Associativity and identity jointly define the category of sequential models, where objects are data spaces and morphisms are (parameterized) transformations. The axioms ensure that any diagram of composable layers commutes, guaranteeing that the model's output depends only on the functional pipeline, not on the order of intermediate operations. This categorical structure underpins the interpretability and correctness of neural architectures, as violations of associativity or identity would introduce ambiguities in forward inference and gradient propagation.

In practice, these axioms are approximately preserved even in stochastic or approximate settings (e.g., dropout, quantization), where associativity holds up to noise terms and identity is realized as a small-norm perturbation. The universality of these properties across architectures—from CNNs to diffusion models—demonstrates their necessity in unifying sequential computation under a rigorous algebraic framework.

### 7.3.3.1 Associativity of Composition

Associativity of composition is a fundamental axiom in category theory that ensures the unambiguous chaining of morphisms, providing the algebraic foundation for sequential processes in mathematics and machine learning. Formally, given three composable morphisms $f: A \to B, g: B \to C$, and $h: C \to D$ in a category $\mathcal{C}$, associativity mandates that the order of composition does not affect the resulting morphism:

$$h \circ (g \circ f) = (h \circ g) \circ f.$$

This equality asserts that the diagram formed by the compositions commutes, meaning the path taken through intermediate objects $B$ and $C$ is irrelevant to the final map from $A$ to $D$. In the context of neural networks, where morphisms are parameterized functions $f_{\theta}: \mathbb{R}^n \to \mathbb{R}^m$, associativity guarantees that the forward pass of a deep network is invariant to the grouping of layers:

$$(h_{\psi} \circ g_{\phi}) \circ f_{\theta}(x) = h_{\psi}(g_{\phi}(f_{\theta}(x))) = h_{\psi} \circ (g_{\phi} \circ f_{\theta})(x).$$

Here, the parameters $\theta, \phi, \psi$ concatenate via the product space $\Theta \times \Phi \times \Psi$, and the equality holds strictly for deterministic layers.

The associativity axiom extends to higher-order compositions in monoidal categories, where parallel and sequential operations interact. For morphisms $f: A \to B, g: B \to C$, and $h: C \to D$, the interchange law in a monoidal category $\mathcal{C}$ ensures:

$$(h \otimes \text{id}_Y) \circ (\text{id}_Z \otimes g) \circ (f \otimes \text{id}_W) = (h \circ g \circ f) \otimes \text{id}_{Y \otimes W},$$

where $\otimes$ is the tensor product representing parallel composition. This generalizes associativity to networks with branching architectures, such as multi-head attention or residual connections, where the coherence of parallel and sequential paths is preserved. In stochastic settings, associativity holds up to measurable isomorphisms. For Markov kernels $f: A \to \Delta(B), g: B \to \Delta(C)$, and $h: C \to \Delta(D)$, the Chapman-Kolmogorov equation enforces:

$$\int_C h(c, D) d(g \circ f)(c) = \int_B \left( \int_C h(c, D) dg(b, c) \right) df(a, b),$$

where $d(g \circ f)(c)$ is the pushforward measure. This probabilistic associativity underpins Bayesian networks and diffusion models, where the order of marginalization is invariant.

Associativity is not merely an algebraic convenience but a universal constraint ensuring the consistency of symbolic and computational representations. It allows complex systems to be decomposed
into modular components, analyzed locally, and recomposed globally without ambiguity. Violations of associativity—such as in approximate or noisy computations—introduce non-trivial curvature in the category’s compositional structure, necessitating corrective mechanisms like regularization or synchronization protocols in distributed learning.

* The associativity law in **Para** ensures that composing layers in any order (while respecting dependencies) yields the same function:

$$h \circ (g \circ f) = (h \circ g) \circ f.$$

* This justifies modular network design: we can group layers into submodules without changing behavior.

Ultimately, associativity of composition is the linchpin of categorical reasoning, enabling the rigorous study of neural networks, programming languages, and physical systems through the lens of morphism coherence and diagrammatic invariance.

### 7.3.3.2 Identity Morphisms and Skip Connections

Identity morphisms and skip connections are fundamental constructs in both category theory and neural network design, ensuring stability and coherence in information propagation. In a category $C$, the identity morphism $\mathrm{id}_X: X \to X$ for an object $X$ satisfies the axioms $\mathrm{id}_X \circ f = f$ and $g \circ \mathrm{id}_X = g$ for all morphisms $f: Y \to X$ and $g: X \to Z$. For parameterized functions in the category **Para**, the identity is realized as $\mathrm{id}_X(x, \emptyset) = x$, where the parameter space is trivial (a singleton), preserving input data under null transformation. In neural networks, skip connections operationalize identity morphisms by creating direct pathways that bypass nonlinear transformations. Given a layer $f_\theta: \mathbb{R}^n \to \mathbb{R}^n$, a residual connection implements $F_\theta(x) = f_\theta(x) + \mathrm{id}_X(x)$, where the identity term $\mathrm{id}_X(x) = x$ ensures gradient stability during backpropagation. The derivative of this operation, $\nabla_x F_\theta(x) = \nabla_x f_\theta(x) + I_n$ (with $I_n$ the $n \times n$ identity matrix), guarantees non-vanishing gradients by preserving the dominant eigenvalue of the Jacobian.

The categorical perspective reveals that skip connections enforce universal properties of identity. For any perturbation $\epsilon: X \to X$ (e.g., noise or dropout), the composite $(\mathrm{id}_X + \epsilon) \circ f_\theta$ approximates $f_\theta$ when $\|\epsilon\| \ll 1$, ensuring robustness. In monoidal categories, the coherence condition $\rho_X: X \otimes I \to X$ justifies the initialization of residual branches as identity maps, as in zero-initialized residual networks where $f_\theta(x) \to 0$ early in training, reducing initial output variance. For stochastic networks, skip connections preserve the identity up to measure: if $f_\theta: X \to \Delta(X)$ is a stochastic layer, the residual version $F_\theta(x) = f_\theta(x) * \delta_x$ (a convex combination with Dirac measure $\delta_x$) ensures the input distribution remains dominant. This aligns with the categorical definition of identity for Markov kernels, where $\mathrm{id}_X(x, A) = \delta_x(A)$ for measurable $A \subseteq X$.

* The identity morphism $\mathrm{id}_X$ allows for **skip connections** (e.g., ResNet):

$$\mathrm{ResBlock}(x, \theta) = f(x, \theta) + \mathrm{id}_X(x) = f(x, \theta) + x.$$

* Identity ensures that a “null” layer (doing nothing) is a valid network component.

The interplay between identity morphisms and skip connections thus provides a rigorous framework for stable deep learning, where categorical axioms ensure both mathematical consistency and empirical performance in architectures like ResNets and Transformers. The identity’s naturality—its preservation under composition and tensor products—underpins the design of modular, scalable networks resistant to degradation.
### 7.3.3.3 Universality of Sequential Models

The universality of sequential models is a rigorous mathematical property asserting their capacity to approximate arbitrary functional mappings to arbitrary precision under specified conditions, formalized through categorical and analytical frameworks. Let $\mathcal{F}$ denote the class of sequential models (e.g., recurrent networks, Transformers) with composition operators ◦, and let $\mathcal{C}(X, Y)$ be the space of continuous functions between compact domains $X \subset \mathbb{R}^n$ and $Y \subset \mathbb{R}^m$. A model family $\mathcal{F}$ is universal if for every $f \in \mathcal{C}(X, Y)$ and $\epsilon > 0$, there exists a finite sequence of parameterized morphisms $\{f_{\theta_i}\}_{i=1}^k \in \mathcal{F}$ such that:

$$ \sup_{x \in X} \|(f_{\theta_k} \circ \cdots \circ f_{\theta_1})(x) - f(x)\| < \epsilon. $$

This approximation property emerges from the Stone-Weierstrass theorem when $\mathcal{F}$ forms an algebra closed under composition, pointwise addition, and scalar multiplication, and separates points in $\mathcal{C}(X, Y)$.

In categorical terms, universality is characterized by the density of $\mathcal{F}$ in the functor category $[\mathcal{C}, \mathcal{D}]$, where $\mathcal{C}$ encodes input data structures and $\mathcal{D}$ represents outputs. For a sequential model $F: \mathcal{C} \to \mathcal{D}$ to be universal, the Yoneda embedding must exhibit a natural isomorphism:

$$ \mathrm{Hom}_{[\mathcal{C}, \mathcal{D}]}(F, G) \cong \mathrm{Hom}_{\mathcal{D}}(F(X), G(X)) $$

for all functors $G$ representing target functions, ensuring that any $G$ can be approximated by a sufficiently deep $F$. The approximation depth is governed by the model's width and nonlinearity. For a ReLU network $F_\theta: \mathbb{R}^n \to \mathbb{R}^m$ with $L$ layers, the minimal width $w$ required for $\epsilon$-approximation of Lipschitz functions scales as $w = O(\epsilon^{-n})$, following the Yarotsky bound. Transformers achieve universality via their attention mechanism, where the self-attention operation $\mathrm{Attn}(Q, K, V) = \mathrm{softmax}(QK^T/\sqrt{d})V$ densely approximates permutation-invariant operators through the universal approximation theorem for kernel integral transforms. Dynamical universality in recurrent models arises from the liquid state machine framework: a continuous-time RNN with sufficient hidden units $h$ can approximate fading-memory filters $F: \mathcal{C}(\mathbb{R}_+, \mathbb{R}^n) \to \mathcal{C}(\mathbb{R}_+, \mathbb{R}^m)$ as:

$$ \|F(x)(t) - W \cdot r(t)\| < \epsilon, \quad r(t) = \sigma(Ar(t-1) + Bx(t)), $$

where $r \in \mathbb{R}^h$ is the reservoir state and $A, B, W$ are learnable parameters.

The universality property is preserved under categorical limits: if each component $f_{\theta_i}$ in a sequential model is universal, their colimit (e.g., infinite-depth limit) remains universal. This underpins the efficacy of deep architectures, where the functoriality of composition guarantees that approximation bounds compound multiplicatively rather than additively.

* The ability to compose morphisms arbitrarily allows for **universal approximation** (e.g., via deep feedforward networks).

* The categorical framework generalizes to **recurrent networks** by working in a category of dynamical systems, where morphisms are parameterized recurrent cells.

Thus, universality in sequential models is a confluence of analytic approximation theory, categorical density, and dynamical systems theory, ensuring that compositional hierarchies of simple functions can encapsulate the infinite-dimensional space of continuous operators.

### 7.3.4 Conclusion

By formalizing neural networks in the category **Para**, we obtain:
1. Layers as morphisms, with networks as their compositions.

2. A rigorous category of parameterized functions, supporting functorial and monoidal structures.

3. **Associativity and identity** as fundamental properties enabling modular and correct-by-construction network design.

This framework extends naturally to more complex architectures (e.g., transformers, graph networks) by enriching **Para** with additional categorical structure (e.g., symmetric monoidal categories for parallel processing).
# 8 Open Set Learning

Open set learning is a foundational challenge in statistical learning theory and pattern recognition, characterized by the necessity of classifying data instances drawn from both known and unknown distributions. Consider a training dataset $\mathcal{D} = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$, where each instance $\mathbf{x}_i \in \mathbb{R}^d$ is associated with a label $y_i$ from a finite set of known classes $\mathcal{Y}_{\text{train}}$. The fundamental challenge in open set learning arises when test samples are drawn from a distribution $P_{\text{test}}$ that includes both the training distribution $P_{\text{train}}$ and an additional unknown component $P_{\text{unknown}}$, such that

$$P_{\text{test}} = P_{\text{train}} + P_{\text{unknown}}, \quad P_{\text{unknown}} \cap P_{\text{train}} = \emptyset. \qquad (8.1)$$

In classical closed-set classification, one seeks to find a function $f: \mathbb{R}^d \to \mathbb{R}^C$ that minimizes an empirical risk functional of the form

$$\mathcal{R}_{\text{emp}}(f) = \mathbb{E}_{(\mathbf{x},y)\sim P_{\text{train}}} [\ell(f(\mathbf{x}), y)] \qquad (8.2)$$

where $\ell: \mathbb{R}^C \times \mathbb{R}^C \to \mathbb{R}$ is a loss function such as categorical cross-entropy:

$$\ell(f(\mathbf{x}), y) = - \sum_{c=1}^{C} \mathbb{1}[y=c] \log f_c(\mathbf{x}). \qquad (8.3)$$

However, in the presence of unknown classes, minimizing only this risk function leads to over-confident misclassifications when a sample $\mathbf{x}_u \sim P_{\text{unknown}}$ is assigned to one of the known classes. This necessitates the introduction of a rejection mechanism, leading to an extended open set risk functional

$$\mathcal{R}_{\text{open}}(f) = \mathbb{E}_{(\mathbf{x},y)\sim P_{\text{train}}} [\ell(f(\mathbf{x}), y)] + \lambda \mathbb{E}_{\mathbf{x}_u\sim P_{\text{unknown}}} [\ell_{\text{reject}}(f(\mathbf{x}_u))] \qquad (8.4)$$

where $\ell_{\text{reject}}$ is a loss function that penalizes misclassification of unknown samples. A common approach is to introduce a threshold-based rejection criterion where the classification function $f$ satisfies

$$\ell_{\text{reject}}(f(\mathbf{x}_u)) = \max_c f_c(\mathbf{x}_u) - \tau, \qquad (8.5)$$

where $\tau$ is a rejection threshold such that if $\max_c f_c(\mathbf{x}_u) < \tau$, the sample is assigned to an unknown category. The introduction of this threshold formally defines an open space $\mathcal{O}(f)$, given by

$$\mathcal{O}(f) = \{ \mathbf{x} \in \mathbb{R}^d \mid \max_c f_c(\mathbf{x}) < \tau \}, \qquad (8.6)$$

which is crucial in minimizing the open space risk introduced by Scheirer et al. (2013):

$$\mathcal{R}_{\text{open}}(f) = \int_{\mathcal{O}(f)} P(\mathbf{x})d\mathbf{x}. \qquad (8.7)$$

Incorporating open space risk minimization into the learning framework, the overall objective function becomes

$$\mathcal{L}(f) = \mathcal{R}_{\text{emp}}(f) + \beta \mathcal{R}_{\text{open}}(f), \qquad (8.8)$$

where $\beta$ is a hyperparameter that balances classification accuracy with rejection performance. The underlying theoretical framework necessitates consideration of both discriminative and generative
approaches. In discriminative modeling, the classification function $f$ can be regularized to enforce an explicit margin $\gamma$ for open set rejection:

$$ \min_{\mathbf{x}_u \sim P_{\text{unknown}}} \max_{c} f_c(\mathbf{x}_u) \leq \gamma. \qquad (8.9) $$

In generative modeling, the problem is framed as estimating the likelihood $P(\mathbf{x} | y)$ under a known class distribution model, such as a Gaussian Mixture Model (GMM):

$$ P(\mathbf{x} | y = c) = \sum_{k=1}^{K} \pi_k \mathcal{N}(\mathbf{x} | \mu_k, \Sigma_k), \qquad (8.10) $$

where $\mathcal{N}(\mathbf{x} | \mu_k, \Sigma_k)$ represents a multivariate Gaussian distribution with mean $\mu_k$ and covariance matrix $\Sigma_k$, and $\pi_k$ is the prior weight of the $k$-th component. A likelihood-based rejection rule is then applied:

$$ P(\mathbf{x}) < \delta \Rightarrow \text{unknown}. \qquad (8.11) $$

There are several class distribution models in which the open set learning problem can be framed as estimating the likelihood $P(\mathbf{x} | y)$. To state all those class distribution models we first define few quantities. Let $\mathcal{X} \subset \mathbb{R}^d$ be the input feature space and $\mathcal{Y} = \{1, 2, \dots, C\}$ be the known set of class labels. We assume that samples are drawn from a joint probability distribution $P(\mathbf{x}, y)$, where the true distribution of unknown classes is not available during training. Open Set Learning requires estimating a function $f: \mathcal{X} \to \mathcal{Y} \cup \{\emptyset\}$, where $\emptyset$ represents an unknown class.

1. **Gaussian Distribution Model (GDM)** The fundamental assumption is that the feature distribution of each known class follows a multivariate normal distribution parameterized by the mean vector $\mu_c$ and covariance matrix $\Sigma_c$. The likelihood of a given sample $\mathbf{x}$ belonging to class $c$ is:

$$ P(\mathbf{x} | y = c) = \frac{1}{(2\pi)^{d/2} |\Sigma_c|^{1/2}} \exp \left( -\frac{1}{2} (\mathbf{x} - \mu_c)^T \Sigma_c^{-1} (\mathbf{x} - \mu_c) \right). \qquad (8.12) $$

To quantify the confidence in assigning $\mathbf{x}$ to class $c$, we compute the Mahalanobis distance:

$$ D_M(\mathbf{x}, \mu_c, \Sigma_c) = \sqrt{(\mathbf{x} - \mu_c)^T \Sigma_c^{-1} (\mathbf{x} - \mu_c)}. \qquad (8.13) $$

A sample is rejected as unknown if:

$$ \min_{c} D_M(\mathbf{x}, \mu_c, \Sigma_c) > \delta, \qquad (8.14) $$

where $\delta$ is a threshold chosen based on extreme value statistics.

2. **Gaussian Mixture Model (GMM)** The Gaussian assumption can be generalized using a Gaussian Mixture Model (GMM), which represents each class as a weighted sum of multiple Gaussian components:

$$ P(\mathbf{x} | y = c) = \sum_{k=1}^{K} \pi_k \mathcal{N}(\mathbf{x} | \mu_k, \Sigma_k), \qquad (8.15) $$

where $\pi_k$ are the mixture weights satisfying $\sum_{k=1}^{K} \pi_k = 1$, and each Gaussian component $\mathcal{N}(\mathbf{x} | \mu_k, \Sigma_k)$ is given by:

$$ \mathcal{N}(\mathbf{x} | \mu_k, \Sigma_k) = \frac{1}{(2\pi)^{d/2} |\Sigma_k|^{1/2}} \exp \left( -\frac{1}{2} (\mathbf{x} - \mu_k)^T \Sigma_k^{-1} (\mathbf{x} - \mu_k) \right). \qquad (8.16) $$

Unknown samples are rejected based on low maximum likelihood estimation (MLE) scores:

$$ \max_{c} P(\mathbf{x} | y = c) < \tau. \qquad (8.17) $$

where $\tau$ is a predefined threshold.
3. Dirichlet Process Gaussian Mixture Model (DP-GMM) A Dirichlet Process (DP) prior can be introduced to allow the number of mixture components K to grow dynamically. The prior over the mixture weights follows:

$$ \pi_k \sim \text{Dir}(\alpha), \qquad (8.18) $$

where α is the concentration parameter controlling cluster sparsity. This enables automatic adaptation of the number of mixture components to better capture class distributions. The likelihood of x follows:

$$ P(\mathbf{x} | y = c) = \sum_{k=1}^{K} \pi_k \mathcal{N}(\mathbf{x} | \mu_k, \Sigma_k), \qquad (8.19) $$

but with a nonparametric prior that ensures more flexible decision boundaries.

4. Extreme Value Theory (EVT) Models The tail distribution of softmax probabilities is modeled using an Extreme Value Theorem (EVT) approach. Given softmax scores $\sigma_c(\mathbf{x})$, we fit a Weibull distribution to the tail:

$$ P_{\text{Weibull}}(x) = 1 - \exp\left(-\left(\frac{x - \lambda}{\kappa}\right)^{\beta}\right). \qquad (8.20) $$

A sample x is rejected if:

$$ \max_{c} \sigma_{c}(\mathbf{x}) < \tau_{\text{EVT}}. \qquad (8.21) $$

5. Bayesian Neural Networks (BNNs) BNNs introduce uncertainty estimation by placing priors over network weights:

$$ w \sim P(w). \qquad (8.22) $$

Posterior inference is performed via Bayesian updating:

$$ P(y | \mathbf{x}, \mathcal{D}) = \int P(y | \mathbf{x}, w) P(w | \mathcal{D}) dw. \qquad (8.23) $$

A sample is rejected if the entropy of the predictive distribution is high:

$$ H[P(y | \mathbf{x})] = - \sum_{c} P(y = c | \mathbf{x}) \log P(y = c | \mathbf{x}). \qquad (8.24) $$

6. Support Vector Models (OC-SVM and SVDD) One-Class SVM (OC-SVM): Finds a separating hyperplane $w^T\phi(\mathbf{x}) + b = 0$ such that:

$$ \min_{w,b} \frac{1}{2} \|w\|^2 \qquad (8.25) $$

subject to $w^T\phi(\mathbf{x}) + b \ge \rho$. A sample is rejected if:

$$ w^T\phi(\mathbf{x}) + b < \rho. \qquad (8.26) $$

7. Support Vector Data Description (SVDD): Finds a minimum enclosing hypersphere with center c and radius R:

$$ \min_{R,c} R^2 + C \sum_{i} \xi_i \qquad (8.27) $$

subject to:

$$ \|\phi(\mathbf{x}_i) - \mathbf{c}\|^2 \le R^2 + \xi_i. \qquad (8.28) $$

A sample is rejected if:

$$ \|\phi(\mathbf{x}) - \mathbf{c}\| > R. \qquad (8.29) $$