This differential learning rate for frequency components leads to the spectral regularization phenomenon, where the network learns the low-frequency components of the function first, and the high-frequency modes only begin to adapt once the low-frequency ones have been approximated with sufficient accuracy. In a more formal setting, the spectral bias can also be understood in terms of Sobolev spaces. A neural network function $f_{NN}$ can be seen as a function in a Sobolev space

<figure><img src="image_51.png" /><figcaption>Figure 4.15: Singular Value Spectrum of First Layer (MNIST)</figcaption></figure>

$W^{m,2}$, where the norm of a function $f$ in this space is defined as

$$ \|f\|_{W^{m,2}}^2 = \int_{\mathbb{R}^d} (1 + \|\xi\|^2)^m |\hat{f}(\xi)|^2 d\xi \quad (4.85) $$

When training a neural network, the optimization process implicitly regularizes the higher-order Sobolev norms, meaning that the network will initially approximate the target function in terms of lower-order derivatives (which correspond to low-frequency modes). This can be expressed by introducing a regularization term in the loss function:

$$ \mathcal{L}_{\text{eff}}(\boldsymbol{\theta}) = \mathcal{L}(\boldsymbol{\theta}) + \lambda \|f_{\text{NN}}\|_{W^{m,2}}^2 \quad (4.86) $$

where $\lambda$ is a regularization parameter that controls the trade-off between data fidelity and smoothness in the approximation.

Thus, spectral regularization emerges as a consequence of the network's architecture, the nature of gradient descent optimization, and the inherent smoothness of the functions that neural networks are capable of learning. The mathematical structure of the NTK and the regularization properties of the Sobolev spaces provide a rigorous framework for understanding why neural networks prioritize the learning of low-frequency modes, reinforcing the idea that neural networks are implicitly biased toward smooth, low-frequency approximations at the beginning of training. This insight has profound implications for the generalization behavior of neural networks and their capacity to approximate complex functions.
```python
import torch.nn(functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

# --- Define simple MLP ---
class SimpleNet nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = F relu(self.fc1(x))
        return self.fc2(x)

# --- Training function with spectral regularization ---
def train_model(lambda_reg=0.0, epochs=5):
    model = SimpleNet()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('', train=True, download=True, transform=
            transforms.ToTensor()),
        batch_size=128, shuffle=True
    )

    singular vals = []

    for epoch in range(epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zeroGrad()
            output = model(data)
            loss = F cross-projection output, target)

            # --- spectral regularization ---
            W = model.fc1.weight
            u, s, v = torch.svd(W) # singular values
            spectral_norm = s[0] # largest singular value
            loss = loss + lambda_reg * spectral_norm**2

            loss backward()
            optimizer.step()

            # Track singular values at epoch end
            with torch.noGrad():
                u, s, v = torch.svd(model.fc1.weight)
                singular vals.append(s.cpu(). numpy())

    return singular vals
```

```python
# --- Run training with and without spectral regularization ---
sing vals_no_reg = train_model(lambda_reg=0.0, epochs=5)
sing vals_reg = train_model(lambda_reg=1e-3, epochs=5)
```
57

58

59

60

61

62

63

64

65

66

67

68

69

70

# --- Plot singular value spectrum ---

plt.figure(figsize=(10,6))

for i, s in enumerate(sing vals no reg):

plt .plot(s, 'o-', alpha=0.5, label=f"No Reg Epoch {i+1}" if i==0 else "")

for i, s in enumerate(sing vals reg):

plt .plot(s, 's-', alpha=0.5, label=f"Reg Epoch {i+1}" if i==0 else "")

plt .xlabel("Singular value index")

plt .ylabel("Magnitude")

67 plt .title("Singular Value Spectrum of First Layer (MNIST)")

68 plt .legend()

69 plt .grid(True)

70 plt .show()
# 5 Game-theoretic formulations of Deep Neural Networks

The training of deep neural networks (DNNs) can be rigorously analyzed through game-theoretic frameworks, where optimization is interpreted as a competitive or cooperative interaction between different components of the model. In this setting, the learning process is modeled as a multi-player game, where each player corresponds to a subset of parameters (e.g., layers, neurons, or auxiliary networks), and their strategies are parameter updates that collectively influence the network's performance. Let $\theta = (\theta_1, \dots, \theta_K)$ denote the parameters partitioned among K players. The objective is to minimize a global loss function $\mathcal{L}(\theta)$, which induces a potential game where each player's cost function aligns with the global loss. The Nash equilibrium of this game corresponds to a stationary point of $\mathcal{L}$, satisfying:

$$ \nabla_{\theta_k} \mathcal{L}(\theta) = 0, \quad \forall k \in \{1, \dots, K\}. \qquad (5.1) $$

However, since exact equilibria are computationally intractable, training dynamics are often studied through online learning in games, where players iteratively adjust their strategies via gradient-based updates. For example, in generative adversarial networks (GANs), the generator $G_{\theta_G}$ and discriminator $D_{\theta_D}$ engage in a zero-sum game with value function:

$$ V(\theta_G, \theta_D) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}} [\log D(\mathbf{x})] + \mathbb{E}_{\mathbf{z} \sim p_{\mathbf{z}}} [\log(1 - D(G(\mathbf{z})))], \qquad (5.2) $$

where the generator seeks to minimize $V$ while the discriminator maximizes it. The equilibrium corresponds to a saddle point $(\theta_G^*, \theta_D^*)$ satisfying:

$$ V(\theta_G^*, \theta_D) \leq V(\theta_G^*, \theta_D^*) \leq V(\theta_G, \theta_D^*), \quad \forall \theta_G, \theta_D. \qquad (5.3) $$

The training dynamics of DNNs can also be formalized using differential game theory, where parameter updates follow continuous-time gradient flows. Consider a two-player game between feature extractors and classifiers in a deep network, where the state evolves according to the coupled ODEs:

$$ \frac{d\theta_1}{dt} = -\nabla_{\theta_1} \mathcal{L}(\theta_1, \theta_2), \quad \frac{d\theta_2}{dt} = -\nabla_{\theta_2} \mathcal{L}(\theta_1, \theta_2). \qquad (5.4) $$

The stability and convergence of such dynamics depend on the spectral properties of the game's Jacobian matrix $\mathbf{J}$, defined as:

$$ \mathbf{J} = \begin{pmatrix} \nabla_{\theta_1}^2 \mathcal{L} & \nabla_{\theta_1} \nabla_{\theta_2} \mathcal{L} \\ \nabla_{\theta_2} \nabla_{\theta_1} \mathcal{L} & \nabla_{\theta_2}^2 \mathcal{L} \end{pmatrix}. \qquad (5.5) $$

If $\mathbf{J}$ has eigenvalues with negative real parts, the system converges to a stable equilibrium; otherwise, limit cycles or divergence may occur, as seen in GAN training. Furthermore, Stackelberg game formulations introduce hierarchical optimization, where one player (the leader, e.g., hyperparameters) commits to a strategy first, and the other (the follower, e.g., model parameters) responds optimally. This leads to bilevel optimization problems of the form:

$$ \min_{\theta_1} \mathcal{L}_{\text{outer}}(\theta_1, \theta_2^*(\theta_1)), \quad \text{where} \quad \theta_2^*(\theta_1) = \arg \min_{\theta_2} \mathcal{L}_{\text{inner}}(\theta_1, \theta_2). \qquad (5.6) $$

These game-theoretic perspectives unify adversarial training, meta-learning, and multi-agent RL under a common framework, providing insights into convergence, robustness, and implicit regularization in deep learning.
## 5.1 Literature Review of Game-theoretic formulations of Deep Neural Networks

The game-theoretic formulations of deep neural networks have been explored extensively in recent literature, offering novel perspectives on optimization, adversarial robustness, and multi-agent interactions. One seminal contribution is the work by Goodfellow et al. (2014) [1472] on generative adversarial networks (GANs), which frames the training process as a two-player minimax game between a generator and a discriminator. This formulation has revolutionized unsupervised learning by enabling the generation of highly realistic data. The theoretical foundations of GANs have been further refined by Arjovsky et al. (2017) [1459], who introduced Wasserstein GANs to address convergence issues, providing a more stable training framework through the lens of optimal transport theory. Building on this, Gulrajani et al. (2017) [1473] proposed gradient penalty methods to enforce Lipschitz constraints, improving training stability. Another critical contribution comes from Shalev-Shwartz et al. (2012) [1495], who interpret supervised learning as a convex-concave game, establishing connections between empirical risk minimization and game-theoretic equilibria. This perspective has led to deeper insights into the generalization properties of neural networks, further explored by Hardt et al. (2016) [1475] through the lens of algorithmic stability.

The interplay between game theory and deep learning is also evident in adversarial robustness, where Madry et al. (2018) [1486] model adversarial attacks as a zero-sum game between a classifier and an adversary. Their work formalizes robust training as a min-max optimization problem, leading to the development of adversarially robust models. Similarly, Sinha et al. (2018) [1497] leverage distributionally robust optimization to frame adversarial training as a game against nature, providing probabilistic guarantees on model performance. These approaches highlight the utility of game-theoretic frameworks in addressing security challenges in deep learning, with additional insights from Raghunathan et al. (2018) [1492] on certified defenses and Zhang et al. (2019) [1502] on trade-offs between robustness and accuracy. Multi-agent reinforcement learning (MARL) is another domain where game theory plays a pivotal role. Littman's (1994) [1483] formulation of Markov games extends single-agent reinforcement learning to competitive and cooperative settings, enabling the study of Nash equilibria in multi-agent systems. This has been expanded upon by Foerster et al. (2018) [1469], who introduce learning dynamics with opponent modeling, combining deep Q-learning with game-theoretic equilibria. Their work demonstrates how agents can adapt to others' strategies in complex environments, a theme further developed by Lowe et al. (2017) [1484] with multi-agent actor-critic methods. Another notable contribution is by Lanctot et al. (2017) [1480], who propose the use of fictitious play and regret minimization in deep MARL, bridging the gap between classical game theory and modern deep learning techniques. Recent advances by Silver et al. (2017) [1496] in AlphaGo and OpenAI (2019) [1488] in multi-agent hide-and-seek showcase the practical success of these frameworks.

The optimization of neural networks has also been reinterpreted through game-theoretic lenses. For instance, Balduzzi et al. (2018) [1460] analyze gradient descent in deep learning as a dynamical system of interacting particles, drawing parallels to potential games. This perspective sheds light on the convergence behavior of neural networks, complemented by the work of Chizat & Bach (2018) [1462] on lazy training in wide networks. Meanwhile, Mertikopoulos et al. (2018) [1487] study the role of non-convexity in game-theoretic optimization, providing convergence guarantees for gradient-based methods in adversarial settings. Their work has implications for training GANs and other adversarial models, with additional theoretical grounding provided by Daskalakis et al. (2018) [1464] on the complexity of finding equilibria in non-convex games. Cooperative game theory has also found applications in explainability and feature attribution. Lundberg and Lee's (2017) [1485] SHAP values, rooted in Shapley values from cooperative game theory, provide a principled framework for interpreting model predictions. This approach has been extended by Sundararajan
et al. (2017) [1499] through integrated gradients, offering another game-theoretic perspective on feature importance. These methods have become standard tools for explainable AI, with further refinements by Janzing et al. (2020) [1477] on causal Shapley values and Frye et al. (2020) [1470] on contextual decomposition.

The intersection of game theory and deep learning continues to evolve, with recent work exploring mean-field games for large-scale multi-agent systems, as seen in the contributions of Carmona et al. (2019) [1461]. Their framework scales game-theoretic analysis to populations of agents, enabling applications in economics and robotics, with extensions by Perrin et al. (2020) [1491] on deep mean-field reinforcement learning. Another emerging direction is the use of auction theory in neural architecture search, where Liu et al. (2019) [1313] model the selection of network architectures as a combinatorial auction, optimizing resource allocation efficiently. These advancements underscore the breadth of game-theoretic applications in deep learning, from theoretical foundations to practical implementations, as seen in the work of Feng et al. (2020) [1467] on automated mechanism design. Further developments include the study of incentive alignment in federated learning, as explored by Li et al. (2020) [1482], who use mechanism design to ensure truthful participation among distributed agents. This line of work addresses challenges in collaborative learning environments, where strategic behavior must be accounted for, with additional insights from Kang et al. (2020) [1478] on reputation-based incentives. Additionally, the work of Letcher et al. (2019) [1481] on differentiable game theory provides tools for analyzing and optimizing the interactions of multiple loss functions in machine learning, offering a unified framework for multi-objective optimization, further refined by Schäfer & Anandkumar (2019) [1494] on competitive gradient descent.

The literature also delves into the role of information asymmetry in deep learning games, as exemplified by the work of Raghu et al. (2019) [1493] on signaling games in neural networks. Their analysis reveals how information bottlenecks emerge during training, influencing model behavior, with connections to the information bottleneck theory of Tishby et al. (2015) [1500]. Another significant contribution is by Perolat et al. (2022) [1489], who apply extensive-form games to hierarchical reinforcement learning, enabling sophisticated decision-making in structured environments, building on earlier work by Heinrich & Silver (2016) [1476] on fictitious self-play. Other notable contributions include the work of Czarnecki et al. (2020) [1463] on open-ended learning in games, which explores how agents can continuously adapt in dynamic environments, and the study of meta-learning in games by Al-Shedivat et al. (2018) [1458], which frames few-shot learning as a multi-agent interaction problem. The role of evolutionary dynamics in deep learning has been examined by Such et al. (2017) [1498], who apply population-based training to optimize hyperparameters, while Wang et al. (2019) [1501] explore the game-theoretic foundations of self-play in reinforcement learning.

Recent advances in differentiable economics by Dütting et al. (2019) [1466] and strategic classification by Hardt et al. (2016) [1475] further illustrate the growing synergy between game theory and deep learning. The work of Gemp et al. (2021) [1471] on resolving bargaining games with deep networks and the study of cooperative inverse reinforcement learning by Hadfield-Menell et al. (2016) [1474] highlight the potential for game-theoretic principles to guide AI alignment. Overall, the game-theoretic formulations of deep neural networks span a wide array of topics, from adversarial training and multi-agent systems to optimization and explainability. The contributions discussed here represent a fraction of the rich literature in this interdisciplinary field, which continues to grow as researchers uncover new connections between game theory and deep learning. Each of these works provides rigorous theoretical foundations, practical algorithms, or novel applications, collectively advancing our understanding of neural networks through the lens of game theory. The ongoing exploration of these ideas promises to yield further breakthroughs in the design, analysis, and deployment of intelligent systems.
## 5.2 Analysis of Game-theoretic formulations of Deep Neural Networks

Deep Neural Networks (DNNs) can be rigorously formulated within a game-theoretic framework by considering the interplay between competing objectives in optimization, adversarial robustness, and equilibrium strategies. The fundamental nature of DNNs involves a multi-agent optimization problem where various components interact within a high-dimensional function space. Formally,

<figure><img src="image_52.png" /><figcaption>Figure 5.1: Min-Max Game Formulation (Adversarial Training)</figcaption></figure>

a deep neural network is parameterized by a set of weights $\theta \in \mathbb{R}^d$ and seeks to optimize a loss function $L: \mathbb{R}^d \times \mathbb{R}^n \to \mathbb{R}$, where the training data $(x, y) \sim P(x, y)$ is drawn from an unknown distribution. The optimization problem can be expressed as

$$ \theta^* = \arg \min_{\theta} \mathbb{E}_{(x,y) \sim P}[L(f_{\theta}(x), y)]. \qquad (5.7) $$

A fundamental game-theoretic formulation of DNNs arises in adversarial training, where a deep network $f_{\theta}(x)$ is trained to minimize an objective function while an adversary seeks to maximize it. This setting leads to a two-player zero-sum game where the adversary introduces perturbations $\delta \in \mathcal{D}$ such that the resulting optimization problem is

$$ \min_{\theta} \max_{\delta \in \mathcal{D}} \mathbb{E}_{(x,y) \sim P} [L(f_{\theta}(x+\delta), y)]. \qquad (5.8) $$

The equilibrium of this adversarial game corresponds to a minimax solution, satisfying the condition

$$ \theta^* = \arg \min_{\theta} \max_{\delta \in \mathcal{D}} L(f_{\theta}(x + \delta), y). \qquad (5.9) $$

If the loss function satisfies convexity in $\theta$ and concavity in $\delta$, then the minimax theorem guarantees the existence of an equilibrium solution:

$$ \min_{\theta} \max_{\delta} L(\theta, \delta) = \max_{\delta} \min_{\theta} L(\theta, \delta). \qquad (5.10) $$
A notable game-theoretic application of DNNs is in Generative Adversarial Networks (GANs), where two neural networks, a generator $G(z)$ and a discriminator $D(x)$, play a competitive game. The generator aims to produce realistic samples $G(z) \sim P_G$, while the discriminator seeks to distinguish between real and generated samples. The objective function for this adversarial game is

$$ \min_G \max_D \mathbb{E}_{x \sim P_{\text{data}}} [\log D(x)] + \mathbb{E}_{z \sim P_z} [\log(1 - D(G(z)))] \quad (5.11) $$

The Nash equilibrium condition for this game is given by

$$ D^*(x) = \frac{P_{\text{data}}(x)}{P_{\text{data}}(x) + P_G(x)}. \quad (5.12) $$

At equilibrium, the optimal generator satisfies

$$ P_G(x) = P_{\text{data}}(x), \quad (5.13) $$

leading to a situation where the discriminator is unable to distinguish real from generated samples. If $D(x)$ is parameterized by $\phi$, then training the discriminator corresponds to solving

$$ \phi^* = \arg \max_{\phi} \mathbb{E}_{x \sim P_{\text{data}}} [\log D_{\phi}(x)] + \mathbb{E}_{z \sim P_z} [\log(1 - D_{\phi}(G_{\theta}(z)))] \quad (5.14) $$

Given that neural networks can approximate arbitrary functions, their training can also be formulated in terms of multi-agent differential game theory. Consider a dynamic system where the

<figure><img src="image_53.png" /><figcaption>Figure 5.2: Regularization as Game Formulation</figcaption></figure>

evolution of parameters follows

$$ \frac{d\theta}{dt} = -\nabla_{\theta} \mathcal{L}(\theta, \phi), \quad (5.15) $$

$$ \frac{d\phi}{dt} = \nabla_{\phi} \mathcal{L}(\theta, \phi). \quad (5.16) $$
The Hamiltonian for this system is given by

$$H(\theta, p) = \min_{u} [\nabla_{\theta}\mathcal{L} \cdot u + p \cdot f(\theta, u)], \qquad (5.17)$$

where $p$ is the costate variable and $u$ represents control parameters. The Hamilton-Jacobi-Bellman equation characterizing optimal control in a reinforcement learning-based DNN training scenario is

$$\frac{\partial V}{\partial t} + \min_{u} [\nabla_{\theta}\mathcal{L} \cdot u + p \cdot f(\theta, u)] = 0. \qquad (5.18)$$

Another approach to studying DNNs in a game-theoretic framework is through evolutionary game dynamics. Here, learning dynamics follow replicator equations given by

$$\frac{dx_i}{dt} = x_i(f_i - \bar{f}), \qquad (5.19)$$

where $x_i$ represents the probability of selecting strategy $i$, $f_i$ is the fitness function for strategy $i$, and $\bar{f} = \sum_j x_j f_j$ is the average fitness. In the context of neural networks, the replicator equation governs weight updates as follows:

$$\frac{d\theta_i}{dt} = \theta_i \left( \nabla_{\theta_i} J - \frac{1}{n} \sum_j \nabla_{\theta_j} J \right). \qquad (5.20)$$

This provides a mechanism for gradient-based optimization that incorporates selection and mutation principles. Another critical formulation involves Nash equilibria in non-cooperative learning

<figure><img src="image_54.png" /><figcaption>Figure 5.3: Multi-Agent Game: Layers as Players</figcaption></figure>

settings, such as federated learning, where multiple learners optimize separate but interdependent loss functions. The Nash equilibrium conditions require that

$$\theta_i^* = \arg \min_{\theta_i} L_i(\theta_i, \theta_{-i}), \qquad (5.21)$$
where $\theta_{-i}$ denotes the set of parameters for all players except $i$. A Nash equilibrium is attained when

$$ \frac{\partial L_i}{\partial \theta_i} = 0, \quad \forall i. \qquad (5.22) $$

Under convexity assumptions, Nash learning dynamics can be characterized by a gradient projection algorithm:

$$ \theta^{(t+1)} = \theta^{(t)} - \eta \nabla_{\theta} L(\theta^{(t)}). \qquad (5.23) $$

Beyond Nash equilibria, deep neural networks can also be analyzed through variational inequalities, which generalize equilibrium concepts. Given a function $F: \mathbb{R}^d \to \mathbb{R}^d$, the variational inequality problem seeks $\theta^*$ such that

$$ \langle F(\theta^*), \theta - \theta^* \rangle \geq 0, \quad \forall \theta \in \mathbb{R}^d. \qquad (5.24) $$

For adversarial training, the corresponding variational inequality formulation is

$$ \langle \nabla_{\theta} L(\theta, \delta^*(\theta)), \theta - \theta^* \rangle \geq 0. \qquad (5.25) $$

where $\delta^*(\theta)$ denotes the optimal adversarial perturbation satisfying

$$ \delta^*(\theta) = \arg \max_{\delta \in \mathcal{D}} L(\theta, \delta). \qquad (5.26) $$

This variational inequality provides a principled characterization of adversarial robustness in deep networks. The application of game-theoretic formulations to deep learning thus allows for a mathematically rigorous understanding of optimization, equilibrium conditions, and strategic interactions among network components.

## 5.2.1 Python Code to generate Figures 5.1, 5.2, and 5.3 illustrating Game-theoretic formulations of Deep Neural Networks

The Python code below produces the Figures 5.1, 5.2, and 5.3 illustrating Game-theoretic formulations of Deep Neural Networks.

```python
import numpy as np
import matplotlib.pyplot as plt

# --- 1. Min-Max (Adversarial Training Game) ---
epsilon = np.linspace(0, 1, 100)
empirical_loss = 0.2 + 0.3 * epsilon  # increasing with epsilon
robust_loss = empirical_loss + 0.1 * np.sin(3*epsilon)  # adversarial perturbation

fig, ax = plt.subplots(figsize=(6,4))
ax.set(xlim=(0, 1), ylim=(-0.5, 0.5))
ax.set(xlabel="Perturbation Size ($\epsilon$)", ylabel="Loss")
ax.set(xlim=(0, 1), ylim=(-0.5, 0.5))
ax.set(xlabel="Min-Max Game Formulation (Adversarial Training)", ylabel="Adversarial Loss")
ax.legend()
ax.grid(True)
ax.show()

# --- 2. Regularization as Game (Player: Model vs Regularizer) ---
weights = np.linspace(-3, 3, 100)
loss = (weights - 1)**2
reg = 0.5 * weights**2
```
23 total = loss + reg
24
25 plt.figure(figsize=(6,4))
26 plt.plot weights, loss, label="Task Loss")
27 plt.plot weights, reg, label="Regularization")
28 plt plot weights, total, label="Total Objective")
29 plt xlabel("Weight Parameter")
30 plt ylabel("Loss Value")
31 plt title("Regularization as Game Formulation")
32 plt.legend()
33 plt grid(True)
34 plt show()
35
36 # --- 3. Multi-Agent View (Layers as Players) ---
37 epochs = np arange 1, 11)
38 layer1 = np exp(-0.2 * epochs) + 0.05* np random randn 10)
39 layer2 = np exp(-0.15 * epochs) + 0.05* np random randn 10)
40 layer3 = np exp(-0.1 * epochs) + 0.05* np random randn 10)
41
42 plt figure(figsize=(6,4))
43 plt plot epochs, layer1, 'o-,' label="Layer 1 Loss")
44 plt plot epochs, layer2, 's-,' label="Layer 2 Loss")
45 plt plot epochs, layer3, '^-,' label="Layer 3 Loss")
46 plt xlabel("Epoch")
47 plt ylabel("Layer Objective")
48 plt title("Multi-Agent Game: Layers as Players")
49 plt legend()
50 plt grid(True)
51 plt show()

5.2.2 Min-Max (Saddle) dynamics: $L(x, y) = x * y$

The function that serves as a canonical, analytically tractable model for studying the min-max
saddle point dynamics that is central to the training of Generative Adversarial Networks (GANs)
and other adversarial learning formulations within deep learning is

$$
L(x, y) = x \cdot y \tag{5.27}
$$

In this game-theoretic framework, the two players are the generator and the discriminator, whose
parameters are abstracted here by the variables *y* and *x*, respectively. The discriminator’s objective
is to maximize the payoff function *L*(*x*, *y*), while the generator’s objective is to minimize it. This
defines a zero-sum game with the value function

$$
\phi(x) = \max_{y} L(x, y) \tag{5.28}
$$

and the overall goal is to find a Nash equilibrium, which is a point $(x^*, y^*)$ satisfying the saddle
point condition

$$
L(x^*, y) \leq L(x^*, y^*) \leq L(x, y^*) \tag{5.29}
$$

for all *x* and *y* in a local neighborhood.

The dynamics of the training process are typically modeled by simultaneous or alternating gra-
dient descent-ascent. The vector field governing these dynamics is given by the gradients of the
<figure><img src="image_55.png" /><figcaption>Figure 5.4: Min-Max (Saddle) dynamics: $L(x, y) = x * y$</figcaption></figure>

objective with respect to each player's variable. For the bilinear objective $L(x, y) = x \cdot y$, these gradients are

$$ \nabla_x L(x, y) = y \quad \text{and} \quad \nabla_y L(x, y) = x \qquad (5.30) $$

The simultaneous gradient descent-ascent update rule with learning rate $\eta$ is

$$ x_{t+1} = x_t + \eta \cdot y_t \qquad (5.31) $$

and

$$ y_{t+1} = y_t - \eta \cdot x_t \qquad (5.32) $$

This system of equations is linear and can be written in matrix form as

$$ \mathbf{z}_{t+1} = A\mathbf{z}_t \quad \text{where} \quad \mathbf{z}_t = (x_t, y_t)^\top \qquad (5.33) $$

and

$$ A = \begin{bmatrix} 1 & \eta \\ -\eta & 1 \end{bmatrix} \qquad (5.34) $$

The eigenvalues of this update matrix are $\lambda = 1 \pm i\eta$, which have magnitude $\sqrt{1 + \eta^2} > 1$. This implies that the dynamics are unstable and will diverge for any $\eta > 0$, causing the players' parameters to orbit the equilibrium point $(0, 0)$ on a trajectory of exponentially increasing radius. This simple model thus captures a fundamental pathology in adversarial training: the vector field

$$ \omega(z) = (\nabla_x L, -\nabla_y L) = (y, -x) \qquad (5.35) $$

is non-conservative and has a non-zero rotational component, preventing convergence under naive gradient-based methods.

The divergence of this simple dynamic motivates the need for more sophisticated optimization algorithms in practical deep neural network training. The bilinear term $x \cdot y$ is often a component
of a larger loss landscape, and its presence induces rotational forces that can destabilize training. Algorithms such as Optimistic Mirror Descent (OMD) and the closely related Extra Gradient (EG) method are designed to counteract this instability. For example, the Optimistic Mirror Descent update for this problem is given by

$$x_{t+1} = x_t + 2\eta y_t - \eta y_{t-1} \qquad (5.36)$$

and

$$y_{t+1} = y_t - 2\eta x_t + \eta x_{t-1} \qquad (5.37)$$

The spectral analysis of this system reveals that it can converge for a sufficiently small learning rate $\eta$, unlike the divergent simultaneous gradient method. The core mathematical insight is that the stability of the learning dynamics depends critically on the structure of the game’s Jacobian,

$$J = \begin{bmatrix} 0 & \nabla_{xy}^2 L \\ -(\nabla_{yx}^2 L)^\top & 0 \end{bmatrix} = \begin{bmatrix} 0 & 1 \\ -1 & 0 \end{bmatrix} \qquad (5.38)$$

For the bilinear game, this Jacobian is a skew-symmetric matrix, and convergence guarantees for algorithms like OMD and EG are derived by showing that their update operators can mitigate the rotational effects encoded by this skew-symmetric component, thereby allowing the dynamics to spiral inward towards the saddle point $x^* = 0, y^* = 0$ rather than outward to infinity.

### 5.2.3 Potential/Cooperative game (gradient descent on shared Phi)

In the game-theoretic analysis of deep neural network training, a potential or cooperative game represents a foundational scenario where the interests of all parameter blocks are perfectly aligned. This framework is defined by the existence of a global, differentiable potential function $\Phi(\theta)$, where $\theta = (\theta_1, \theta_2, ..., \theta_N)$ represents the concatenated parameters of all $N$ players or network components. The individual loss function for each player $i$ is precisely this shared potential

$$\mathcal{L}_i(\theta) = \Phi(\theta) \qquad (5.39)$$

for all $i$. Consequently, the game is one of pure cooperation, as all players share the identical objective of minimizing $\Phi(\theta)$. The dynamics are governed by the simultaneous gradient, which is simply the full gradient of the potential function:

$$\xi(\theta) = (\nabla_{\theta_1} \mathcal{L}_1, \nabla_{\theta_2} \mathcal{L}_2, ..., \nabla_{\theta_N} \mathcal{L}_N) = (\nabla_{\theta_1} \Phi, \nabla_{\theta_2} \Phi, ..., \nabla_{\theta_N} \Phi) = \nabla_{\theta} \Phi \qquad (5.40)$$

The learning process in this setting is equivalent to gradient descent on the shared objective $\Phi$. The continuous-time dynamics are described by the ordinary differential equation

$$\dot{\theta} = -\eta \nabla_{\theta} \Phi(\theta) \qquad (5.41)$$

where $\eta > 0$ is the learning rate. The discrete-time update with a fixed learning rate $\eta$ is given by

$$\theta_{t+1} = \theta_t - \eta \nabla_{\theta} \Phi(\theta_t) \qquad (5.42)$$

A first-order Taylor expansion reveals the change in the potential function:

$$\Phi(\theta_{t+1}) = \Phi(\theta_t - \eta \nabla_{\theta} \Phi(\theta_t)) \approx \Phi(\theta_t) - \eta \|\nabla_{\theta} \Phi(\theta_t)\|^2 + \frac{\eta^2}{2} \nabla_{\theta} \Phi(\theta_t)^{\top} \mathbf{H}(\theta_t) \nabla_{\theta} \Phi(\theta_t) \qquad (5.43)$$

where

$$\mathbf{H}(\theta) = \nabla_{\theta}^{2} \Phi(\theta) \qquad (5.44)$$

is the Hessian matrix of the potential. For sufficiently small $\eta$ such that

$$\eta \mathbf{I} \prec \mathbf{H}(\theta_t)^{-1} \qquad (5.45)$$
<figure><img src="image_56.png" /><figcaption>Figure 5.5: Potential / Cooperative game (gradient descent on shared Phi)</figcaption></figure>

in the positive definite sense, the term linear in $\eta$ dominates, guaranteeing a decrease in the potential:

$$ \Phi(\theta_{t+1}) < \Phi(\theta_t) \qquad (5.46) $$

The game-theoretic properties are analyzed through the lens of the game Jacobian, $\mathbf{J}(\theta)$, which is defined as the matrix of second-order derivatives of the individual losses. For a potential game, this Jacobian is given by

$$ \mathbf{J}_{ij}(\theta) = \nabla_{\theta_j} \nabla_{\theta_i} \mathcal{L}_i(\theta) = \nabla_{\theta_j} \nabla_{\theta_i} \Phi(\theta) \qquad (5.47) $$

This implies that $\mathbf{J}(\theta)$ is equivalent to the Hessian of the potential function, $\mathbf{H}(\theta)$. A critical mathematical consequence is that the Jacobian of a potential game is symmetric

$$ \mathbf{J}(\theta)^{\top} = \mathbf{J}(\theta) \qquad (5.48) $$

This symmetry is the defining characteristic of a potential game and ensures that the dynamics are a gradient flow. The eigenvalues of $\mathbf{J}(\theta)$ are all real, and their signs determine the stability of critical points. A point $\theta^*$ is a Nash equilibrium if

$$ \nabla_{\theta} \Phi(\theta^*) = 0 \qquad (5.49) $$

and it is a locally stable equilibrium under gradient descent if

$$ \mathbf{H}(\theta^*) \succ 0 \qquad (5.50) $$

meaning all eigenvalues of the Hessian are positive.

The convergence analysis leverages Lyapunov stability theory. The potential function $\Phi(\theta)$ itself serves as a Lyapunov function for the gradient dynamics. Its time derivative along the trajectories of the continuous-time system is

$$ \frac{d}{dt} \Phi(\theta(t)) = \nabla_{\theta} \Phi(\theta)^{\top} \dot{\theta} = -\eta \| \nabla_{\theta} \Phi(\theta) \|^2 \le 0 \qquad (5.51) $$
This quantity is strictly negative everywhere except at critical points where $\nabla_{\theta}\Phi(\theta) = 0$, confirming the global stability of the dynamics. In the discrete-time setting, under the assumption that $\Phi$ is $L$-smooth, meaning

$$ \|\nabla_{\theta}\Phi(\theta) - \nabla_{\theta}\Phi(\phi)\| \le L\|\theta - \phi\| \qquad (5.52) $$

which implies

$$ \Phi(\phi) \le \Phi(\theta) + \nabla_{\theta}\Phi(\theta)^{\top}(\phi - \theta) + \frac{L}{2}\|\phi - \theta\|^2 \qquad (5.53) $$

one can derive a quantitative convergence rate. Applying this at

$$ \phi = \theta_{t+1} = \theta_t - \eta \nabla_{\theta}\Phi(\theta_t) \qquad (5.54) $$

yields

$$ \Phi(\theta_{t+1}) \le \Phi(\theta_t) - \eta \|\nabla_{\theta}\Phi(\theta_t)\|^2 + \frac{L\eta^2}{2} \|\nabla_{\theta}\Phi(\theta_t)\|^2 = \Phi(\theta_t) - \eta(1 - \frac{L\eta}{2}) \|\nabla_{\theta}\Phi(\theta_t)\|^2 \qquad (5.55) $$

For a learning rate $\eta < \frac{2}{L}$, this guarantees a monotonic decrease in the potential. After $T$ iterations, the best iterate satisfies

$$ \min_{0 \le t < T} \|\nabla_{\theta}\Phi(\theta_t)\|^2 \le \frac{2(\Phi(\theta_0) - \Phi^*)}{\eta(2 - L\eta)T} \qquad (5.56) $$

where $\Phi^*$ is the minimal value, proving convergence to a stationary point at a rate of $O(1/T)$. This mathematically rigorous framework provides the foundational guarantees for standard, non-adversarial deep learning, where the optimization landscape, despite its non-convexity, is traversed by a single, cohesive gradient flow.

### 5.2.4 Replicator dynamics (Rock-Paper-Scissors) — cyclic behavior

The Replicator Dynamics equation, a cornerstone of evolutionary game theory, provides a continuous-time model for the evolution of strategy frequencies within a population. For a game with $n$ pure strategies and a payoff matrix $A$, where $A_{ij}$ denotes the payoff for strategy $i$ against strategy $j$, the dynamics are defined for a population state $x \in \Delta^n$, the $(n-1)$-dimensional simplex where $x_i \ge 0$ and $\sum_{i=1}^n x_i = 1$. The rate of change of the frequency $x_i$ is given by

$$ \dot{x}_i = x_i ((Ax)_i - x^\top Ax) \qquad (5.57) $$

The term

$$ (Ax)_i = \sum_{j=1}^{n} A_{ij}x_j \qquad (5.58) $$

is the expected payoff for strategy $i$, and

$$ x^\top Ax = \sum_{i=1}^{n} \sum_{j=1}^{n} x_i A_{ij} x_j \qquad (5.59) $$

is the average payoff across the entire population. This equation dictates that strategies yielding above-average payoff proliferate, while those with below-average payoff diminish. The Rock-Paper-Scissors (RPS) game is the canonical example of a zero-sum, non-transitive game that induces cyclic behavior under these dynamics. Its payoff matrix for the row player is given by

$$ A = \begin{bmatrix} 0 & -1 & 1 \\ 1 & 0 & -1 \\ -1 & 1 & 0 \end{bmatrix} \qquad (5.60) $$
<figure><img src="image_57.png" /><figcaption>Figure 5.6: Replicator dynamics (Rock-Paper-Scissors) — cyclic behavior</figcaption></figure>

The average payoff in the population for this antisymmetric matrix ($A = -A^T$) is always zero, $x^T Ax = 0$, simplifying the replicator dynamics to $\dot{x}_i = x_i(Ax)_i$. Writing these equations explicitly for the three strategies (Rock R, Paper P, Scissors S) yields the system:

$$ \dot{x}_R = x_R((Ax)_R) = x_R(0 \cdot x_R + (-1) \cdot x_P + 1 \cdot x_S) = x_R(x_S - x_P) \quad (5.61) $$

$$ \dot{x}_P = x_P((Ax)_P) = x_P(1 \cdot x_R + 0 \cdot x_P + (-1) \cdot x_S) = x_P(x_R - x_S) \quad (5.62) $$

$$ \dot{x}_S = x_S((Ax)_S) = x_S((-1) \cdot x_R + 1 \cdot x_P + 0 \cdot x_S) = x_S(x_P - x_R) \quad (5.63) $$

The dynamics of this system conserve a specific quantity, which can be found by analyzing the time derivative of the logarithm of the product $x_R x_P x_S$. Specifically,

$$ \frac{d}{dt} \ln(x_R x_P x_S) = \frac{\dot{x}_R}{x_R} + \frac{\dot{x}_P}{x_P} + \frac{\dot{x}_S}{x_S} = (x_S - x_P) + (x_R - x_S) + (x_P - x_R) = 0 \quad (5.64) $$

This implies that the function

$$ V(x) = \ln(x_R) + \ln(x_P) + \ln(x_S) \quad (5.65) $$

is a constant of motion, and consequently, the quantity

$$ K = x_R x_P x_S \quad (5.66) $$

remains constant along any trajectory of the dynamics. The orbits of the system are closed curves defined by the level sets of K, resulting in perpetual, neutrally stable cycles around the interior Nash equilibrium $x^* = (\frac{1}{3}, \frac{1}{3}, \frac{1}{3})$.

This cyclic behavior is not merely a pathological example but a fundamental model for understanding non-convergent dynamics in multi-agent learning systems, including certain regimes in the training of deep neural networks, particularly Generative Adversarial Networks (GANs). The vector field $F(x) = (\dot{x}_R, \dot{x}_P, \dot{x}_S)$ defines a flow on the simplex. The Jacobian of this vector field at the interior equilibrium $x^*$ is

$$ J = \begin{bmatrix} 0 & -\frac{1}{3} & \frac{1}{3} \\ \frac{1}{3} & 0 & -\frac{1}{3} \\ -\frac{1}{3} & \frac{1}{3} & 0 \end{bmatrix} \quad (5.67) $$
This matrix is skew-symmetric, $J = -J^\top$, and its eigenvalues are $\lambda = 0, \pm i \frac{1}{\sqrt{3}}$. The purely imaginary conjugate pair $\pm i\omega$ confirm the presence of a center manifold, leading to oscillatory behavior. The connection to deep learning is established by noting that the simultaneous gradient descent-ascent dynamics on a bilinear loss function

$$L(\theta, \phi) = \theta^\top A \phi \qquad (5.68)$$

which abstracts the interaction between two neural networks like a generator and a discriminator, yield a structurally similar vector field

$$(\dot{\theta}, \dot{\phi}) = (-\nabla_{\theta}L, -(-\nabla_{\phi}L)) = (-A\phi, A^{\top}\theta) \qquad (5.69)$$

For an antisymmetric $A$, this becomes

$$\dot{\theta} = -A\phi \quad \text{and} \quad \dot{\phi} = A\theta \qquad (5.70)$$

whose second-order form

$$\ddot{\theta} = -AA^{\top}\theta \qquad (5.71)$$

describes a multi-dimensional harmonic oscillator. This mathematical equivalence demonstrates how the replicator dynamics' inherent cycles manifest as rotational forces and training instability in adversarial neural network training, preventing convergence to a stationary Nash equilibrium.

### 5.2.5 Python Code to generate Figures 5.4, 5.5 and 5.6 illustrating Min-Max (Saddle) dynamics, Potential/Cooperative game (gradient descent on shared Phi), and Replicator dynamics (Rock-Paper-Scissors) — cyclic behavior

The Python code below produces the Figures 5.4, 5.5 and 5.6 illustrating Min-Max (Saddle) dynamics, Potential/Cooperative game (gradient descent on shared Phi), and Replicator dynamics (Rock-Paper-Scissors) — cyclic behavior.

```python
"""
Game-theoretic visualizations for DNN-style games:
    - Saddle (min-max) game (toy 2D)
    - Potential/cooperative game (toy 2D)
    - Replicator dynamics for Rock-Paper-Scissors (3-player simplex cycles)

Run in a Jupyter notebook or plain Python. Requires numpy, matplotlib.
"""
import numpy as np
import matplotlib.pyplot as plt

# ------------------------
# 1) Min-Max (saddle) game
#    Example objective: L(x,y) = x * y (simple bilinear saddle)
#    Player X minimizes in x; Player Y maximizes in y
# ------------------------

def saddle_loss(x, y):
    return x * y

def saddleGrad(x, y):
    # grad_x L = y ; grad_y L = x
    return y, x
```
```python
def simulateGradientDescentAscent(x0, y0, lr_x, lr_y, steps=200):
    x, y = x0, y0
    traj = [(x, y)]

    for _ in range(steps):
        gx, gy = saddleGrad(x, y)
        # X does descent (minimize): x <- x - lr_x * grad_x
        x = x - lr_x * gx
        # Y does ascent (maximize): y <- y + lr_y * grad_y
        y = y + lr_y * gy
        traj.append((x, y))
    return np.array(traj)
```

```python
# -----------------------------
# 2) Potential (cooperative) game
#   Shared objective: Phi(x,y) = 0.5*(x - a)^2 + 0.5*(y - b)^2
#   Both players minimize same objective -> gradient flow toward (a,b)
# -----------------------------

a, b = 1.0, -0.5
def phi(x, y):
    return 0.5*(x - a)**2 + 0.5*(y - b)**2
```

```python
def phiGrad(x, y):
    # grad_x Phi = (x - a) ; grad_y = (y - b)
    return (x - a), (y - b)
```

```python
def simulateCooperative(x0, y0, lr=0.05, steps=200):
    x, y = x0, y0
    traj = [(x,y)]

    for _ in range(steps):
        gx, gy = phiGrad(x, y)
        x = x - lr * gx
        y = y - lr * gy
        traj.append((x,y))
    return np.array(traj)
```

```python
# -----------------------------
# 3) Replicator dynamics for Rock-Paper-Scissors
#   Payoff matrix (zero-sum cyclic):
#     A = [[0, -1, 1],
#           [1,  0, -1],
#           [-1,  1,  0]]
#   Replicator ODE: p_i' = p_i * ((A p)_i - p^T A p)
# -----------------------------

A = np.array([[0, -1, 1],
              [1,  0, -1],
              [-1,  1,  0]], dtypeanteer)
```

```python
def replicatorStep(p, dt=0.01):
    Ap = A.dot(p)
    phi = p.dot(Ap)
    dp = p * (Ap - phi)
    p_next = p + dt * dp

    # project to simplex and clip
    p_next = np.clip(p_next, 1e-12, None)
    p_next = p_next / p_next.sum()
```
132 lr = 0.08
133 for i in range(X.shape[0]):
134 for j in range(X.shape[1]):
135     gx, gy = phiGrad(X[i,j], Y[i,j])
136     U[i,j] = -lr * gx
137     V[i,j] = -lr * gy
138
139 plt streamline(X, Y, U, V, density=1.2)
140 plt.title("Potential / Cooperative game (gradient descent on shared Phi)")
141 plt.xlabel("x")
142 plt.ylabel("y")
143 for p0 in [(-2.0, 1.5), (2.5, -2.5), (0.0, 2.0), (3.0, 1.0)]:
144     traj = simulate_cooperative(p0[0], p0[1], lr=lr, steps=300)
145     plt plot(traj[:,0], traj[:,1],",-", lw=1.5)
146     plt plot(traj[0,0], traj[0,1],",o")
147 # mark the potential minimizer
148 plt plot(a, b,",k*", markersize=10, label='a,b target')
149 plt legend()
150 plt xlim(-2.5, 3.5);
151 plt ylim(-3.0, 2.0)
152
153
154 plt figure(figsize=(6,4))
155 # 3: Replicator dynamics (RPS) on simplex (2D projection)
156 #ax3 = fig.addsubplot(1, 3, 3)
157 # create barycentric coordinates for simplex plotting
158 def barycentric coords:
159     #协调: (n_points,3) -> (x,y)
160     # vertices: (0,0), (1,0), (0.5, sqrt(3)/2)
161     v0 = np.array([0.0, 0.0])
162     v1 = np.array([1.0, 0.0])
163     v2 = np.array([0.5, npsqrt(3)/2.0])
164     return coordinates[::,0, None]*v0 + coordinates[::,1, None]*v1 + coordinates[::,2, None]*v2
165
166 # sample initial points and simulate
167 inits = [np.array([0.9, 0.05, 0.05]), np.array([0.1, 0.8, 0.1]), np.
168         array([0.3, 0.3, 0.4])]
169 for p0 in inits:
170     traj = simulate_replicator(p0, steps=400, dt=0.06)
171     xy = barycentric(traj)
172     plt plot(xy[:0], xy[:1], ",-")
173     plt plot(xy[0,0], xy[0,1], ",o")
174
175 # draw triangle edges
176 verts = np.array([[0.0,0.0],[1.0,0.0],[0.5 npsqrt(3)/2.0],[0.0,0.0]])
177 plt plot(verts[:0], verts[:1], ",k-")
178 plt title("Replicator dynamics (Rock-Paper-Scissors) - cyclic behavior")
179 plt xticks [];
180 plt yticks []
181 plt show()
## 5.3 Game-theoretic formulations of deep neural networks (DNNs) through Evolutionary Game Dynamics

Deep neural networks (DNNs) can be formulated within the framework of evolutionary game dynamics, where the optimization of network parameters is recast as a competitive evolutionary process in which different parameter configurations act as strategies that undergo selection, mutation, and adaptation over time. This perspective allows a rigorous characterization of learning dynamics through replicator equations, Nash equilibria, and mutation-selection balance, thereby linking deep learning to fundamental principles in game theory and statistical physics. Given a neural network parameterized by $\theta \in \Theta$, where $\Theta$ is the space of all possible weight configurations, the performance of the network is quantified by a fitness function $F(\theta)$, which is typically defined as the negative of the empirical loss function $L(\theta)$. This loss function can be expressed as an expectation over a given dataset $D$,

$$L(\theta) = \mathbb{E}_{(x,y)\sim D}[\ell(f_{\theta}(x), y)], \qquad (5.72)$$

where $\ell$ represents the pointwise loss function, such as the squared error loss or cross-entropy loss. The fitness of a given parameter configuration $\theta$ is then defined as

$$F(\theta) = -L(\theta), \qquad (5.73)$$

which ensures that lower loss corresponds to higher fitness. The evolutionary distribution of network parameters, denoted by $p(\theta, t)$, follows an evolutionary replicator-mutator equation,

$$\frac{dp(\theta, t)}{dt} = p(\theta, t)(F(\theta) - \bar{F}) + \sigma^2 \nabla^2 p(\theta, t), \qquad (5.74)$$

where $\bar{F}$ is the mean fitness across all parameter configurations,

$$\bar{F} = \int_{\Theta} p(\theta, t) F(\theta) d\theta. \qquad (5.75)$$

The first term in the equation represents selection dynamics, where parameter configurations with higher fitness proliferate, while the second term models mutations, accounting for stochastic noise in learning updates. The steady-state solution to this equation corresponds to an evolutionarily stable distribution $p^*(\theta)$, which satisfies

$$\mathbb{E}_{\theta' \sim p^*(\theta)}[F(\theta')] > \mathbb{E}_{\theta' \sim p(\theta)}[F(\theta')] \qquad (5.76)$$

for any small perturbation $p(\theta)$ from $p^*(\theta)$. The learning process in deep networks can thus be viewed as an evolutionary competition in which weight distributions evolve toward stable equilibria or limit cycles. The standard stochastic gradient descent (SGD) update rule, given by

$$\theta^{(t+1)} = \theta^{(t)} - \eta \nabla_{\theta} L(\theta), \qquad (5.77)$$

can be reinterpreted as a mean-field approximation of the replicator equation with stochastic diffusion,

$$\frac{d\theta}{dt} = -\nabla_{\theta} L(\theta) + \sigma\xi(t), \qquad (5.78)$$

where $\xi(t)$ represents a stochastic perturbation modeling noise in the gradient updates. In the limit of infinitesimally small learning rate $\eta$, the SGD update can be seen as an Ornstein-Uhlenbeck process,

$$d\theta = -\nabla_{\theta} L(\theta)dt + \sigma dW_t, \qquad (5.79)$$
<figure><img src="image_58.png" /><figcaption>Figure 5.7: Replicator Dynamics as Gradient Descent</figcaption></figure>

where $W_t$ denotes a Wiener process. The equilibrium distribution of network parameters is then governed by the Fokker-Planck equation,

$$ \frac{\partial p(\theta, t)}{\partial t} = -\nabla_{\theta} \cdot (p(\theta, t) \nabla_{\theta} L(\theta)) + \frac{\sigma^2}{2} \nabla^2 p(\theta, t). \quad (5.80) $$

This equation describes the evolution of the probability density of weight configurations under a combination of gradient descent and stochastic exploration, which can be interpreted as a mutation-selection process in an evolutionary game. The steady-state solution $p^*(\theta)$ corresponds to a Boltzmann-Gibbs distribution,

$$ p^*(\theta) = \frac{e^{-\beta L(\theta)}}{Z}, \quad (5.81) $$

where $\beta \propto 1/\sigma^2$ acts as an inverse temperature parameter, and $Z$ is the partition function ensuring normalization. This establishes a direct connection between deep learning optimization and principles of statistical mechanics, whereby learning dynamics resemble an energy minimization process subject to thermal noise. In the case of generative adversarial networks (GANs), the evolutionary perspective is particularly insightful. A GAN consists of a generator $G$ and a discriminator $D$, which play a two-player zero-sum game governed by the minimax objective,

$$ \min_G \max_D \mathbb{E}_{x \sim P_{\text{data}}} [\log D(x)] + \mathbb{E}_{z \sim P_z} [\log(1 - D(G(z)))] \quad (5.82) $$

The training dynamics of GANs exhibit co-evolutionary behavior, which can be captured by the coupled replicator equations,

$$ \frac{dp_G(\theta_G, t)}{dt} = p_G(\theta_G, t)(F_G(\theta_G) - \bar{F}_G) + \sigma^2 \nabla^2 p_G(\theta_G, t), \quad (5.83) $$
<figure><img src="image_59.png" /><figcaption>Figure 5.8: Min-Max Dynamics (Adversarial Training)</figcaption></figure>

$$ \frac{dp_D(\theta_D, t)}{dt} = p_D(\theta_D, t)(F_D(\theta_D) - \bar{F}_D) + \sigma^2 \nabla^2 p_D(\theta_D, t). \qquad (5.84) $$

Unlike standard deep learning, GAN training often fails to converge to an equilibrium due to cycling behavior. This phenomenon can be analyzed through Hamiltonian game dynamics,

$$ \frac{d\theta_G}{dt} = -\nabla_{\theta_G} L_G(\theta_G, \theta_D), \qquad (5.85) $$

$$ \frac{d\theta_D}{dt} = \nabla_{\theta_D} L_D(\theta_G, \theta_D). \qquad (5.86) $$

These equations describe a conservative dynamical system, where learning trajectories follow closed orbits rather than converging to fixed points. This behavior is characteristic of zero-sum games and is mathematically analogous to the Red Queen effect in evolutionary biology, where competing species continually adapt without achieving a stable equilibrium.
<figure><img src="image_60.png" /><figcaption>Figure 5.9: Cooperative Game (Shared Potential)</figcaption></figure>

5.3.1 Python Code to generate Figures 5.7, 5.8 and 5.9 illustrating Game-theoretic formulations of deep neural networks (DNNs) through Evolutionary Game Dynamics

The Python code below produces the Figures 5.7, 5.8 and 5.9 illustrating Game-theoretic formulations of deep neural networks (DNNs) through Evolutionary Game Dynamics.

```python
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

# --- 1. Replicator Dynamics for "weights" ---
A = np.array([[0, -1, 1],
              [1, 0, -1],
              [-1, 1, 0]])  # Rock-Paper-Scissors style payoff

def replicator_dynamics(t, w):
    w = np maximum(w, 1e-8)
    w = w / np.sum(w)
```
13 payoff = A @ w
14 avg payoff = w @ payoff
15 return w * (payoff - avg payoff)
16
17 w0 = np.array([0.3, 0.4, 0.3])
18 sol = solve_ivp(replicator_dynamics, (0, 60), w0, t_eval = np. linspace(0,
19 60, 1000))
20
21 plt. figure (figsize=(7,5))
22 plt. plot (sol.t, sol.y[0], label="Weight 1")
23 plt. plot (sol.t, sol.y[1], label="Weight 2")
24 plt. plot (sol.t, sol.y[2], label="Weight 3")
25
26 plt. title ("Replicator Dynamics as Gradient Descent")
27
28
29
30
31
32 # --- 2. Min-Max Game (Adversarial Training Dynamics) ---
33 def saddle_dynamics (t, z):
34     w, delta = z
dw = -delta  # descent in w
36 ddelta = w   # ascent in delta
37 return [dw, ddelta]
38
39 z0 = [1.0, 0.5]
40 sol = solve_ivp (saddle_dynamics, (0, 30), z0, t_eval = np. linspace(0,
41 1000))
42
43
44
45
46
47
48
49
50
51 # --- 3. Cooperative Game (Potential Dynamics across layers) ---
52 def potential_dynamics (t, z):
53     x, y = z
54     return [-x, -y]  # gradient descent on Phi = (x^2+y^2)/2
55
56 z0 = [2.0, -1.0]
57 sol = solve_ivp (potential_dynamics, (0, 10), z0, t_eval = np. lerp space(0,
58 10, 300))
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
## 5.4 Analysis of Deep neural networks (DNNs) through Variational Inequalities

Deep neural networks (DNNs) can be rigorously analyzed through the mathematical framework of variational inequalities, which provide a fundamental structure for understanding the equilibrium properties of optimization processes underlying deep learning. Variational inequalities (VIs) generalize optimization problems and offer a broader perspective on stability, convergence, and generalization in high-dimensional nonconvex settings. The problem of training a deep neural network can be formulated as a variational inequality by considering the minimization of a loss function $L: \mathbb{R}^d \to \mathbb{R}$, where the goal is to determine the optimal parameters $\theta^*$ satisfying the condition

$$ \langle \nabla L(\theta^*), \theta - \theta^* \rangle \geq 0, \quad \forall \theta \in \mathbb{R}^d. \qquad (5.87) $$

This inequality characterizes the equilibrium conditions in the optimization landscape of deep learning. It expresses the fact that at the optimal point $\theta^*$, any infinitesimal movement in the parameter space does not yield a lower value of the loss function. The gradient descent update rule, given by

$$ \theta^{(k+1)} = \theta^{(k)} - \eta \nabla L(\theta^{(k)}), \qquad (5.88) $$

where $\eta > 0$ is the learning rate, can be rewritten as a **projection-based variational inequality** if constraints exist. Specifically, if the parameter space is constrained to a closed and convex set $\Theta \subset \mathbb{R}^d$, then the variational inequality formulation of the training process is

$$ \langle \nabla L(\theta^*), \theta - \theta^* \rangle \geq 0, \quad \forall \theta \in \Theta. \qquad (5.89) $$

This variational inequality expresses the fact that $\theta^*$ is a stationary point of the constrained optimization problem, meaning that the projected gradient update

$$ \theta^{(k+1)} = \mathrm{Proj}_{\Theta} (\theta^{(k)} - \eta \nabla L(\theta^{(k)})), \qquad (5.90) $$

ensures that the sequence $\{\theta^{(k)}\}$ converges to the solution of the variational inequality under suitable assumptions on the loss function. A key property in variational inequalities is monotonicity. The operator associated with the variational inequality is given by

$$ F(\theta) = \nabla L(\theta), \qquad (5.91) $$

and monotonicity is defined by the condition

$$ \langle F(\theta_1) - F(\theta_2), \theta_1 - \theta_2 \rangle \geq 0, \quad \forall \theta_1, \theta_2 \in \mathbb{R}^d. \qquad (5.92) $$

For convex loss functions, $\nabla L(\theta)$ is a monotone operator, ensuring the uniqueness of the equilibrium point. If $\nabla L(\theta)$ is strongly monotone, satisfying

$$ \langle \nabla L(\theta_1) - \nabla L(\theta_2), \theta_1 - \theta_2 \rangle \geq m \|\theta_1 - \theta_2\|^2, \quad \text{for some } m > 0, \qquad (5.93) $$

then the optimization process exhibits **exponential convergence**, meaning that the iterates satisfy

$$ \|\theta^{(k)} - \theta^*\| \leq C e^{-mk} \|\theta^{(0)} - \theta^*\|, \qquad (5.94) $$

where $C$ is a constant. This strong monotonicity property provides a rigorous foundation for understanding the rate of convergence of deep neural network training. In adversarial training, deep networks are often formulated as **saddle-point problems** of the form

$$ \min_{\theta} \max_{z} \mathcal{L}(\theta, z), \qquad (5.95) $$
where $z$ represents an adversarial perturbation. The equilibrium conditions in this case are described by the system of variational inequalities

$$ \langle \nabla_{\theta} \mathcal{L}(\theta^*, z^*), \theta - \theta^* \rangle \geq 0, \quad \forall \theta, \qquad (5.96) $$

$$ \langle \nabla_z \mathcal{L}(\theta^*, z^*), z - z^* \rangle \leq 0, \quad \forall z. \qquad (5.97) $$

These inequalities characterize the stable equilibrium of the adversarial learning process, ensuring that neither the network nor the adversary can unilaterally improve their respective objectives. The analysis of stochastic gradient descent (SGD) in deep learning can also be conducted via variational inequalities. The expected gradient field in the presence of stochasticity is given by

$$ F(\theta) = \mathbb{E}[\nabla L(\theta, \xi)], \qquad (5.98) $$

where $\xi$ represents a random variable encoding the data distribution. The variational inequality for SGD is then

$$ \langle F(\theta^*), \theta - \theta^* \rangle \geq 0, \quad \forall \theta \in \mathbb{R}^d. \qquad (5.99) $$

Stochastic approximation methods ensure convergence to the equilibrium point under conditions on the variance of the gradient estimator. Specifically, if $F(\theta)$ is **co-coercive**, meaning there exists $L > 0$ such that

$$ \|F(\theta_1) - F(\theta_2)\|^2 \leq L \langle F(\theta_1) - F(\theta_2), \theta_1 - \theta_2 \rangle, \qquad (5.100) $$

then projected SGD satisfies

$$ \mathbb{E}[\|\theta^{(k)} - \theta^*\|^2] \leq \frac{L}{2m} \frac{\sigma^2}{k}, \qquad (5.101) $$

where $\sigma^2$ is the variance of the stochastic gradient noise. This establishes a rigorous bound on the convergence rate of SGD in terms of variational inequalities. The generalization properties of deep networks can also be rigorously studied using variational inequalities. Given a distribution shift perturbation $\delta$, the parameter shift $\theta' = \theta + \delta$ satisfies the **generalization bound**

$$ \mathbb{E}[L(\theta')] - \mathbb{E}[L(\theta)] \leq C \|\delta\|, \qquad (5.102) $$

where the constant $C$ is determined by the Lipschitz continuity of the gradient field. This bound ensures that the network's performance remains stable under small perturbations in the data distribution. The variational inequality framework provides a **unified perspective** on optimization, equilibrium, and generalization in deep learning, offering rigorous theoretical insights into the dynamics of training and stability of neural networks.

### 5.4.1 Monotone/potential operator (gradient of a convex potential)

A monotone operator is a fundamental object in the theory of variational inequalities, convex optimization, and equilibrium problems. Formally, let $(H)$ be a real Hilbert space with inner product $\langle \cdot, \cdot \rangle$ and induced norm $| \cdot |$. An operator $T: H \to H$ is said to be monotone if for all $x, y \in H$, the inequality

$$ \langle T(x) - T(y), x - y \rangle \geq 0 \qquad (5.103) $$

holds. This condition encodes the idea that the operator does not decrease the angle between the difference of outputs and the difference of inputs, thereby generalizing the notion of a non-decreasing function in one dimension. An important class of monotone operators arises as gradients of convex functions. Let $\Phi: H \to \mathbb{R}$ be a convex and differentiable functional. Its gradient operator is defined by

$$ T(x) = \nabla \Phi(x). \qquad (5.104) $$

Since $\Phi$ is convex, it satisfies for all $x, y \in H$ the first-order characterization of convexity:

$$ \Phi(y) \geq \Phi(x) + \langle \nabla \Phi(x), y - x \rangle. \qquad (5.105) $$
<figure><img src="image_61.png" /><figcaption>Figure 5.10: Monotone/Potential F=grad Phi</figcaption></figure>

<figure><img src="image_62.png" /><figcaption>Figure 5.11: Monotone case: Approx VI gap vs iteration</figcaption></figure>

Switching the roles of $x$ and $y$ yields

$$ \Phi(x) \geq \Phi(y) + \langle \nabla \Phi(y), x - y \rangle. \qquad (5.106) $$

Adding these two inequalities gives

$$ \langle \nabla \Phi(x) - \nabla \Phi(y), x - y \rangle \geq 0, \qquad (5.107) $$

which proves that the gradient mapping $\nabla \Phi$ is a monotone operator. Furthermore, if $\Phi$ is strictly convex, then

$$ \langle \nabla \Phi(x) - \nabla \Phi(y), x - y \rangle > 0 \quad \text{for all } x \neq y, \qquad (5.108) $$
which implies strict monotonicity of the operator. This property ensures uniqueness of solutions to variational inequality problems of the form: find $x^* \in K$ such that

$$ \langle \nabla \Phi(x^*), y - x^* \rangle \geq 0 \quad \forall y \in K, \qquad (5.109) $$

where $K \subseteq H$ is a convex feasible set. The monotonicity condition can also be interpreted geometrically. For the gradient of a convex function, the operator $\nabla \Phi$ points in a direction that consistently aligns with the displacement vector between points in the domain. In other words, moving along the direction of the operator never contradicts the structure induced by convexity. This is the key mechanism underlying convergence of gradient-based algorithms in convex optimization.

In summary, the gradient of a convex potential is always a monotone operator, and the converse is also true in the sense that maximal cyclically monotone operators are precisely subdifferentials of convex functions. The equation

$$ T(x) = \nabla \Phi(x) \qquad (5.110) $$

together with the inequality

$$ \langle T(x) - T(y), x - y \rangle \geq 0 \qquad (5.111) $$

completely captures the rigorous mathematical relationship between monotone operators and convex potentials.

## 5.4.2 Saddle/non-monotone bilinear operator (models min-max behavior)

Let $x \in \mathbb{R}^n$, $y \in \mathbb{R}^m$ and denote $z = (x, y) \in \mathbb{R}^{n+m}$. A saddle (min-max) objective is a function $L: \mathbb{R}^n \times \mathbb{R}^m \to \mathbb{R}$ that is convex in $x$ for every fixed $y$ and concave in $y$ for every fixed $x$. Associated to $L$ one defines the operator

$$ F(z) = \begin{pmatrix} \nabla_x L(x, y) \\ -\nabla_y L(x, y) \end{pmatrix} \qquad (5.112) $$

so that a saddle point $(x^*, y^*)$ satisfies the variational inequality

$$ \langle F(z^*), z - z^* \rangle \geq 0 \quad \forall z = (x, y) \in \mathbb{R}^{n+m} \qquad (5.113) $$

This VI formulation is equivalent to the saddle condition $L(x^*, y) \le L(x^*, y^*) \le L(x, y^*)$ for all $x, y$ whenever (L) is convex-concave and sufficiently regular. A canonical and instructive special case is a bilinear objective

$$ L(x, y) = x^\top A y, \quad A \in \mathbb{R}^{n \times m}, \qquad (5.114) $$

for which

$$ F(x, y) = \begin{pmatrix} Ay \\ -A^\top x \end{pmatrix}. \qquad (5.115) $$

The stationary (saddle) conditions $\nabla_x L(x^*, y^*) = 0$ and $\nabla_y L(x^*, y^*) = 0$ reduce to

$$ Ay^* = 0, \quad A^\top x^* = 0, \qquad (5.116) $$

so the solution set is $(x, y): y \in \ker A, x \in \ker A^\top$; in particular multiplicity of solutions occurs whenever (A) is rank-deficient.

For two arbitrary points $z = (x, y)$ and $z' = (x', y')$ the bilinear operator satisfies the exact identity

$$ \langle F(z) - F(z'), z - z' \rangle = (x - x')^\top A (y - y') - (y - y')^\top A^\top (x - x') = 0 \qquad (5.117) $$
<figure><img src="image_63.png" /><figcaption>Figure 5.12: Saddle Bilinear F (min-max)</figcaption></figure>

hence (F) is (weakly) monotone in the sense that the left-hand side is identically zero for all $z, z'$. More generally for a smooth convex-concave $L$ one has the monotonicity inequality

$$ \langle F(z) - F(z'), z - z' \rangle \ge 0 \qquad (5.118) $$

but the bilinear case is the extreme where the quantity is exactly zero rather than strictly positive. The bilinear operator is therefore not a gradient field of any scalar potential on $\mathbb{R}^{n+m}$ unless it is identically zero, because for any gradient field $G = \nabla \Phi$ the Jacobian $\nabla G$ must be symmetric, whereas the Jacobian of the bilinear operator is the block matrix

$$ J(z) = \begin{pmatrix} 0 & A \\ -A^{\top} & 0 \end{pmatrix} \qquad (5.119) $$

which satisfies $J^{\top} = -J$, $J$ is skew-symmetric (block anti-symmetric) and thus cannot equal a symmetric matrix unless $J \equiv 0$. Consider the continuous-time gradient-flow-ascent dynamics that arise from (F),

$$ \dot{z}(t) = -F(z(t)) \qquad (5.120) $$

which for the bilinear case is the linear ODE

$$ \dot{z}(t) = -Jz(t). \qquad (5.121) $$

Because $J^{\top} = -J$, all eigenvalues of (J) are purely imaginary (or zero); equivalently (J) generates a one-parameter family of orthogonal transformations and the quadratic energy $\frac{1}{2}|z|^2$ is conserved along trajectories:

$$ \frac{d}{dt} \frac{1}{2} |z(t)|^2 =; z(t)^{\top} \dot{z}(t) = -z(t)^{\top} J z(t) = 0. \qquad (5.122) $$

In the simplest scalar toy $n = m = 1$ with $L(x, y) = xy$ one obtains $F(x, y) = (y, -x)$, Jacobian $J = \begin{pmatrix} 0 & 1 \\ -1 & 0 \end{pmatrix}$, and the continuous dynamics

$$ \dot{x} = -y, \quad y = x, \qquad (5.123) $$

which is the harmonic oscillator; solutions are circular periodic orbits and the radius $\sqrt{x^2 + y^2}$ is constant in time.

The behaviour in discrete-time simultaneous gradient steps (explicit Euler) is markedly different
and typically unstable. For a constant linear operator $F(z) = Jz$ the explicit simultaneous update with step-size $\eta > 0$ reads

$$z_{k+1} = z_k - \eta F(z_k) = (I - \eta J)z_k. \qquad (5.124)$$

For the scalar bilinear example $J = (0 \ 11 \ 0)$ the update matrix is

$$M = \begin{pmatrix} 1 & -\eta[4pt] \\ \eta & 1 \end{pmatrix} \qquad (5.125)$$

whose eigenvalues are $\lambda = 1 \pm i\eta$ and satisfy $|\lambda| = \sqrt{1 + \eta^2} > 1$. Thus the explicit (simultaneous) method amplifies the norm and the iterates diverge for any fixed $\eta > 0$. This instability is the discrete-time manifestation of attempting to integrate a conservative rotational field with an explicit integrator: energy conservation in continuous time becomes energy growth under the explicit Euler map.

Stabilization of these rotational saddle dynamics can be achieved in several mathematically transparent ways. Adding a symmetric dissipative term $S = S^\top \succeq \mu I$ (for $\mu > 0$) to the operator, i.e. replacing $F(z) = Jz$ by $F_{\text{diss}}(z) = (J + S)z$, yields a Jacobian whose symmetric part is (S) and hence the operator becomes strongly monotone with modulus at least $\mu$:

$$\langle F_{\text{diss}}(z) - F_{\text{diss}}(z'), z - z' \rangle = (z - z')^\top S (z - z') \geq \mu |z - z'|^2 \qquad (5.126)$$

Strong monotonicity implies uniqueness of the VI solution and gives (linear) contraction under appropriate implicit or proximal algorithms. Algorithmically, the extragradient method applied to a Lipschitz monotone operator $F$ with stepsize $\eta$ produces the linear update for $F(z) = Jz$,

$$z_{k+1} = (I - \eta J + \eta^2 J^2) z_k \qquad (5.127)$$

and for a skew-symmetric $J$ the term $J^2$ is symmetric negative semidefinite; in the scalar bilinear example $J^2 = -I$ giving

$$z_{k+1} = ((1 - \eta^2)I - \eta J)z_k \qquad (5.128)$$

whose eigenvalues are $(1 - \eta^2) \pm i\eta$ and hence have modulus $\sqrt{1 - \eta^2 + \eta^4} < 1$ for $0 < \eta < 1$ consequently the extragradient step produces a contraction for sufficiently small $\eta$ and stabilizes the iteration where explicit simultaneous updates blow up. More generally, for monotone and L-Lipschitz operators the extragradient method converges with step-size $\eta \le 1/(2L)$ (linear algebraic detail omitted here).

In summary, bilinear saddle operators arising in min-max problems furnish linear, skew-adjoint interaction terms that are not gradients of any scalar potential and which induce purely rotational continuous dynamics and discrete-time instabilities under naive simultaneous gradient updates. The mathematical diagnostics are

1. The block-Jacobian

$$J = \begin{pmatrix} 0 & A \\ -A^\top & 0 \end{pmatrix} \qquad (5.129)$$

with $J^\top = -J$

2. The identity

$$\langle F(z) - F(z'), z - z' \rangle \equiv 0 \qquad (5.130)$$

for the bilinear case

3. Conservation of quadratic energies in continuous time, and
4. Discrete amplification

$$|\lambda(I - \eta J)| > 1 \qquad (5.131)$$

for explicit Euler.

Algorithmic remedies—addition of symmetric dissipative terms, implicit/proximal discretizations, or extrag gradient/optimistic methods—convert the spectral properties so that the symmetric part of the Jacobian is positive (strong monotonicity) or the discrete map becomes contractive, thereby restoring convergence to equilibrium.

5.4.3 Skew (rotation)

A skew operator is a linear operator whose defining property is antisymmetry with respect to the underlying inner product. In the setting of a finite-dimensional real Hilbert space $\mathbb{R}^d$ with standard Euclidean inner product $\langle u, v \rangle = u^\top v$, a matrix $J \in \mathbb{R}^{d \times d}$ is said to be skew-symmetric if

$$J^\top = -J. \qquad (5.132)$$

Equivalently, for all vectors $u, v \in \mathbb{R}^d$, one has

$$\langle Ju, v \rangle = -\langle u, Jv \rangle. \qquad (5.133)$$

This antisymmetry implies that the quadratic form associated with $J$ vanishes identically, because for any $u \in \mathbb{R}^d$

$$\langle Ju, u \rangle = u^\top Ju = 0. \qquad (5.134)$$

Thus, skew operators generate purely rotational effects in the vector space, preserving lengths and inner products. The eigenstructure of skew-symmetric matrices is determined by their antisymmetry. If $J$ is real skew-symmetric, then for any eigenvalue–eigenvector pair $(\lambda, u)$ with $u \neq 0$, one has

$$Ju = \lambda u. \qquad (5.135)$$

Taking inner products with $u$, it follows that

$$u^\top Ju = \lambda |u|^2. \qquad (5.136)$$

But the left-hand side vanishes since $J$ is skew-symmetric, implying $\lambda|u|^2 = 0$. Therefore, every eigenvalue is either zero or purely imaginary. In fact, the spectrum of $J$ is symmetric with respect to the imaginary axis, and the nonzero eigenvalues occur in complex conjugate pairs $\pm i\omega$. This fact reveals that $J$ generates infinitesimal rotations, since purely imaginary eigenvalues correspond to circular motions. The matrix exponential of a skew-symmetric operator describes finite rotations. For any $t \in \mathbb{R}$,

$$R(t) = e^{tJ} = \sum_{k=0}^{\infty} \frac{t^k}{k!} J^k, \qquad (5.137)$$

is an orthogonal matrix with determinant one, hence a rotation matrix. Orthogonality follows from the fact that

$$R(t)^\top R(t) = e^{tJ^\top} e^{tJ} = e^{-tJ} e^{tJ} = I, \qquad (5.138)$$

since $J^\top = -J$. Thus, the flow generated by the differential equation

$$\dot{z}(t) = Jz(t) \qquad (5.139)$$

is given by

$$z(t) = e^{tJ} z(0) \qquad (5.140)$$
<figure><img src="image_64.png" /><figcaption>Figure 5.13: Skew (rotation) F: R z</figcaption></figure>

which traces a circular or rotational trajectory in the invariant subspaces associated with each conjugate pair of imaginary eigenvalues. The simplest case is the two-dimensional rotation operator. Let

$$J = \begin{pmatrix} 0 & -1 \\ 1 & 0 \end{pmatrix}. \qquad (5.141)$$

Then the ODE $\dot{z} = Jz$ expands to

$$\dot{x} = -y, \quad \dot{y} = x, \qquad (5.142)$$

which is the harmonic oscillator system. Its solution is

$$\begin{pmatrix} x(t) \\ y(t) \end{pmatrix} = \begin{pmatrix} \cos t & -\sin t \\ \sin t & \cos t \end{pmatrix} \begin{pmatrix} x(0) \\ y(0) \end{pmatrix}. \qquad (5.143)$$

Thus, the trajectory is a circle centered at the origin with radius $\sqrt{x(0)^2 + y(0)^2}$, and the motion is uniform rotation with unit angular velocity. The energy

$$E(t) = \frac{1}{2}|z(t)|^2 = \frac{1}{2}(x(t)^2 + y(t)^2) \qquad (5.144)$$

is conserved exactly, because

$$\frac{d}{dt} E(t) = \langle z(t), \dot{z}(t) \rangle = z(t)^{\top} J z(t) = 0 \qquad (5.145)$$

In higher dimensions, skew-symmetric operators generate rotations within orthogonal two-dimensional subspaces corresponding to each $\pm i\omega$ pair of eigenvalues, while leaving the null space ker $J$ invariant. That is, $\mathbb{R}^d$ decomposes into a direct sum of invariant planes where circular motions occur, and invariant lines where the dynamics are static. More explicitly, there exists an orthogonal transformation $Q$ such that

$$Q^{\top} J Q = \mathrm{diag} \left( \begin{bmatrix} 0 & -\omega_1 \\ \omega_1 & 0 \end{bmatrix}, \dots, \begin{bmatrix} 0 & -\omega_k \\ \omega_k & 0 \end{bmatrix}, 0, \dots, 0 \right) \qquad (5.146)$$

with $\omega_j > 0$. Each $2 \times 2$ block generates uniform rotation in its associated plane with angular speed $\omega_j$.

The geometric effect of skew operators is thus entirely conservative: they preserve Euclidean distances and angles, merely redistributing directions through rotation. In optimization and game-theoretic dynamics, skew components in an operator correspond to cycling or rotational forces
around equilibria rather than dissipative movement toward equilibria. This is reflected in the fact that monotonicity of skew operators is degenerate: for all $u, v \in \mathbb{R}^d$

$$ \langle J(u - v), u - v \rangle = 0 \qquad (5.147) $$

so skew operators are monotone but not strictly monotone. Their essential role is to encode conservative rotational dynamics rather than convergence-inducing gradients.

### 5.4.4 Skew + Dissipation

Let $H = \mathbb{R}^d$ with the Euclidean inner product $\langle u, v \rangle = u^\top v$. Let $J : \mathbb{R}^d \to \mathbb{R}^{d \times d}$ denote a linearization of an operator $F$ at a point (or consider the linear operator $F(z) = Jz$). Decompose $J$ into its symmetric and skew-symmetric parts

$$ J = S + K, \quad S = \frac{1}{2}(J + J^\top), \quad K = \frac{1}{2}(J - J^\top), \qquad (5.148) $$

with $S^\top = S$ and $K^\top = -K$. The phrase “Skew + Dissipation” refers precisely to this decomposition when $S$ has positive semidefinite (or positive definite) character and $(K)$ is skew-symmetric: $S \succeq 0$ (or $S \succeq \mu I$ with $\mu > 0$) and $K^\top = -K$. The algebraic and dynamical consequences follow from elementary identities and comparison inequalities.

For any $x, y \in \mathbb{R}^d$ one has the exact identity

$$ \langle J(x) - J(y), x - y \rangle = (x - y)^\top S (x - y), \qquad (5.149) $$

so the skew part $K$ does not enter the quadratic form that governs monotonicity. If $S \succeq 0$ then $J$ is monotone in the sense that

$$ \langle J(x) - J(y), x - y \rangle \geq 0 \quad \forall x, y, \qquad (5.150) $$

and if $S \succeq \mu I$ for some $\mu > 0$ then $J$ is $\mu$-strongly monotone:

$$ \langle J(x) - J(y), x - y \rangle \geq \mu |x - y|^2 \quad \forall x, y. \qquad (5.151) $$

Strong monotonicity immediately yields uniqueness of zeros: if $J(x^*) = 0$ and $J(y^*) = 0$ then $(x^* - y^*)^\top S (x^* - y^*) = 0$ and hence $x^* = y^*$ when $\mu > 0$. Consider the continuous-time linear

<figure><img src="image_65.png" /><figcaption>Figure 5.14: Skew + Dissipation F: R + alpha I</figcaption></figure>
dynamics $\dot{z}(t) = -Jz(t) = -Sz(t) - Kz(t)$. The instantaneous rate of change of the quadratic energy $E(t) = \frac{1}{2}|z(t)|^2$ satisfies, using the skew property $z^\top K z = 0$,

$$ \frac{d}{dt} E(t) = \langle z(t), \dot{z}(t) \rangle = -z(t)^\top S z(t) \leq -\mu |z(t)|^2, \qquad (5.152) $$

if $S \succeq \mu I$. By Grönwall's inequality this yields exponential decay

$$ |z(t)| \leq e^{-\mu t} |z(0)|, \qquad (5.153) $$

so the dissipation captured by $S$ converts the conservative rotational action of (K) into an overall contracting flow. If $S \succeq 0$ with $\mu = 0$ then $E(t)$ is nonincreasing and one obtains Lyapunov stability but not exponential convergence in general.

Spectrally, if $Jv = \lambda v$ for a (complex) eigenpair $(\lambda, v)$ then

$$ \lambda, v^*v = v^*Jv = v^*Sv + v^*Kv, \qquad (5.154) $$

and taking real parts gives the Rayleigh-type relation

$$ \Re(\lambda) = \frac{v^*Sv}{v^*v}. \qquad (5.155) $$

Hence the real parts of eigenvalues of $J$ lie in the numerical range of $S$; in particular $S \succeq \mu I$ implies $\Re(\lambda) \geq \mu$ for every eigenvalue $\lambda$ of $J$. Consequently the linear semigroup generated by $-J$ is exponentially stable with decay rate at least $\mu$. The skew part $K$ only shifts eigenvalues along the imaginary axis (sets the oscillation frequencies) but cannot move their real parts away from those determined by $S$.

In discrete time the explicit Euler map with step-size $\eta > 0$,

$$ z_{k+1} = (I - \eta J)z_k, \qquad (5.156) $$

is stable iff the spectral radius $\rho(I - \eta J) < 1$. The purely skew case $S = 0$ yields eigenvalues $1 \pm i\eta\omega$ in canonical $2 \times 2$ rotational blocks and hence $|\lambda| = \sqrt{1 + \eta^2\omega^2} > 1$, so explicit steps amplify and diverge. By contrast, when $S = \alpha I$ (scalar dissipation $\alpha \geq 0$) and $K$ is a rotation with frequency $\omega$, the update matrix in the $2 \times 2$ invariant plane is

$$ M = I - \eta(\alpha I + K) = (1 - \eta\alpha \quad \eta\omega[4pt] - \eta\omega \quad 1 - \eta\alpha), \qquad (5.157) $$

whose eigenvalues $\lambda = (1 - \eta\alpha) \pm i\eta\omega$ satisfy

$$ |\lambda|^2 = (1 - \eta\alpha)^2 + \eta^2\omega^2. \qquad (5.158) $$

Thus a sufficient stability condition is

$$ (1 - \eta\alpha)^2 + \eta^2\omega^2 < 1 \iff \eta < \frac{2\alpha}{\alpha^2 + \omega^2}, \qquad (5.159) $$

showing quantitatively how dissipation $\alpha$ permits a finite explicit step-size that stabilizes the discrete map; when $\alpha \downarrow 0$ no positive $\eta$ suffices.

Algorithmic stabilization may also be effected by higher-fidelity integrators. The extragradient (EG) linearized update applied to a constant $J$ produces the map

$$ M_{EG} = I - \eta J + \eta^2 J^2. \qquad (5.160) $$
If $J = K$ is pure skew with $K^2 = -\omega^2 I$ on a rotation plane, then

$$M_{EG} = (1 - \eta^2 \omega^2) I - \eta K, \qquad (5.161)$$

whose eigenvalues are $(1 - \eta^2 \omega^2) \pm i\eta\omega$ and hence

$$|\lambda|^2 = (1 - \eta^2 \omega^2)^2 + \eta^2 \omega^2 = 1 - \eta^2 \omega^2 + \eta^4 \omega^4 < 1 \qquad (5.162)$$

for sufficiently small $\eta > 0$. Thus EG introduces an effective symmetric contraction term $-\eta^2 J^2$ (which is positive semidefinite when $J$ is skew) and converts rotation into contraction; more generally for $J = S + K$ the $J^2$ term couples $S$ and $K$ and can improve stability relative to simultaneous explicit updates.

From the variational inequality viewpoint, adding dissipation upgrades monotonicity: if $F(z) = Jz$ with $J = S + K$ and $S \succeq \mu I$, then for the VI problem $\langle F(z^*), z - z^* \rangle \ge 0$ the operator is strongly monotone and the VI has a unique solution. The resolvent (proximal) operator $(I + \eta F)^{-1}$ is then a contraction: if $u = (I + \eta F)^{-1}x$ and $v = (I + \eta F)^{-1}y$ then

$$u - v + \eta(F(u) - F(v)) = x - y, \qquad (5.163)$$

and inner-producting with $u - v$ and using strong monotonicity gives

$$\langle u - v, x - y \rangle \ge (1 + \eta\mu)|u - v|^2, \qquad (5.164)$$

hence

$$|u - v| \le \frac{1}{1 + \eta\mu}|x - y|. \qquad (5.165)$$

This resolvent contraction is the continuous-time analogue of how dissipation regularizes and guarantees stable implicit updates for nonconservative dynamics.

In summary, the mathematical role of “Skew + Dissipation” is: $K$ encodes conservative rotational forces (purely imaginary spectral displacements, energy-preserving infinitesimal action) while $S$ provides symmetric dissipation that determines real parts of the spectrum, monotonicity, energy decay, uniqueness, and contractivity. The interplay is quantitative: the magnitude of $K$ sets frequencies and oscillation amplitudes, the minimal eigenvalue $\mu$ of $S$ sets exponential decay rates in continuous time and allowable explicit step-sizes in discrete time, and algorithmic devices (EG, implicit resolvents, preconditioning) exploit the positive symmetric component or introduce effective symmetric terms to neutralize the destabilizing influence of the skew part.

### 5.4.5 Python Code to generate Figures 5.10, 5.11 illustrating Monotone/potential operator and Figure 5.12 illustrating Saddle Bilinear F (min-max) and Figure 5.13 illustrating Skew (rotation) F: R z and Figure 5.14 illustrating Skew + Dissipation F: R + alpha I

The Python code below produces the Figures 5.10, 5.11 illustrating Monotone/potential operator and Figure 5.12 illustrating Saddle Bilinear F (min-max) and Figure 5.13 illustrating Skew (rotation) F: R z and Figure 5.14 illustrating Skew + Dissipation F: R + alpha I.

```python
# """
# Variational Inequality illustrations for DNN-style dynamics.
# 
# Shows:
# - Monotone (potential) operator: F = grad Phi, with Phi convex.
```
- Saddle (bilinear) operator: $F$ from $L(x,y) = x*y$ (min-max).

- Skew/rotational operator: pure rotation -> cyclic.

Compares:

- Gradient step: $x_{k+1} = x_k - lr * F(x_k)$

- Extragradient (EG): $x_{\text{half}} = x_k - lr * F(x_k)$; $x_{k+1} = x_k - lr * F(x_{\text{half}})$

Also computes an approximate VI gap on a grid for the monotone case:

$\text{gap}(x) \approx \max_y < F(x), x - y >$

(we maximize over a finite grid to approximate the sup)

"
"

import numpy as np

import matplotlib.pyplot as plt

# -----------------------

# Helper / algorithms

# -----------------------

def grad_step(x, F, lr):
    return x - lr * F(x)

def extragradient_step(x, F, lr):
    xhalf = x - lr * F(x)
    return x - lr * F(xhalf)

def simulate_iterates(x0, F, lr, steps, method="grad"):

x = np.array(x0, dtypefloat)"

traj = [x.copy()]

for k in range(steps):

if method == "grad":

x = grad_step(x, F, lr)

elif method == "eg":

x = extragradient_step(x, F, lr)

else:

raise ValueError("method must be 'grad' or 'eg'")

traj.append(x.copy())

return np.array(traj)

# -----------------------

# Operators (2D)

# -----------------------

# 1) Monotone / Potential: $F = \text{grad} \Phi_i, \Phi_i(x,y) = 0.5*( (x-a)^2 + (y-b)^2 )$

a, b = 0.8, -0.5

def F Coronary(z):

x, y = z

return np.array([x - a, y - b])

# 2) Saddle / Bilinear: L(x,y)=x*y, treat z=(x,y) with operator F = (d/dx L, -d/dy L)

# For a min in x and max in y, the VI operator is F = [d/dx L, -d/dy L]
] = [y, x]

def F_saddle(z):

x, y = z
5.4. ANALYSIS OF DEEP NEURAL NETWORKS (DNNS) THROUGH VARIATIONAL INEQUALITIES

```python
return np.array([y, x])  # note: non-monotone, leads to cycles
for naive methods

# 3) Skew / rotation: F(z) = R z, R is skew-symmetric -> purely
rotational (no potential)
theta = np.pi/2 * 0.7  # rotation-like skew (not a pure rotation matrix;
keep skew)
R = np.array([[0.0, -1.0],
               [1.0,  0.0]]) * 1.0  # pure skew gives rotation effect

def F_skew(z):
    return R.dot(z)

# For richer behavior we will include a small symmetric (dissipative)
# part optionally:
def make_skew_PLUS_DISSIP(alpha=0.0):
    S = np.eye(2) * alpha  # symmetric positive definite part (
                          # dissipation)
    M = R + S
    return lambda z: M.dot(z)

# ------------------------
# Plotting utilities
# ------------------------
def plot_vector_field(ax, F, xrange=(-2.5,2.5), yrange=(-2.5,2.5),
    density=20):
    xs = np.lliavel(xrange[0], xrange[1], density)
    ys = np.lliavel(yrange[0], yrange[1], density)
    X, Y = np.meshgrid(xs, ys)
    U = np(zeros_like(X))
    V = np(zeros_like(Y))
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            f = F(np.array([X[i,j], Y[i,j]]))
            U[i,j] = f[0]
            V[i,j] = f[1]
    # normalize for nicer arrows
    M = np.sqrt(U**2 + V**2)
    M[M==0] = 1.0
    ax.streamplot(X, Y, -U, -V, density=1.2, linewidth=1)  # show
    descent direction -F
    ax.quiver(X, Y, -U/M, -V/M, angles='xy', width=0.003, alpha=0.6)
def plot_trjectories(F, name, lr=0.05, steps=500, initial_points=None,
    xrange=(-2.5,2.5), yrange=(-2.5,2.5)):
    if initial_points is None:
        initial_points = [np.array([1.5, 0.5]), np.array([-1.5, 1.2]),
                          np.array([0.5, -1.5]), np.array([-0.8, -0.6])]
    fig, axes = plt.subplots(1,2, figsize=(12,5))
    # vector field
    plot_vector_field(axes[0], F, xrange=xrange, yrange=yrange)
    axes[0].set_title(f"{name": vector field (-F shown for descent)")
    axes[0].set_xlim(xrange); axes[0].set_ylim(yrange)
    axes[0].axhline(0, color='k', lw=0.4); axes[0].axvline(0, color='k',
        lw=0.4)
```
```python
102 # trajectories: gradient vs extragrwardient
103 for p0 in initial_points:
104     traj_g = simulate_iterates(p0, F, lr, steps, method="grad")
105     traj eg = simulate_iterates(p0, F, lr, steps, method="eg")
106     axes[1].plot(traj_g[:,0], traj_g[:,1],",-", alpha=0.6, label="
107         Grad" if (p0==initial_points[0]).all() else "")
108     axes[1].plot(traj eg [:,0], traj eg [:,1],"--", alpha=0.9, label="
109         Extragrwardient" if (p0==initial_points[0]).all() else "")
110     axes[1].plot(p0[0], p0[1], 'o', markersize=4)
111     axes[1].set_title(f"{name}: trajectories (lr={lr})")
112     axes[1].set_xlim(xrange); axes[1].set_ylim(yrange)
113     axes[1].legend()
114     axes[1].grid(True)
115     plt.tight_layout()
116     plt.show()
117
118 # ------------------------
119 # 1) Monotone / Potential example
120 # ------------------------
121 print("Plotting MONOTONE/POTENTIAL operator")
122 plot trajectories(F_potential, "Monotone/Potential F=grad Phi", lr=0.1,
123                  steps=200,
124                  initial_points=[np.array([2.0, 1.5]), np.array([-2.0,
125                   -1.0])], np.array([1.0,-2.0])],
125                   xrange=(-3,3), yrange=(-3,3))
126
127 # Compute approximate VI gap over grid for monotone case to compare
128 algorithms
129 def approx_vi-gap(F, x, grid_points):
130     # approximate gap(x) = sup_y <F(x), x - y> by maximizing over a
131     # finite grid of y values
132     Fx = F(x)
133     diffs = x.reshape(1,2) - grid_points  # (m,2)
134     inner = diffs.dot(Fx)              # (m,)
135     return np.max(inner)
136
137 # prepare grid once
138 gridx = np.linspace(-3,3,41)
139 gridy = np.linspace(-3,3,41)
140 YY, XX = np.meshgrid(gridy, gridx)
141 grid_points = np.vstack([XX.ravel(), YY.ravel()]).T
142
143 # simulate and track approximate gap
144 lr = 0.1
145 steps = 200
146 x0 = np.array([2.0,1.0])
147 traj_g = simulate_iterates(x0, F_potential, lr, steps, method="grad")
148 traj eg = simulate_iterates(x0, F_potential, lr, steps, method="eg")
149 gaps_g = [approx_vi-gap(F_potential, traj_g[k], grid_points) for k in range
150            (len(traj_g))]
151 gaps eg = [approx_vi-gap(F_potential, traj eg[k], grid_points) for k in
152            range(len(traj eg))]
153
154 plt.figure(figsize=(6,4))
155 plt plot(gaps_g, label="GD gap")
```
```python
149 plt.plot(gaps_eg, label="Extragradients gap")
150 plt.yscale('log')
151 plt.xlabel("Iteration")
152 plt.ylabel("Approx VI gap (log scale)")
153 plt.title("Monotone case: Approx VI gap vs iteration")
154 plt.legend()
155 plt.grid(True, which='both', ls='-')
156 plt.show()
157
158 # ------------------------
159 # 2) Saddle / Bilinear example
160 # ------------------------
161 print("Plotting SADDLE / BILINEAR operator")
162 plot_trjectories(F_saddle, "Saddle/Bilinear F (min-max)", lr=0.05,
163 steps=800,
164 initial_points=[np.array([1.5, 0.5]), np.array([-1.2,
165 1.5]), np.array([0.2,-1.7])],
166 xrange=(-3,3), yrange=(-3,3))
167
168 # ------------------------
169 # 3) Skew / Rotation example
170 # ------------------------
171 print("Plotting SKEW / ROTATIONAL operator (purely rotational behavior)")
172 F_skew_PLUS = make_skew_PLUS_dissip(alpha=0.0) # pure skew -> cycles
173
174 plot_trjectories(F_skew_PLUS, "Skew (rotation) F: R z", lr=0.05, steps
175 =700,
176 initial_points=[np.array([1.5,0.2]), np.array
177 ([1.2,0.8]), np.array([0.5,-1.2])],
178 xrange=(-2.5,2.5), yrange=(-2.5,2.5))
179
180 # Also show effect of adding small dissipation (symmetric positive part)
181 -> stabilized dynamics
182 print("Plotting SKEW + SMALL DISSIPATION (stabilized)")
183 F_skew_stable = make_skew_PLUS_dissip(alpha=0.3)
184
185 plot_trjectories(F_skew_stable, "Skew + Dissipation F: R + alpha I", lr
186 =0.08, steps=300,
187 initial_points=[np.array([1.2,0.6]), np.array
188 ([1.0,1.1])],
189 xrange=(-2.5,2.5), yrange=(-2.5,2.5))
190
```

## 5.5 Optimal control in reinforcement learning (RL)-based deep neural network (DNN) training

Let $\theta_t \in \mathbb{R}^n$ represent the weight parameters of a deep neural network (DNN) at time step $t$. The evolution of these parameters over discrete time is governed by a controlled dynamical system, where the control input $u_t \in \mathbb{R}^n$ represents the weight update at each iteration. The fundamental objective of reinforcement learning (RL)-based DNN training is to determine an optimal control policy $u_t^*$ that minimizes a given performance criterion over a finite or infinite training horizon. The system dynamics are given by the equation

$$ \theta_{t+1} = \theta_t + u_t. \qquad (5.166) $$
Let $J(\theta_0, u)$ denote the cumulative loss function associated with a given control sequence $\{u_t\}_{t=0}^T$, defined as

$$J(\theta_0, u) = \sum_{t=0}^{T} c(\theta_t, u_t), \qquad (5.167)$$

where $c(\theta_t, u_t)$ is the instantaneous cost function. In the context of neural network training, the cost function is typically the empirical risk over a dataset $\mathcal{D}$, given by

$$c(\theta_t, u_t) = \mathbb{E}_{(x,y)\sim\mathcal{D}}[\ell(f_{\theta_t}(x), y)], \qquad (5.168)$$

where $\ell$ is a loss function such as cross-entropy or mean squared error, and $f_{\theta_t}(x)$ represents the neural network's output for an input $x$. The optimal control problem consists of finding a control sequence $\{u_t^*\}_{t=0}^T$ such that

$$J^*(\theta_0) = \min_u J(\theta_0, u). \qquad (5.169)$$

Applying the principle of dynamic programming, the optimal cost-to-go function is defined recursively as

$$V(\theta_t) = \min_{u_t} \{c(\theta_t, u_t) + V(\theta_{t+1})\}. \qquad (5.170)$$

Expanding $\theta_{t+1} = \theta_t + u_t$, the Bellman equation takes the form

$$V(\theta_t) = \min_{u_t} \{c(\theta_t, u_t) + V(\theta_t + u_t)\}. \qquad (5.171)$$

Considering the continuous-time limit where the weight updates are infinitesimal, i.e., $u_t \to \dot{\theta}(t)dt$, we obtain the Hamilton-Jacobi-Bellman (HJB) equation

$$\frac{\partial V}{\partial t} + \min_{u} \left[ \frac{\partial V}{\partial \theta} u + c(\theta, u) \right] = 0. \qquad (5.172)$$

The optimal control $u^*(t)$ satisfies

$$u^*(t) = \arg \min_{u} \left[ \frac{\partial V}{\partial \theta} u + c(\theta, u) \right]. \qquad (5.173)$$

Alternatively, the problem can be analyzed using Pontryagin's Maximum Principle. Define the Hamiltonian

$$H(\theta, u, \lambda) = c(\theta, u) + \lambda^{\top}(\theta + u), \qquad (5.174)$$

where $\lambda$ is the costate variable that evolves according to the adjoint equation

$$\frac{d\lambda}{dt} = -\frac{\partial H}{\partial \theta}. \qquad (5.175)$$

For optimality, the control must satisfy

$$\frac{\partial H}{\partial u} = \frac{\partial c}{\partial u} + \lambda = 0. \qquad (5.176)$$

Solving for $u^*$, we obtain

$$u^* = - \left( \frac{\partial^2 c}{\partial u^2} \right)^{-1} \left( \frac{\partial c}{\partial u} + \lambda \right). \qquad (5.177)$$

In the reinforcement learning framework, the optimal control formulation is connected to policy optimization, where the policy $\pi(a_t|s_t)$ determines the action $a_t$ given the state $s_t$. The expected return under policy $\pi$ is given by

$$J(\pi) = \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t=0}^{T} \gamma^t R_t \right], \qquad (5.178)$$
where $R_t$ is the reward function and $\gamma \in (0, 1]$ is the discount factor. The policy gradient theorem states that

$$ \nabla_{\theta} J(\pi) = \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t=0}^{T} \nabla_{\theta} \log \pi(a_t | s_t) Q^{\pi}(s_t, a_t) \right], \quad (5.179) $$

where $Q^{\pi}(s_t, a_t)$ is the state-action value function satisfying the Bellman equation

$$ Q^{\pi}(s_t, a_t) = \mathbb{E}[R_t + \gamma Q^{\pi}(s_{t+1}, a_{t+1})|s_t, a_t]. \quad (5.180) $$

By applying the HJB equation, the optimal policy update rule is derived as

$$ \theta_{t+1} = \theta_t - \eta \nabla_{\theta} J(\pi), \quad (5.181) $$

where $\eta$ is the learning rate, often adaptively adjusted based on second-order curvature information such as in natural gradient descent. More formally, using the Fisher information matrix $F(\theta)$, the optimal update rule can be written as

$$ \theta_{t+1} = \theta_t - \eta F^{-1}(\theta) \nabla_{\theta} J(\pi). \quad (5.182) $$

This establishes the link between optimal control and policy optimization in reinforcement learning, providing a rigorous theoretical foundation for RL-based DNN training.

### 5.5.1 Backpropagation as Pontryagin's Maximum Principle (PMP)

We begin by formulating the supervised learning problem for a deep neural network as a discrete-time optimal control problem. Consider a network with $L$ layers. Let the state $\mathbf{x}_k \in \mathbb{R}^{n_k}$ represent the input to layer $k$, for $k = 0, 1, \dots, L$. The initial state $\mathbf{x}_0$ is the input data. The control variables $\mathbf{u}_k \in \mathbb{R}^{m_k}$ are the parameters (weights and biases) of layer $k$. The state dynamics, which describe the forward propagation, are given by the equation

$$ \mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \quad (5.183) $$

where $\mathbf{f}_k$ is the activation function (e.g., a linear transformation followed by a non-linear element-wise function like ReLU or tanh). The objective of the training is to minimize a total cost function $J$ that comprises a terminal cost and a running cost. The terminal cost is the loss function $\ell(\mathbf{x}_L, \mathbf{y})$ which measures the discrepancy between the final state (the network output) $\mathbf{x}_L$ and the target label $\mathbf{y}$. The running cost can incorporate regularization terms on the controls, such as $L^2$ weight decay, leading to the full cost function

$$ J(\mathbf{u}) = \ell(\mathbf{x}_L, \mathbf{y}) + \sum_{k=0}^{L-1} \mathcal{R}_k(\mathbf{u}_k) \quad (5.184) $$

The optimal control problem is to find the sequence of controls $\mathbf{u}_0, \mathbf{u}_1, \dots, \mathbf{u}_{L-1}$ that minimizes $J$ subject to the dynamical constraints

$$ \mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \quad (5.185) $$

Pontryagin's Maximum Principle (PMP) provides necessary conditions for optimality for this class of problems. To apply PMP, we first define the Hamiltonian function $H_k$ for each stage $k$. The Hamiltonian is given by

$$ H_k(\mathbf{x}_k, \mathbf{u}_k, \mathbf{p}_{k+1}) = \mathbf{p}_{k+1}^\top \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) - \mathcal{R}_k(\mathbf{u}_k) \quad (5.186) $$

where $\mathbf{p}_{k+1} \in \mathbb{R}^{n_{k+1}}$ is the costate vector. The PMP states that for an optimal control sequence $\mathbf{u}_k^*$ and the corresponding optimal state trajectory $\mathbf{x}_k^*$, there exists a corresponding costate trajectory $\mathbf{p}_k^*$ satisfying the following canonical equations. The state evolution is given by the forward equation

$$ \mathbf{x}_{k+1}^* = \nabla_{\mathbf{p}_{k+1}} H_k = \mathbf{f}_k(\mathbf{x}_k^*, \mathbf{u}_k^*) \quad (5.187) $$
which is identical to the network's forward propagation. The costate evolution is governed by the backward equation

$$ \mathbf{p}_k^* = \nabla_{\mathbf{x}_k} H_k = \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^*, \mathbf{u}_k^*)}{\partial \mathbf{x}_k} \right)^\top \mathbf{p}_{k+1}^* \quad (5.188) $$

for $k = 1, \dots, L-1$. The boundary condition for the costate at the final time is determined by the terminal cost,

$$ \mathbf{p}_L^* = -\nabla_{\mathbf{x}_L} \ell(\mathbf{x}_L^*, \mathbf{y}) \quad (5.189) $$

The maximization condition of PMP requires that the optimal control $\mathbf{u}_k^*$ maximizes the Hamiltonian at every step along the optimal trajectory, that is,

$$ H_k(\mathbf{x}_k^*, \mathbf{u}_k^*, \mathbf{p}_{k+1}^*) \geq H_k(\mathbf{x}_k^*, \mathbf{u}_k, \mathbf{p}_{k+1}^*) \quad (5.190) $$

for all admissible $\mathbf{u}_k$. For unconstrained controls, this often leads to the condition

$$ \nabla_{\mathbf{u}_k} H_k(\mathbf{x}_k^*, \mathbf{u}_k^*, \mathbf{p}_{k+1}^*) = \mathbf{0} \quad (5.191) $$

which expands to

$$ \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^*, \mathbf{u}_k^*)}{\partial \mathbf{u}_k} \right)^\top \mathbf{p}_{k+1}^* - \nabla_{\mathbf{u}_k} \mathcal{R}_k(\mathbf{u}_k^*) = \mathbf{0} \quad (5.192) $$

The connection to backpropagation emerges from the process of solving these necessary conditions. The standard method of successive approximation (MSA) involves a two-step iterative process. First, given a current control sequence $\mathbf{u}_k$, the state equation

$$ \mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \quad (5.193) $$

is propagated forward from $k=0$ to $k=L$ to compute the state trajectory. Second, the costate equation

$$ \mathbf{p}_k = \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k)}{\partial \mathbf{x}_k} \right)^\top \mathbf{p}_{k+1} \quad (5.194) $$

is propagated backward from $k=L$ to $k=1$, initialized with

$$ \mathbf{p}_L = -\nabla_{\mathbf{x}_L} \ell(\mathbf{x}_L, \mathbf{y}) \quad (5.195) $$

This backward pass is precisely the backpropagation algorithm. The vector $\mathbf{p}_k$ is exactly the error signal, or the gradient of the loss with respect to the state $\mathbf{x}_k$, that is

$$ \mathbf{p}_k = - \frac{\partial \ell}{\partial \mathbf{x}_k} \quad (5.196) $$

The final step of MSA is to update the controls to satisfy the PMP maximization condition. For a simple case with no regularization, $\mathcal{R}_k = 0$, the condition $\nabla_{\mathbf{u}_k} H_k = \mathbf{0}$ becomes

$$ \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k)}{\partial \mathbf{u}_k} \right)^\top \mathbf{p}_{k+1} = \mathbf{0} \quad (5.197) $$

In practice, this condition is not enforced directly but is used to compute the gradient for a gradient-based optimizer. The gradient of the total cost $J$ with respect to the control $\mathbf{u}_k$ is given by

$$ \nabla_{\mathbf{u}_k} J = -\nabla_{\mathbf{u}_k} H_k = - \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k)}{\partial \mathbf{u}_k} \right)^\top \mathbf{p}_{k+1} \quad (5.198) $$

This is the same gradient computed by the backpropagation through time (BPTT) algorithm, where the term $\left( \frac{\partial \mathbf{f}_k}{\partial \mathbf{u}_k} \right)^\top \mathbf{p}_{k+1}$ is the backward pass through the layer's parameter Jacobian. Therefore, the backward pass of backpropagation is mathematically equivalent to the costate backward evolution of PMP, and the parameter gradient update is a single step (e.g., gradient ascent on the Hamiltonian) towards satisfying the PMP maximization condition. This rigorous mathematical correspondence reveals that the training of deep neural networks via gradient descent and backpropagation is an implementation of the Method of Successive Approximation for solving the necessary optimality conditions dictated by Pontryagin's Maximum Principle.
5.5.2 The Method of Successive Approximation (MSA)

The Method of Successive Approximation (MSA) is a numerical algorithm for solving the necessary optimality conditions arising from Pontryagin's Maximum Principle (PMP) for a discrete-time optimal control problem. In the context of training a deep neural network, the problem is formulated as minimizing a total cost

$$
J(\mathbf{u}) = \ell(\mathbf{x}_L, \mathbf{y}) + \sum_{k=0}^{L-1} \mathcal{R}_k(\mathbf{u}_k) \qquad (5.199)
$$

subject to the state dynamics

$$
\mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \qquad (5.200)
$$

for $k = 0, 1, \dots, L-1$, where $\mathbf{x}_0$ is the input data, $\mathbf{x}_L$ is the network output, $\mathbf{y}$ is the target, $\mathbf{u}_k$ are the parameters of layer $k$, and $\mathcal{R}_k$ is a regularization function. The Hamiltonian for this system is defined as

$$
H_k(\mathbf{x}_k, \mathbf{u}_k, \mathbf{p}_{k+1}) = \mathbf{p}_{k+1}^\top \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) - \mathcal{R}_k(\mathbf{u}_k) \quad (5.201)
$$

where $\mathbf{p}_{k+1}$ is the costate variable. The PMP provides the necessary conditions for optimality: the state equation

$$
\mathbf{x}_{k+1}^* = \nabla_{\mathbf{p}_{k+1}} H_k = \mathbf{f}_k(\mathbf{x}_k^*, \mathbf{u}_k^*) \quad (5.202)
$$

the costate equation

$$
\mathbf{p}_k^* = \nabla_{\mathbf{x}_k} H_k = \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^*, \mathbf{u}_k^*)}{\partial \mathbf{x}_k} \right)^{\top} \mathbf{p}_{k+1}^* \quad (5.203)
$$

with the boundary condition

$$
\mathbf{p}_L^* = -\nabla_{\mathbf{x}_L} \ell(\mathbf{x}_L^*, \mathbf{y}) \qquad (5.204)
$$

and the maximization condition

$$
H_k(\mathbf{x}_k^*, \mathbf{u}_k^*, \mathbf{p}_{k+1}^*) \geq H_k(\mathbf{x}_k^*, \mathbf{u}_k, \mathbf{p}_{k+1}^*) \quad (5.205)
$$

for all admissible $\mathbf{u}_k$, which for unconstrained, differentiable problems implies the stationarity condition

$$
\nabla_{\mathbf{u}_k} H_k(\mathbf{x}_k^*, \mathbf{u}_k^*, \mathbf{p}_{k+1}^*) = \mathbf{0} \qquad (5.206)
$$

or explicitly

$$
\left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i)})}{\partial \mathbf{u}_k} \right)^{\top} \mathbf{p}_{k+1}^{(i)} - \nabla_{\mathbf{u}_k} \mathcal{R}_k(\mathbf{u}_k^{(i)}) = \mathbf{0} \qquad (5.207)
$$

The MSA algorithm iteratively refines a sequence of control policies {$\mathbf{u}_k^{(i)}$} to satisfy these condi-
tions. Given an initial guess for the control sequence $\mathbf{u}_k^{(0)}$ for $k = 0, \dots, L-1$, the $i$-th iteration
of MSA consists of two sequential passes. The first pass is the forward state propagation, which
computes the state trajectory $\mathbf{x}_k^{(i)}$ corresponding to the current control guess $\mathbf{u}_k^{(i)}$. This is executed
by initializing with the input data $\mathbf{x}_0^{(i)} = \mathbf{x}_0$ and recursively applying the dynamics

$$
\mathbf{x}_{k+1}^{(i)} = \mathbf{f}_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i)}) \qquad (5.208)
$$

for $k = 0$ to $L-1$. This forward pass is identical to a forward propagation through the neural
network using the current parameters. The second pass is the backward costate propagation, which
computes the costate trajectory $\mathbf{p}_k^{(i)}$ using the computed states $\mathbf{x}_k^{(i)}$ and controls $\mathbf{u}_k^{(i)}$. This pass is
initialized with the gradient of the terminal loss

$$
\mathbf{p}_L^{(i)} = -\nabla_{\mathbf{x}_L} \ell(\mathbf{x}_L^{(i)}, \mathbf{y}) \qquad (5.209)
$$

and proceeds backwards recursively using the costate equation

$$
\mathbf{p}_k^{(i)} = \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i)})}{\partial \mathbf{x}_k} \right)^{\top} \mathbf{p}_{k+1}^{(i)} \quad (5.210)
$$
for $k = L-1, L-2, \dots, 1$. This backward pass is mathematically equivalent to the backpropagation algorithm, as the costate $\mathbf{p}_k^{(i)}$ represents the negative gradient of the total cost with respect to the state,

$$ \mathbf{p}_k^{(i)} = -\frac{\partial J}{\partial \mathbf{x}_k} \qquad (5.211) $$

The third and final step of the MSA iteration is the control update, which aims to satisfy the PMP maximization condition. A new control sequence $\mathbf{u}_k^{(i+1)}$ is computed to maximize the Hamiltonian $H_k$ given the newly computed states $\mathbf{x}_k^{(i)}$ and costates $\mathbf{p}_{k+1}^{(i)}$. The update is performed pointwise for each $k$ by solving

$$ \mathbf{u}_k^{(i+1)} = \arg \max_{\mathbf{u}_k} H_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k, \mathbf{p}_{k+1}^{(i)}) = \arg \max_{\mathbf{u}_k} \left[ \mathbf{p}_{k+1}^{(i)\top} \mathbf{f}_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k) - \mathcal{R}_k(\mathbf{u}_k) \right] \qquad (5.212) $$

For the common case where the dynamics are differentiable with respect to $\mathbf{u}_k$ and the regularization is convex, the maximization condition leads to the implicit equation

$$ \nabla_{\mathbf{u}_k} H_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i+1)}, \mathbf{p}_{k+1}^{(i)}) = \mathbf{0} \qquad (5.213) $$

which expands to

$$ \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i+1)})}{\partial \mathbf{u}_k} \right)^{\top} \mathbf{p}_{k+1}^{(i)} - \nabla_{\mathbf{u}_k} \mathcal{R}_k(\mathbf{u}_k^{(i+1)}) = \mathbf{0} \qquad (5.214) $$

Solving this equation exactly for $\mathbf{u}_k^{(i+1)}$ is often computationally intractable. Instead, a common approximation is to perform a single gradient ascent step on the Hamiltonian with respect to the control. This involves computing the gradient

$$ \nabla_{\mathbf{u}_k} H_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i)}, \mathbf{p}_{k+1}^{(i)}) = \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i)})}{\partial \mathbf{u}_k} \right)^{\top} \mathbf{p}_{k+1}^{(i)} - \nabla_{\mathbf{u}_k} \mathcal{R}_k(\mathbf{u}_k^{(i)}) \qquad (5.215) $$

and then updating the controls as

$$ \mathbf{u}_k^{(i+1)} = \mathbf{u}_k^{(i)} + \eta \nabla_{\mathbf{u}_k} H_k(\mathbf{x}_k^{(i)}, \mathbf{u}_k^{(i)}, \mathbf{p}_{k+1}^{(i)}) \qquad (5.216) $$

where $\eta$ is a learning rate. Crucially, the gradient $\nabla_{\mathbf{u}_k} H_k$ is precisely the negative of the gradient required to minimize the total cost $J$, since

$$ \nabla_{\mathbf{u}_k} J = -\nabla_{\mathbf{u}_k} H_k \qquad (5.217) $$

when the state and costate equations are satisfied. Therefore, the control update

$$ \mathbf{u}_k^{(i+1)} = \mathbf{u}_k^{(i)} - \eta \nabla_{\mathbf{u}_k} J \qquad (5.218) $$

is recovered, which is the standard gradient descent update used in backpropagation-based training. Thus, the MSA algorithm, with a single gradient step for the control update, is mathematically equivalent to the training of a deep neural network via forward propagation, backpropagation of errors, and a gradient descent step on the parameters, providing a rigorous optimal control interpretation of the entire learning process. The iteration continues until a convergence criterion on the control sequence or the total cost is met, yielding a solution that approximates the PMP necessary conditions.
### 5.5.3 Differentiable Optimization Layers

Differentiable Optimization Layers represent a profound integration of mathematical optimization into the fabric of deep neural networks, where the output of a layer is defined as the solution to a constrained optimization problem. The forward pass of such a layer, at a given depth $k$, computes the state $\mathbf{x}_{k+1}$ by solving a parameterized problem:

$$ \mathbf{x}_{k+1} = \arg \min_{\mathbf{z} \in \mathcal{C}(\mathbf{x}_k, \mathbf{u}_k)} \mathcal{L}_k(\mathbf{z}, \mathbf{x}_k, \mathbf{u}_k) \qquad (5.219) $$

where $\mathcal{L}_k$ is a cost function internal to the layer, $\mathcal{C}$ is a constraint set defined by parameters $\mathbf{x}_k$ (the previous state) and $\mathbf{u}_k$ (the layer's control parameters or weights), and $\mathbf{z}$ is the optimization variable. This formulation generalizes standard layers; for example, a linear layer followed by a ReLU activation can be interpreted as solving a quadratic program with non-negativity constraints. The primary challenge lies in differentiating through this argmin operation to compute the Jacobians $\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{x}_k}$ and $\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{u}_k}$ required for the backward pass during training. This is not a standard application of the chain rule because the mapping from parameters to solution is implicitly defined by the optimality conditions of the optimization problem.

The rigorous treatment of backpropagation through these layers relies on the implicit function theorem applied to the Karush-Kuhn-Tucker (KKT) conditions. For a convex optimization problem with inequality constraints, the Lagrangian is

$$ \mathcal{L}_{k}^{\text{Lagrange}}(\mathbf{z}, \mathbf{\lambda}, \mathbf{\nu}; \mathbf{x}_{k}, \mathbf{u}_{k}) = \mathcal{L}_{k}(\mathbf{z}, \mathbf{x}_{k}, \mathbf{u}_{k}) + \mathbf{\lambda}^{\top}\mathbf{g}(\mathbf{z}, \mathbf{x}_{k}, \mathbf{u}_{k}) + \mathbf{\nu}^{\top}\mathbf{h}(\mathbf{z}, \mathbf{x}_{k}, \mathbf{u}_{k}) \qquad (5.220) $$

where $\mathbf{g}(\mathbf{z}, \mathbf{x}_k, \mathbf{u}_k) \le 0$ represents the inequality constraints, $\mathbf{h}(\mathbf{z}, \mathbf{x}_k, \mathbf{u}_k) = 0$ represents the equality constraints, and $\mathbf{\lambda} \ge 0$, $\mathbf{\nu}$ are the dual variables. The KKT optimality conditions, which must hold at the solution $(\mathbf{x}_{k+1}, \mathbf{\lambda}^*, \mathbf{\nu}^*)$, are stationarity:

$$ \nabla_{\mathbf{z}} \mathcal{L}_{k}^{\text{Lagrange}} = \nabla_{\mathbf{z}} \mathcal{L}_{k}(\mathbf{x}_{k+1}, \mathbf{x}_{k}, \mathbf{u}_{k}) + \left( \frac{\partial \mathbf{g}(\mathbf{x}_{k+1}, \mathbf{x}_{k}, \mathbf{u}_{k})}{\partial \mathbf{z}} \right)^{\top} \mathbf{\lambda}^{*} + \left( \frac{\partial \mathbf{h}(\mathbf{x}_{k+1}, \mathbf{x}_{k}, \mathbf{u}_{k})}{\partial \mathbf{z}} \right)^{\top} \mathbf{\nu}^{*} = \mathbf{0} \qquad (5.221) $$

primal feasibility:

$$ \mathbf{g}(\mathbf{x}_{k+1}, \mathbf{x}_{k}, \mathbf{u}_{k}) \le 0, \quad \mathbf{h}(\mathbf{x}_{k+1}, \mathbf{x}_{k}, \mathbf{u}_{k}) = 0 \qquad (5.222) $$

dual feasibility:

$$ \mathbf{\lambda}^* \ge 0 \qquad (5.223) $$

and complementary slackness:

$$ \mathbf{\lambda}^* \odot \mathbf{g}(\mathbf{x}_{k+1}, \mathbf{x}_{k}, \mathbf{u}_{k}) = \mathbf{0} \qquad (5.224) $$

where $\odot$ denotes the Hadamard product. The implicit function theorem states that if the Jacobian of the system of equations formed by the stationarity and active constraints with respect to the primal and dual variables $(\mathbf{z}, \mathbf{\lambda}, \mathbf{\nu})$ is invertible at the solution, then the solution and the active dual variables are differentiable functions of the parameters $(\mathbf{x}_k, \mathbf{u}_k)$.

To compute the necessary gradients for the backward pass, we consider the total derivative of the loss function $\ell$ with respect to the parameters. Let $\mathbf{p}_{k+1} = -\frac{\partial \ell}{\partial \mathbf{x}_{k+1}}$ be the incoming gradient from the subsequent layer, which is the negative costate from the PMP viewpoint. The goal is to compute

$$ \frac{d\ell}{d\mathbf{x}_k} = \left(\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{x}_k}\right)^{\top} \frac{\partial \ell}{\partial \mathbf{x}_{k+1}} = -\left(\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{x}_k}\right)^{\top} \mathbf{p}_{k+1} \qquad (5.225) $$

and

$$ \frac{d\ell}{d\mathbf{u}_k} = \left(\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{u}_k}\right)^{\top} \frac{\partial \ell}{\partial \mathbf{x}_{k+1}} = -\left(\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{u}_k}\right)^{\top} \mathbf{p}_{k+1} \qquad (5.226) $$
Differentiating the system of KKT conditions that hold at the optimum with respect to the pa- rameters provides a linear system to solve for these Jacobians. Let $\mathbf{y} = (\mathbf{z}, \boldsymbol{\lambda}_A, \boldsymbol{\nu})$ represent the concatenated vector of the primal variable and the dual variables for the active constraints (where $A$ is the set of indices for which $g_i(\mathbf{x}_{k+1}, \mathbf{x}_k, \mathbf{u}_k) = 0$ and $\lambda_i > 0$). The active KKT conditions form a system

$$ \mathbf{F}(\mathbf{y}, \mathbf{x}_k, \mathbf{u}_k) = \mathbf{0} \qquad (5.227) $$

which includes the stationarity condition and the active equality constraints

$$ \mathbf{g}_A(\mathbf{z}, \mathbf{x}_k, \mathbf{u}_k) = \mathbf{0} \quad \text{and} \quad \mathbf{h}(\mathbf{z}, \mathbf{x}_k, \mathbf{u}_k) = \mathbf{0} \qquad (5.228) $$

The total derivative is

$$ \frac{d\mathbf{F}}{d\mathbf{x}_k} = \frac{\partial \mathbf{F}}{\partial \mathbf{y}} \frac{\partial \mathbf{y}}{\partial \mathbf{x}_k} + \frac{\partial \mathbf{F}}{\partial \mathbf{x}_k} = \mathbf{0} \qquad (5.229) $$

which yields

$$ \frac{\partial \mathbf{y}}{\partial \mathbf{x}_k} = - \left( \frac{\partial \mathbf{F}}{\partial \mathbf{y}} \right)^{-1} \frac{\partial \mathbf{F}}{\partial \mathbf{x}_k} \qquad (5.230) $$

The Jacobian $\frac{\partial \mathbf{F}}{\partial \mathbf{y}}$ is the KKT matrix at the optimum. The desired Jacobian $\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{x}_k}$ is a block of $\frac{\partial \mathbf{y}}{\partial \mathbf{x}_k}$.

A more efficient and direct method, which avoids explicitly forming and inverting the full KKT matrix, involves solving a linear system that yields the gradient directly. The differential of the stationarity condition is

$$ d(\nabla_{\mathbf{z}} \mathcal{L}_k + \mathbf{J}_{\mathbf{g}}^{\top} \boldsymbol{\lambda} + \mathbf{J}_{\mathbf{h}}^{\top} \boldsymbol{\nu}) = \mathbf{0} \qquad (5.231) $$

where $\mathbf{J}_g = \frac{\partial \mathbf{g}}{\partial \mathbf{z}}$ and $\mathbf{J}_h = \frac{\partial \mathbf{h}}{\partial \mathbf{z}}$. Expanding this differential with respect to perturbations $d\mathbf{x}_k, d\mathbf{u}_k$, and the resulting changes $d\mathbf{z}, d\boldsymbol{\lambda}_A, d\boldsymbol{\nu}$ gives:

$$ \mathbf{H}d\mathbf{z} + \mathbf{J}_{\mathbf{g}}^{\top}d\boldsymbol{\lambda}_{A} + \mathbf{J}_{\mathbf{h}}^{\top}d\boldsymbol{\nu} + \frac{\partial}{\partial \mathbf{x}_k}(\nabla_{\mathbf{z}}\mathcal{L}_k)d\mathbf{x}_k + \frac{\partial}{\partial \mathbf{x}_k}(\mathbf{J}_{\mathbf{g}}^{\top}\boldsymbol{\lambda})d\mathbf{x}_k + \frac{\partial}{\partial \mathbf{x}_k}(\mathbf{J}_{\mathbf{h}}^{\top}\boldsymbol{\nu})d\mathbf{x}_k = \mathbf{0} \qquad (5.232) $$

where

$$ \mathbf{H} = \frac{\partial^2 \mathcal{L}_k^{\text{Lagrange}}}{\partial \mathbf{z}^2} \qquad (5.233) $$

is the Hessian of the Lagrangian with respect to $\mathbf{z}$. Simultaneously, the differentials of the active constraints must be zero:

$$ d\mathbf{g}_A = \mathbf{J}_{gA}d\mathbf{z} + \frac{\partial \mathbf{g}_A}{\partial \mathbf{x}_k}d\mathbf{x}_k = \mathbf{0} \qquad (5.234) $$

and

$$ d\mathbf{h} = \mathbf{J}_{h}d\mathbf{z} + \frac{\partial \mathbf{h}}{\partial \mathbf{x}_k}d\mathbf{x}_k = \mathbf{0} \qquad (5.235) $$

This can be written as a linear system:

$$ \begin{bmatrix} \mathbf{H} & \mathbf{J}_{gA}^{\top} & \mathbf{J}_{h}^{\top} \\ \mathbf{J}_{gA} & \mathbf{0} & \mathbf{0} \\ \mathbf{J}_{h} & \mathbf{0} & \mathbf{0} \end{bmatrix} \begin{bmatrix} d\mathbf{z} \\ d\boldsymbol{\lambda}_A \\ d\boldsymbol{\nu} \end{bmatrix} = - \begin{bmatrix} \frac{\partial}{\partial \mathbf{x}_k} (\nabla_{\mathbf{z}} \mathcal{L}_k^{\text{Lagrange}}) \\ \frac{\partial \mathbf{g}_A}{\partial \mathbf{x}_k} \\ \frac{\partial \mathbf{h}}{\partial \mathbf{x}_k} \end{bmatrix} d\mathbf{x}_k. \qquad (5.236) $$

To compute the vector-Jacobian product $(\frac{\partial \mathbf{x}_{k+1}}{\partial \mathbf{x}_k})^{\top} \mathbf{p}_{k+1}$, one sets $d\mathbf{x}_k$ to be the identity in the direction of interest and solves the adjoint of this system. The standard approach is to solve for the adjoint variables $(\bar{\mathbf{z}}, \bar{\boldsymbol{\lambda}}_A, \bar{\boldsymbol{\nu}})$ such that:

$$ \begin{bmatrix} \mathbf{H} & \mathbf{J}_{gA}^{\top} & \mathbf{J}_{h}^{\top} \\ \mathbf{J}_{gA} & \mathbf{0} & \mathbf{0} \\ \mathbf{J}_{h} & \mathbf{0} & \mathbf{0} \end{bmatrix} \begin{bmatrix} \bar{\mathbf{z}} \\ \bar{\boldsymbol{\lambda}}_A \\ \bar{\boldsymbol{\nu}} \end{bmatrix} = - \begin{bmatrix} \mathbf{p}_{k+1} \\ \mathbf{0} \\ \mathbf{0} \end{bmatrix}. \qquad (5.237) $$
The solution $\bar{z}$ to this adjoint system then gives the desired gradient with respect to the previous state:

$$ \frac{d\ell}{dx_k} = \bar{z}^\top \qquad (5.238) $$

An analogous procedure, solving a similar adjoint system but with the right-hand side terms derived from differentiation with respect to $u_k$, yields the gradient $\frac{d\ell}{du_k}$. This formalism provides a complete and rigorous mechanism for embedding optimization problems as differentiable layers within a deep network, allowing for end-to-end training using gradient-based methods while maintaining a direct interpretation as an optimal control step within the network's dynamics.

### 5.5.4 Neural Ordinary Differential Equations (Neural ODEs)

Neural Ordinary Differential Equations (Neural ODEs) formulate the forward propagation through a deep neural network as the solution to a continuous-time dynamical system defined by an ordinary differential equation. The state of the network, represented by a vector $\mathbf{x}(t) \in \mathbb{R}^n$, evolves from an initial time $t_0$ to a final time $t_1$ according to the differential equation

$$ \frac{dx(t)}{dt} = f(x(t), t, \theta) \qquad (5.239) $$

where $f$ is a function parameterized by the network weights $\theta$, and the initial condition is the input data $\mathbf{x}(t_0) = \mathbf{x}_0$. The output of the network is the state at the terminal time, $\mathbf{x}(t_1)$. This formulation generalizes residual networks, where each layer performs an update

$$ \mathbf{x}_{k+1} = \mathbf{x}_k + \mathbf{f}_k(\mathbf{x}_k, \theta_k) \qquad (5.240) $$

to the continuous limit as the number of layers approaches infinity and the step size approaches zero. The fundamental problem of training such a network is an optimal control problem, where the goal is to find the control parameters $\theta$ that minimize a loss function $\ell(\mathbf{x}(t_1))$ subject to the dynamical constraint of the ODE. The cost functional is

$$ J(\theta) = \ell(\mathbf{x}(t_1)) + \int_{t_0}^{t_1} \mathcal{L}(\mathbf{x}(t), \theta, t) dt \qquad (5.241) $$

where $\mathcal{L}$ is a running cost, often a regularization term on the weights or the state.

The necessary conditions for optimality are derived from the calculus of variations and Pontryagin's Maximum Principle in continuous time. The Hamiltonian for this system is defined as

$$ H(\mathbf{x}, \mathbf{p}, \theta, t) = \mathbf{p}^\top \mathbf{f}(\mathbf{x}, t, \theta) - \mathcal{L}(\mathbf{x}, \theta, t) \qquad (5.242) $$

where $\mathbf{p}(t) \in \mathbb{R}^n$ is the costate variable. The PMP states that for an optimal control $\theta^*$ and the corresponding optimal state trajectory $\mathbf{x}^*(t)$, there exists a costate trajectory $\mathbf{p}^*(t)$ satisfying the canonical equations. The state equation is the forward ODE:

$$ \frac{d\mathbf{x}^*(t)}{dt} = \nabla_{\mathbf{p}} H = \mathbf{f}(\mathbf{x}^*(t), t, \theta^*) \qquad (5.243) $$

The costate equation is the backward ODE:

$$ \frac{d\mathbf{p}^*(t)}{dt} = -\nabla_{\mathbf{x}} H = -\left(\frac{\partial \mathbf{f}(\mathbf{x}^*(t), t, \theta^*)}{\partial \mathbf{x}}\right)^\top \mathbf{p}^*(t) + \nabla_{\mathbf{x}} \mathcal{L}(\mathbf{x}^*(t), \theta^*, t) \qquad (5.244) $$

with the terminal boundary condition

$$ \mathbf{p}^*(t_1) = -\nabla_{\mathbf{x}} \ell(\mathbf{x}^*(t_1)) \qquad (5.245) $$
The optimality condition for the control requires that the Hamiltonian is maximized, leading to

$$ \nabla_{\theta} H(\mathbf{x}^*(t), \mathbf{p}^*(t), \mathbf{\theta}^*, t) = \left( \frac{\partial \mathbf{f}(\mathbf{x}^*(t), t, \mathbf{\theta}^*)}{\partial \mathbf{\theta}} \right)^{\top} \mathbf{p}^*(t) - \nabla_{\theta} \mathcal{L}(\mathbf{x}^*(t), \mathbf{\theta}^*, t) = \mathbf{0} \quad (5.246) $$

for all $t \in [t_0, t_1]$.

The central innovation for training Neural ODEs is the adjoint sensitivity method, which provides a memory-efficient way to compute the gradient of the loss with respect to the parameters, $\frac{dJ}{d\theta}$, by solving a backward ODE. The adjoint state is defined as $\mathbf{a}(t) = \frac{d\mathbf{f}(\mathbf{x}(t), t, \mathbf{\theta}^*)}{d\mathbf{x}(t)}$, which is exactly the negative of the costate, $\mathbf{a}(t) = -\mathbf{p}(t)$. Its dynamics are derived by differentiating the loss with respect to time. Consider the total derivative

$$ \frac{d\ell}{dt} = \left( \frac{d\mathbf{x}(t)}{dt} \right)^{\top} \frac{d\ell}{d\mathbf{x}(t)} = \mathbf{f}(\mathbf{x}(t), t, \mathbf{\theta})^{\top} \mathbf{a}(t) \quad (5.247) $$

However, $\ell$ is a constant function of $t$, so $\frac{d\ell}{dt} = 0$. A more rigorous derivation considers the gradient

$$ \frac{d\ell}{d\theta} = \int_{t_0}^{t_1} \frac{\partial \mathbf{f}(\mathbf{x}(t), t, \mathbf{\theta})^{\top}}{\partial \mathbf{\theta}} \mathbf{p}(t) dt \quad (5.248) $$

from the PMP and constructs an augmented state

$$ \mathbf{s}(t) = [\mathbf{a}(t), \frac{d\ell}{d\theta}(t)] \quad (5.249) $$

whose dynamics can be computed backwards. The definitive result is that the adjoint state $\mathbf{a}(t)$ satisfies the differential equation

$$ \frac{d\mathbf{a}(t)}{dt} = - \left( \frac{\partial \mathbf{f}(\mathbf{x}(t), t, \mathbf{\theta})}{\partial \mathbf{x}} \right)^{\top} \mathbf{a}(t) \quad (5.250) $$

with the initial condition

$$ \mathbf{a}(t_1) = \nabla_{\mathbf{x}} \ell(\mathbf{x}(t_1)) \quad (5.251) $$

To compute the total gradient $\frac{d\ell}{d\theta}$, an auxiliary variable $\mathbf{g}(t) = \frac{d\ell}{d\theta}(t)$ is introduced, which evolves as

$$ \frac{d\mathbf{g}(t)}{dt} = - \left( \frac{\partial \mathbf{f}(\mathbf{x}(t), t, \mathbf{\theta})}{\partial \mathbf{\theta}} \right)^{\top} \mathbf{a}(t) \quad (5.252) $$

with $\mathbf{g}(t_1) = \mathbf{0}$. The total gradient is then

$$ \frac{d\ell}{d\theta} = \int_{t_1}^{t_0} \frac{d\mathbf{g}(t)}{dt} dt = \mathbf{g}(t_0) - \mathbf{g}(t_1) = \mathbf{g}(t_0) \quad (5.253) $$

In practice, the forward pass is performed by numerically integrating the ODE

$$ \frac{d\mathbf{x}(t)}{dt} = \mathbf{f}(\mathbf{x}(t), t, \mathbf{\theta}) \quad (5.254) $$

from $t_0$ to $t_1$ using an ODE solver to obtain the final state $\mathbf{x}(t_1)$ and the loss $\ell(\mathbf{x}(t_1))$. The backward pass does not require storing the intermediate states of the forward pass. Instead, the combined adjoint system is solved backwards in time from $t_1$ to $t_0$. The system to be solved is defined for the augmented state

$$ \mathbf{s}(t) = \left[ \mathbf{a}(t) \quad \frac{d\ell}{d\theta}(t) \right] \quad (5.255) $$
and its dynamics are given by:

$$ \frac{d}{dt} \begin{bmatrix} \mathbf{a}(t) \\ \frac{d\ell}{d\theta}(t) \end{bmatrix} = \begin{bmatrix} - \left( \frac{\partial f(\mathbf{x}(t),t,\theta)}{\partial \mathbf{x}} \right)^{\top} \mathbf{a}(t) \\ - \left( \frac{\partial f(\mathbf{x}(t),t,\theta)}{\partial \theta} \right)^{\top} \mathbf{a}(t) \end{bmatrix}. \qquad (5.256) $$

The initial conditions for this backward integration are

$$ \mathbf{a}(t_1) = \nabla_{\mathbf{x}} \ell(\mathbf{x}(t_1)) \qquad (5.257) $$

and

$$ \frac{d\ell}{d\theta}(t_1) = \mathbf{0} \qquad (5.258) $$

Crucially, to compute the right-hand side of this adjoint ODE, the state trajectory $\mathbf{x}(t)$ is required. It can be recomputed backwards from $\mathbf{x}(t_1)$ by solving the original ODE backwards, or more efficiently, it can be reconstructed during the backward pass if a reversible ODE solver is used, or by using a continuous-time checkpoints and interpolation based on the solver's internal steps. The final value $\frac{d\ell}{d\theta}(t_0)$ provides the necessary gradient for updating the parameters $\theta$ via gradient descent. This entire procedure, where the backward pass is computed by solving a differential equation, is a direct and rigorous implementation of the continuous-time PMP, establishing Neural ODE training as a canonical example of optimal control in deep learning. The absence of a discrete computational graph for backpropagation allows for adaptive computation and constant memory cost relative to depth, as the memory requirement is $O(1)$ instead of $O(L)$ for a network with $L$ layers.

### 5.5.5 Stochastic Optimal Control for RL and Robust Training

Stochastic Optimal Control provides the foundational framework for Reinforcement Learning by considering dynamical systems governed by stochastic differential equations or Markov Decision Processes, where the state evolution incorporates inherent uncertainty. The continuous-time formulation models the state dynamics as

$$ d\mathbf{x}_t = \mathbf{f}(\mathbf{x}_t, \mathbf{u}_t, t)dt + \mathbf{G}(\mathbf{x}_t, \mathbf{u}_t, t)d\mathbf{w}_t \qquad (5.259) $$

where $\mathbf{x}_t \in \mathbb{R}^n$ is the state, $\mathbf{u}_t \in \mathbb{R}^m$ is the control, $\mathbf{w}_t$ is a standard Wiener process representing Brownian motion, and $\mathbf{G}$ is the diffusion matrix. The objective is to find a control policy $\pi(\mathbf{x}_t, t)$ that minimizes the expected cost

$$ J^{\pi}(\mathbf{x}, t) = \mathbb{E}_{\mathbf{x}_t=\mathbf{x}} \left[ \phi(\mathbf{x}_T) + \int_t^T \mathcal{L}(\mathbf{x}_s, \mathbf{u}_s, s) ds \right] \qquad (5.260) $$

where $\phi$ is the terminal cost and $\mathcal{L}$ is the running cost. The fundamental equation of stochastic optimal control is the Hamilton-Jacobi-Bellman equation, a nonlinear partial differential equation for the optimal value function

$$ V^*(\mathbf{x}, t) = \min_{\pi} J^{\pi}(\mathbf{x}, t) \qquad (5.261) $$

The HJB equation is derived via the dynamic programming principle and Itô's lemma, yielding

$$ -\frac{\partial V^*(\mathbf{x}, t)}{\partial t} = \min_{\mathbf{u}} \left[ \mathcal{L}(\mathbf{x}, \mathbf{u}, t) + \mathbf{f}(\mathbf{x}, \mathbf{u}, t)^{\top} \nabla_{\mathbf{x}} V^* + \frac{1}{2} \mathrm{Tr} \left( \mathbf{G}(\mathbf{x}, \mathbf{u}, t) \mathbf{G}(\mathbf{x}, \mathbf{u}, t)^{\top} \nabla_{\mathbf{x}}^2 V^* \right) \right] \qquad (5.262) $$

with the terminal condition $V^*(\mathbf{x}, T) = \phi(\mathbf{x})$. The term $\frac{1}{2} \mathrm{Tr} (\mathbf{G}\mathbf{G}^{\top}\nabla_{\mathbf{x}}^2 V^*)$ is the second-order differential operator arising from the variance of the Brownian motion, which is absent in deterministic control.
The connection to Reinforcement Learning emerges from the discrete-time, finite-horizon formulation known as a Markov Decision Process, defined by the tuple $(\mathcal{X}, \mathcal{U}, P, R, \gamma)$, where $\mathcal{X}$ is the state space, $\mathcal{U}$ is the action space, $P(\mathbf{x}_{k+1}|\mathbf{x}_k, \mathbf{u}_k)$ is the transition probability kernel, $R(\mathbf{x}_k, \mathbf{u}_k)$ is the reward function, and $\gamma \in (0, 1]$ is the discount factor. The objective is to maximize the expected discounted return

$$J^{\pi} = \mathbb{E}_{\pi,P} \left[ \sum_{k=0}^{\infty} \gamma^k R(\mathbf{x}_k, \mathbf{u}_k) \right] \qquad (5.263)$$

The optimal value function

$$V^*(\mathbf{x}) = \max_{\pi} \mathbb{E} \left[ \sum_{k=0}^{\infty} \gamma^k R(\mathbf{x}_k, \mathbf{u}_k) \mid \mathbf{x}_0 = \mathbf{x} \right] \qquad (5.264)$$

satisfies the discrete-time HJB equation, known as the Bellman optimality equation:

$$V^*(\mathbf{x}) = \max_{\mathbf{u}} \left[ R(\mathbf{x}, \mathbf{u}) + \gamma \mathbb{E}_{\mathbf{x}' \sim P(\cdot|\mathbf{x},\mathbf{u})} [V^*(\mathbf{x}')] \right] \qquad (5.265)$$

Similarly, the optimal state-action value function

$$Q^*(\mathbf{x}, \mathbf{u}) = R(\mathbf{x}, \mathbf{u}) + \gamma \mathbb{E}_{\mathbf{x}' \sim P(\cdot|\mathbf{x},\mathbf{u})} \left[ \max_{\mathbf{u}'} Q^*(\mathbf{x}', \mathbf{u}') \right] \qquad (5.266)$$

provides the basis for Q-learning algorithms. Policy iteration and value iteration are dynamic programming algorithms that solve these equations, and when the state and action spaces are large, deep neural networks are used as function approximators, leading to Deep RL. The parameters $\theta$ of a value network $V(\mathbf{x}; \theta)$ or a policy network $\pi(\mathbf{x}; \phi)$ are trained to minimize the Bellman error

$$\mathbb{E} \left[ \left( R(\mathbf{x}_k, \mathbf{u}_k) + \gamma V(\mathbf{x}_{k+1}; \mathbf{\theta}^{-}) - V(\mathbf{x}_k; \mathbf{\theta}) \right)^2 \right] \qquad (5.267)$$

or to maximize the expected return

$$J(\phi) = \mathbb{E}_{\pi_{\phi},P} \left[ \sum_{k} \gamma^k R(\mathbf{x}_k, \mathbf{u}_k) \right] \qquad (5.268)$$

using the policy gradient theorem

$$\nabla_{\phi} J(\phi) = \mathbb{E}_{\pi_{\phi}} \left[ \sum_{k} \gamma^k Q^{\pi}(\mathbf{x}_k, \mathbf{u}_k) \nabla_{\phi} \log \pi(\mathbf{u}_k | \mathbf{x}_k; \phi) \right] \qquad (5.269)$$

For robust training of deep neural networks in supervised learning, the problem is framed as a minimax game or a robust optimal control problem where the controller must perform well under worst-case perturbations. Given a training dataset $\{(\mathbf{x}_0^{(i)}, \mathbf{y}^{(i)})\}_{i=1}^N$, the dynamics of a deterministic network are

$$\mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \qquad (5.270)$$

and the objective is to minimize the loss $\ell(\mathbf{x}_L, \mathbf{y})$. In the presence of adversarial perturbations $\delta$ applied to the input or intermediate features, the robust training problem is

$$\min_{\mathbf{u}} \max_{\delta \in \Delta} \frac{1}{N} \sum_{i=1}^{N} \ell(\mathbf{x}_{L}^{(i)}(\delta), \mathbf{y}^{(i)}) \qquad (5.271)$$

where $\Delta$ is a bounded set of admissible perturbations, such as $\|\delta\|_p \le \epsilon$. This can be viewed as a zero-sum dynamic game where nature chooses the worst-case disturbance. The Hamiltonian for this robust control problem incorporates the disturbance

$$H_k(\mathbf{x}_k, \mathbf{u}_k, \mathbf{p}_{k+1}, \delta_k) = \mathbf{p}_{k+1}^{\top} \mathbf{f}_k(\mathbf{x}_k + \delta_k, \mathbf{u}_k) \qquad (5.272)$$
and the robust optimality conditions involve a maximization over the disturbance followed by a minimization over the control. Solving the inner maximization exactly is often intractable, leading to approximations like Projected Gradient Descent to find adversarial examples

$$ \mathbf{x}_0^{adv} = \mathbf{x}_0 + \arg \max_{\|\boldsymbol{\delta}\| \le \epsilon} \ell(\mathbf{x}_L(\boldsymbol{\delta}), \mathbf{y}) \quad (5.273) $$

which are then used in training, a method known as Adversarial Training.

A more formal approach uses stochastic control theory with uncertain dynamics. The dynamics are modeled as

$$ \mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) + \boldsymbol{\sigma}_k(\mathbf{x}_k, \mathbf{u}_k)\boldsymbol{\xi}_k \quad (5.274) $$

where $\boldsymbol{\xi}_k$ is a random noise vector. The goal is to find a control policy that minimizes the expected loss while being robust to the noise distribution. This leads to risk-sensitive optimal control, which optimizes a risk-measure of the cost, such as the exponential utility

$$ J^{RS}(\mathbf{u}) = \frac{1}{\beta} \log \mathbb{E}[\exp(\beta \ell(\mathbf{x}_L))] \quad (5.275) $$

where $\beta > 0$ is a risk-aversion parameter. Using the Taylor expansion

$$ \frac{1}{\beta} \log \mathbb{E}[e^{\beta Z}] = \mathbb{E}[Z] + \frac{\beta}{2} \mathrm{Var}[Z] + O(\beta^2) \quad (5.276) $$

it is clear this penalizes high variance in the loss. The Path Integral Control framework provides a solution for a specific class of stochastic problems where the dynamics are affine in the control and the cost is quadratic in the control. For a system with dynamics

$$ d\mathbf{x}_t = (\mathbf{f}(\mathbf{x}_t, t) + \mathbf{G}(\mathbf{x}_t, t)\mathbf{u}_t)dt + \mathbf{G}(\mathbf{x}_t, t)d\mathbf{w}_t \quad (5.277) $$

and cost

$$ J = \mathbb{E} \left[ \phi(\mathbf{x}_T) + \int_0^T \left( \mathcal{L}(\mathbf{x}_t, t) + \frac{1}{2} \mathbf{u}_t^\top \mathbf{R} \mathbf{u}_t \right) dt \right] \quad (5.278) $$

the HJB equation can be linearized via a logarithmic transformation

$$ V(\mathbf{x}, t) = -\lambda \log \Psi(\mathbf{x}, t) \quad (5.279) $$

leading to a linear PDE for $\Psi$. The optimal control is then given by

$$ \mathbf{u}^*(\mathbf{x}, t) = \mathbf{R}^{-1} \mathbf{G}(\mathbf{x}, t)^\top \frac{\nabla_{\mathbf{x}} \Psi(\mathbf{x}, t)}{\Psi(\mathbf{x}, t)} \quad (5.280) $$

which can be estimated by sampling trajectories, forming the basis of the Model Predictive Path Integral (MPPI) control algorithm used in RL for policy improvement. This direct sampling-based approach bypasses the need for value function approximation and is highly suitable for complex, high-dimensional problems where DNNs represent the policy.

### 5.5.6 Stable and Efficient Network Design

The design of stable and efficient neural networks from an optimal control perspective necessitates the enforcement of structural constraints on the network's dynamics to guarantee desirable properties such as Lyapunov stability, passivity, or contractivity, which in turn ensure robustness to input perturbations and prevent exploding gradients during training. Consider the discrete-time dynamical system representation of a neural network

$$ \mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \quad (5.281) $$
where $\mathbf{x}_k \in \mathbb{R}^n$ is the state and $\mathbf{u}_k \in \mathbb{R}^m$ are the parameters. A fundamental stability concept is that of a contraction mapping, which requires the existence of a matrix measure $\mu$ such that the Lipschitz constant of each layer is strictly less than one, $\text{Lip}(\mathbf{f}_k) < 1$, ensuring that the distance between any two trajectories contracts over time. A more general and powerful framework is provided by Lyapunov theory. A Lyapunov function $V: \mathbb{R}^n \to \mathbb{R}_{\ge 0}$ for the autonomous system

$$ \mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k) \qquad (5.282) $$

must satisfy $V(\mathbf{x}) > 0$ for all $\mathbf{x} \neq \mathbf{0}$ and $V(\mathbf{0}) = 0$, and the condition

$$ V(\mathbf{x}_{k+1}) - V(\mathbf{x}_k) \leq -\alpha V(\mathbf{x}_k) \qquad (5.283) $$

for some $\alpha > 0$, guaranteeing global exponential stability of the equilibrium point. To design layers that inherently satisfy such properties, one can parameterize the layer $\mathbf{f}_k$ as the solution to a constrained optimization problem or as the forward Euler discretization of a continuous-time system with known stability properties.

A prominent approach is to design layers that are nonexpansive, meaning they are 1-Lipschitz continuous

$$ \|\mathbf{f}_k(\mathbf{x}) - \mathbf{f}_k(\mathbf{y})\| \leq \|\mathbf{x} - \mathbf{y}\| \qquad (5.284) $$

for all $\mathbf{x}, \mathbf{y}$ which prevents the amplification of input perturbations and gradient norms. This can be achieved by constraining the weights to be orthogonal or using spectral normalization, which enforces

$$ \sigma_{\max}(\nabla_{\mathbf{x}} \mathbf{f}_k) \leq 1 \qquad (5.285) $$

where $\sigma_{\max}$ denotes the maximum singular value. A more rigorous method involves modeling the network as the discretization of a continuous-time Ordinary Differential Equation

$$ \frac{d\mathbf{x}(t)}{dt} = \mathbf{g}(\mathbf{x}(t), \boldsymbol{\theta}(t)) \qquad (5.286) $$

and ensuring the continuous dynamics are contractive. A nonlinear system is contractive if there exists a uniformly positive definite metric

$$ \mathbf{M}(\mathbf{x}, t) = \boldsymbol{\Theta}(\mathbf{x}, t)^{\top} \boldsymbol{\Theta}(\mathbf{x}, t) \qquad (5.287) $$

such that the generalized Jacobian

$$ \mathbf{F} = \left( \dot{\boldsymbol{\Theta}} + \boldsymbol{\Theta} \frac{\partial \mathbf{g}}{\partial \mathbf{x}} \right) \boldsymbol{\Theta}^{-1} \qquad (5.288) $$

is uniformly negative definite, i.e.,

$$ \frac{1}{2}(\mathbf{F} + \mathbf{F}^{\top}) \preceq -\lambda \mathbf{I} \qquad (5.289) $$

for some $\lambda > 0$. A sufficient condition for contractivity in the identity metric $\mathbf{M} = \mathbf{I}$ is that the symmetric part of the Jacobian is uniformly negative definite

$$ \frac{1}{2} \left( \frac{\partial \mathbf{g}}{\partial \mathbf{x}} + \left( \frac{\partial \mathbf{g}}{\partial \mathbf{x}} \right)^{\top} \right) \preceq -\lambda \mathbf{I} \qquad (5.290) $$

This insight leads to the design of stable continuous-time blocks, such as

$$ \mathbf{g}(\mathbf{x}) = -\mathbf{W}^{\top} \boldsymbol{\sigma}(\mathbf{W}\mathbf{x} + \mathbf{b}) \qquad (5.291) $$

where $\boldsymbol{\sigma}$ is a monotone activation function, which can be shown to be contracting under certain conditions on $\mathbf{W}$.
The training of such stable architectures can be framed as an optimal control problem with stability constraints. The objective is to minimize the loss

$$J(\mathbf{u}) = \ell(\mathbf{x}_L) \qquad (5.292)$$

subject to the dynamics

$$\mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \qquad (5.293)$$

and the algebraic stability constraints $\mathbf{h}_k(\mathbf{u}_k) \le 0$, where $\mathbf{h}_k$ encodes constraints like spectral norm bounds. The Lagrangian for this constrained optimization is

$$\mathcal{L}(\mathbf{x}, \mathbf{u}, \boldsymbol{\lambda}) = \ell(\mathbf{x}_L) + \sum_{k=0}^{L-1} \mathbf{p}_{k+1}^{\top} (\mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) - \mathbf{x}_{k+1}) + \sum_{k=0}^{L-1} \boldsymbol{\lambda}_k^{\top} \mathbf{h}_k(\mathbf{u}_k) \qquad (5.294)$$

where $\boldsymbol{\lambda}_k \ge 0$ are the Lagrange multipliers for the stability constraints. The necessary conditions from the Karush-Kuhn-Tucker theorem include the state equation

$$\mathbf{x}_{k+1} = \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k) \qquad (5.295)$$

the costate equation

$$\mathbf{p}_k = \left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k)}{\partial \mathbf{x}_k} \right)^{\top} \mathbf{p}_{k+1} \qquad (5.296)$$

with

$$\mathbf{p}_L = \nabla_{\mathbf{x}_L} \ell(\mathbf{x}_L) \qquad (5.297)$$

and the stationarity condition with respect to the controls

$$\left( \frac{\partial \mathbf{f}_k(\mathbf{x}_k, \mathbf{u}_k)}{\partial \mathbf{u}_k} \right)^{\top} \mathbf{p}_{k+1} + \left( \frac{\partial \mathbf{h}_k(\mathbf{u}_k)}{\partial \mathbf{u}_k} \right)^{\top} \boldsymbol{\lambda}_k = \mathbf{0} \qquad (5.298)$$

complemented by the complementary slackness conditions

$$\boldsymbol{\lambda}_k \odot \mathbf{h}_k(\mathbf{u}_k) = \mathbf{0} \qquad (5.299)$$

In practice, the stability constraints $\mathbf{h}_k(\mathbf{u}_k) \le 0$ are often enforced via projected gradient descent or by reparameterizing the weights. For example, one can parameterize a weight matrix $\mathbf{W}_k$ as

$$\mathbf{W}_k = \gamma_k \tilde{\mathbf{W}}_k \qquad (5.300)$$

where $\tilde{\mathbf{W}}_k$ has its spectral norm constrained to 1 via iterative power methods, and $\gamma_k \le 1$ is a trainable scalar, ensuring

$$\sigma_{\max}(\mathbf{W}_k) \le \gamma_k \le 1 \qquad (5.301)$$

For recurrent networks and those with skip connections, more advanced concepts from robust control, such as the use of Linear Matrix Inequalities (LMIs), are required to guarantee stability. Consider a system with uncertain or time-varying parameters,

$$\mathbf{x}_{k+1} = \mathbf{A}_k \mathbf{x}_k + \mathbf{B}_k \mathbf{u}_k \qquad (5.302)$$

where $[\mathbf{A}_k \mathbf{B}_k] \in \Omega$, a bounded uncertainty set. The system is quadratically stable if there exists a matrix $\mathbf{P} \succ 0$ such that the LMI

$$\mathbf{A}^{\top} \mathbf{P} \mathbf{A} - \mathbf{P} \prec 0 \qquad (5.303)$$

holds for all $[\mathbf{A} \mathbf{B}] \in \Omega$. This can be extended to nonlinear systems via sector bounds on the activation functions. Let $\phi(\mathbf{z})$ be a slope-restricted activation function in the sector $[\alpha, \beta]$, meaning

$$(\phi(\mathbf{z}) - \alpha \mathbf{z})^{\top} (\beta \mathbf{z} - \phi(\mathbf{z})) \ge 0 \qquad (5.304)$$
for all $\mathbf{z}$. A network layer

$$ \mathbf{x}_{k+1} = \phi(\mathbf{W}\mathbf{x}_k + \mathbf{b}) \qquad (5.305) $$

can be analyzed using the discrete-time version of the Kalman-Yakubovich-Popov lemma. Stability can be certified by finding a diagonal positive definite matrix $\mathbf{D}$ and a matrix $\mathbf{P} > 0$ such that the following LMI is feasible:

$$ \begin{bmatrix} \mathbf{W}^{\top}\mathbf{P}\mathbf{W} - \mathbf{P} - 2\alpha\mathbf{W}^{\top}\mathbf{D} & \mathbf{W}^{\top}\mathbf{D} + \alpha\beta\mathbf{W}^{\top}\mathbf{D} \\ \mathbf{D}\mathbf{W} + \alpha\beta\mathbf{D}\mathbf{W} & -2\mathbf{D} \end{bmatrix} \preceq 0. \qquad (5.306) $$

Training a network to satisfy such LMIs during optimization is challenging but can be approximated by incorporating a regularization term that penalizes the violation of the stability condition, leading to a robust and stable network architecture whose training dynamics are themselves well-behaved, mitigating issues like vanishing and exploding gradients and enhancing generalization performance by enforcing a smooth input-output map. The integration of these control-theoretic stability constraints directly into the architecture and the training loop represents a principled approach to designing deep neural networks with predictable and reliable behavior, which is paramount for their deployment in safety-critical reinforcement learning applications.

### 5.5.7 Optimal Control of the Training Process Itself

The optimal control of the training process itself constitutes a meta-optimization problem where the parameters of the learning algorithm, such as the learning rate, momentum coefficients, or even the entire update rule, become the control variables in a higher-level dynamical system whose state is the entire set of model parameters and optimizer states. Let $\theta_t \in \mathbb{R}^\Theta$ represent the parameters of the deep neural network at optimization step $t$, and let $\mathbf{s}_t \in \mathbb{R}^S$ represent the state of the optimizer, which may include momentum vectors $\mathbf{m}_t$, variance estimates $\mathbf{v}_t$, and other auxiliary variables. The combined state is

$$ \mathbf{X}_t = [\theta_t^\top, \mathbf{s}_t^\top]^\top \qquad (5.307) $$

The update rule of the optimizer defines the dynamics of this meta-system. For a generic optimizer, the dynamics can be expressed as

$$ \mathbf{X}_{t+1} = \mathbf{F}(\mathbf{X}_t, \alpha_t, \mathbf{g}_t) \qquad (5.308) $$

where $\alpha_t \in \mathbb{R}^A$ is a vector of hyperparameters (e.g., learning rate, momentum decay factor) acting as the control signal, and

$$ \mathbf{g}_t = \nabla_{\theta} \mathcal{L}_t(\theta_t) \qquad (5.309) $$

is the stochastic gradient computed on a minibatch at step $t$. The fundamental challenge is that the gradient $\mathbf{g}_t$ is not an exogenous input but a function of the state $\theta_t$ and the data distribution, making the system a closed-loop feedback system with stochastic disturbances.

The objective of the meta-optimization is to find a hyperparameter policy

$$ \alpha_t = \pi(\mathbf{X}_t, t) \qquad (5.310) $$

that minimizes the expected total loss over a horizon $T$, which is

$$ J(\pi) = \mathbb{E} \left[ \sum_{t=0}^{T} \gamma^t \mathcal{L}_t(\theta_t) \right] \qquad (5.311) $$

where $\gamma \in (0, 1]$ is a discount factor. This is a classic stochastic optimal control problem. The corresponding Bellman equation for the optimal meta-value function $V^*(\mathbf{X})$ is

$$ V^*(\mathbf{X}_t) = \min_{\alpha_t} \mathbb{E}_{\mathbf{g}_t} [\mathcal{L}_t(\theta_t) + \gamma V^*(\mathbf{X}_{t+1}) | \mathbf{X}_t, \alpha_t] \qquad (5.312) $$
where the expectation is over the stochasticity in the gradient $g_t$. The optimal hyperparameter control is then

$$ \alpha_t^* = \arg \min_{\alpha_t} \mathbb{E}_{g_t} [\mathcal{L}_t(\theta_t) + \gamma V^*(\mathbf{F}(\mathbf{X}_t, \alpha_t, g_t))] \quad (5.313) $$

Solving this exactly is intractable, but it provides a formal framework for gradient-based hyper-parameter optimization. By considering the hyperparameters $\alpha$ as fixed over a horizon, one can compute the hypergradient $\frac{dJ}{d\alpha}$ to perform meta-descent. This involves differentiating through the optimization process. The dynamics for a simple gradient descent optimizer are

$$ \theta_{t+1} = \theta_t - \alpha g_t \quad (5.314) $$

The hypergradient for a horizon of one step is trivial:

$$ \frac{d\mathcal{L}_{t+1}}{d\alpha} = \frac{\partial \mathcal{L}_{t+1}}{\partial \theta_{t+1}} \top \frac{d\theta_{t+1}}{d\alpha} = -g_{t+1}^\top g_t \quad (5.315) $$

For a longer horizon $T$, the hypergradient requires backpropagation through the entire training trajectory (through-time), given by

$$ \frac{d\mathcal{L}_T}{d\alpha} = \sum_{t=0}^{T-1} \frac{\partial \mathcal{L}_T}{\partial \theta_T} \frac{d\theta_T}{d\theta_t} \frac{\partial \theta_{t+1}}{\partial \alpha} \quad (5.316) $$

where

$$ \frac{d\theta_T}{d\theta_t} = \prod_{k=t}^{T-1} \frac{\partial \theta_{k+1}}{\partial \theta_k} = \prod_{k=t}^{T-1} (I - \alpha \mathbf{H}_k) \quad (5.317) $$

and

$$ \mathbf{H}_k = \nabla_{\theta}^2 \mathcal{L}_k(\theta_k) \quad (5.318) $$

is the Hessian at step $k$, and

$$ \frac{\partial \theta_{t+1}}{\partial \alpha} = -g_t \quad (5.319) $$

This product of Hessians is computationally prohibitive and numerically unstable, leading to approximations such as ignoring the higher-order terms and using the approximation

$$ \frac{d\mathcal{L}_T}{d\alpha} \approx -\sum_{t=0}^{T-1} g_T^\top g_t \quad (5.320) $$

For more complex optimizers like Momentum or Adam, the state $s_t$ and its dynamics become crucial. The momentum optimizer has states $m_t$ and dynamics

$$ \mathbf{m}_{t+1} = \beta \mathbf{m}_t + (1 - \beta) \mathbf{g}_t \quad (5.321) $$

$$ \theta_{t+1} = \theta_t - \alpha \mathbf{m}_{t+1} \quad (5.322) $$

where the controls are $\alpha_t = (\alpha_t, \beta_t)$. The hypergradient for the learning rate $\alpha$ involves differentiating through these dynamics:

$$ \frac{\partial \theta_{t+1}}{\partial \alpha} = -\mathbf{m}_{t+1} - \alpha \frac{\partial \mathbf{m}_{t+1}}{\partial \alpha} \quad (5.323) $$

where

$$ \frac{\partial \mathbf{m}_{t+1}}{\partial \alpha} = \beta \frac{\partial \mathbf{m}_t}{\partial \alpha} + (1 - \beta) \frac{\partial \mathbf{g}_t}{\partial \theta_t} \frac{\partial \theta_t}{\partial \alpha} \quad (5.324) $$

This creates a recurrent system for the hypergradient that mirrors the parameter update itself. The Adam optimizer introduces further complexity with states $m_t$, $v_t$ and dynamics

$$ \mathbf{m}_{t+1} = \beta_1 \mathbf{m}_t + (1 - \beta_1) \mathbf{g}_t \quad (5.325) $$
$$ \mathbf{v}_{t+1} = \beta_2 \mathbf{v}_t + (1 - \beta_2) \mathbf{g}_t^2 \qquad (5.326) $$

$$ \hat{\mathbf{m}}_{t+1} = \mathbf{m}_{t+1}/(1 - \beta_1^{t+1}) \qquad (5.327) $$

$$ \hat{\mathbf{v}}_{t+1} = \mathbf{v}_{t+1}/(1 - \beta_2^{t+1}) \qquad (5.328) $$

$$ \boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \alpha \hat{\mathbf{m}}_{t+1} / (\sqrt{\hat{\mathbf{v}}_{t+1}} + \epsilon) \qquad (5.329) $$

where the division and square root are element-wise. Differentiating through this requires careful application of the chain rule, and the hypergradients for $\alpha$, $\beta_1$, and $\beta_2$ can be computed recursively, albeit with significant memory and computational cost if done exactly through a long horizon.

A more general and rigorous approach formulates this as a continuous-time optimal control problem for the training trajectory. The parameters $\boldsymbol{\theta}(\tau)$ evolve during training "time" $\tau$ according to the optimizer's ODE, such as

$$ \frac{d\boldsymbol{\theta}}{d\tau} = -\alpha(\tau)\mathbf{g}(\boldsymbol{\theta}(\tau), \tau) \qquad (5.330) $$

for gradient flow. The meta-control problem is to find the time-varying learning rate $\alpha(\tau)$ that minimizes

$$ J(\alpha(\tau)) = \int_{0}^{T} \mathcal{L}(\boldsymbol{\theta}(\tau)) d\tau \qquad (5.331) $$

The Hamiltonian for this meta-problem is

$$ H(\boldsymbol{\theta}, \mathbf{p}, \alpha, \tau) = \mathbf{p}^{\top}(-\alpha\mathbf{g}(\boldsymbol{\theta}, \tau)) - \mathcal{L}(\boldsymbol{\theta}(\tau)) \qquad (5.332) $$

where $\mathbf{p}(\tau)$ is the meta-costate. The meta-costate equation is

$$ \frac{d\mathbf{p}}{d\tau} = -\nabla_{\boldsymbol{\theta}} H = \alpha (\nabla_{\boldsymbol{\theta}}^2 \mathcal{L}) \mathbf{p} + \nabla_{\boldsymbol{\theta}} \mathcal{L} \qquad (5.333) $$

with the terminal condition $\mathbf{p}(T) = 0$. The optimality condition for the control is

$$ \frac{\partial H}{\partial \alpha} = -\mathbf{p}^{\top} \mathbf{g} = 0 \qquad (5.334) $$

implying that the optimal learning rate is chosen such that the meta-costate is orthogonal to the gradient. This continuous-time formulation elegantly captures the trade-off: the term $\alpha(\nabla^2 \mathcal{L})\mathbf{p}$ in the costate dynamics represents the interaction between the learning rate and the curvature of the loss landscape, guiding how $\alpha$ should be adjusted to navigate sharp minima and flat regions effectively. This rigorous optimal control perspective provides a unifying theoretical foundation for understanding and designing adaptive learning rate schedules and other meta-optimization strategies, framing the entire training procedure of a deep neural network as a single, coherent control problem over the space of parameters.

### 5.5.8 Connections to Reinforcement Learning

The connections between optimal control and reinforcement learning are foundational and profound, with RL being essentially a framework for solving stochastic optimal control problems in the absence of a known dynamical model. The canonical formulation is the Markov Decision Process defined by the tuple $(\mathcal{S}, \mathcal{A}, P, R, \gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $P(s'|s, a)$ is the transition probability kernel, $R(s, a)$ is the reward function, and $\gamma \in (0, 1]$ is the discount factor. The objective is to find a policy $\pi(a|s)$ that maximizes the expected discounted return

$$ J(\pi) = \mathbb{E}_{\pi, P} \left[ \sum_{t=0}^{\infty} \gamma^t R(s_t, a_t) \right] \qquad (5.335) $$
The optimal value function

$$V^*(s) = \max_{\pi} \mathbb{E} \left[ \sum_{k=0}^{\infty} \gamma^k R(s_{t+k}, a_{t+k}) \mid s_t = s \right] \quad (5.336)$$

satisfies the Bellman optimality equation

$$V^*(s) = \max_{a \in \mathcal{A}} \left[ R(s, a) + \gamma \mathbb{E}_{s' \sim P(\cdot|s, a)} [V^*(s')] \right] \quad (5.337)$$

which is the discrete-time, stochastic analogue of the Hamilton-Jacobi-Bellman equation from continuous optimal control. The Hamiltonian for the discrete MDP can be defined as

$$H(s, a, V) = R(s, a) + \gamma \mathbb{E}_{s' \sim P(\cdot|s, a)} [V(s')] \quad (5.338)$$

and the Bellman equation is then

$$V^*(s) = \max_a H(s, a, V^*) \quad (5.339)$$

directly mirroring the HJB condition

$$-\frac{\partial V^*}{\partial t} = \max_u H(\mathbf{x}, \mathbf{u}, \nabla_{\mathbf{x}} V^*, t) \quad (5.340)$$

for the deterministic case.

Policy iteration and value iteration are the dynamic programming algorithms that solve the Bellman equations, and when the state and action spaces are large, function approximation is required, leading to Deep Reinforcement Learning where deep neural networks parameterize the value function $V(s; \mathbf{w})$ or the policy $\pi(a|s; \boldsymbol{\theta})$. The training of these networks is itself an optimal control problem where the parameters $\mathbf{w}$ or $\boldsymbol{\theta}$ are the controls, and the dynamics are defined by the optimization algorithm. The value network parameters $\mathbf{w}$ are trained to minimize the expected Temporal Difference error, which is derived from the Bellman equation. The TD error for a transition $(s_t, a_t, r_t, s_{t+1})$ is

$$\delta_t = r_t + \gamma V(s_{t+1}; \mathbf{w}) - V(s_t; \mathbf{w}) \quad (5.341)$$

and the mean-squared projected Bellman error objective is

$$J(\mathbf{w}) = \mathbb{E}[(\delta_t)^2] \quad (5.342)$$

The gradient of this objective is

$$\nabla_{\mathbf{w}} J(\mathbf{w}) = -2\mathbb{E}[\delta_t \nabla_{\mathbf{w}} V(s_t; \mathbf{w})] \quad (5.343)$$

which can be seen as a stochastic gradient descent update aiming to satisfy the Bellman orthogonality condition, a form of policy evaluation in the linear function approximation setting.

For policy-based methods, the policy gradient theorem provides the foundational update rule. The objective is

$$J(\boldsymbol{\theta}) = \mathbb{E}_{\tau \sim p(\tau; \boldsymbol{\theta})} [R(\tau)] = \int p(\tau; \boldsymbol{\theta}) R(\tau) d\tau \quad (5.344)$$

where $\tau = (s_0, a_0, s_1, a_1, \dots)$ is a trajectory and

$$R(\tau) = \sum_{t=0}^{\infty} \gamma^t R(s_t, a_t) \quad (5.345)$$
The policy gradient is

$$ \nabla_{\boldsymbol{\theta}} J(\boldsymbol{\theta}) = \mathbb{E}_{\tau \sim p(\tau; \boldsymbol{\theta})} \left[ \left( \sum_{t=0}^{\infty} \nabla_{\boldsymbol{\theta}} \log \pi(a_t | s_t; \boldsymbol{\theta}) \right) R(\tau) \right] \quad (5.346) $$

A more common and lower-variance form uses the state-action value function

$$ Q^{\pi}(s, a) = \mathbb{E}_{\pi} \left[ \sum_{k=0}^{\infty} \gamma^k R(s_{t+k}, a_{t+k}) \mid s_t = s, a_t = a \right] \quad (5.347) $$

yielding

$$ \nabla_{\boldsymbol{\theta}} J(\boldsymbol{\theta}) = \mathbb{E}_{s \sim d^{\pi}, a \sim \pi(\cdot|s)} [Q^{\pi}(s, a) \nabla_{\boldsymbol{\theta}} \log \pi(a|s; \boldsymbol{\theta})] \quad (5.348) $$

where $d^{\pi}(s)$ is the stationary state distribution under policy $\pi$. This gradient can be interpreted from an optimal control perspective where the policy update is a step towards maximizing a performance metric, and the Q-function acts as a costate-like variable that evaluates the quality of the action. Advanced policy gradient methods like Trust Region Policy Optimization (TRPO) and Proximal Policy Optimization (PPO) explicitly incorporate a constraint on the policy change, framing the update as a constrained optimization problem

$$ \boldsymbol{\theta}_{k+1} = \arg \max_{\boldsymbol{\theta}} \mathbb{E} \left[ \frac{\pi(a|s; \boldsymbol{\theta})}{\pi(a|s; \boldsymbol{\theta}_k)} A^{\pi_k}(s, a) \right] \quad (5.349) $$

subject to

$$ \mathbb{E}[D_{KL}(\pi(\cdot|s; \boldsymbol{\theta}_k) \| \pi(\cdot|s; \boldsymbol{\theta}))] \le \delta \quad (5.350) $$

where

$$ A^{\pi}(s, a) = Q^{\pi}(s, a) - V^{\pi}(s) \quad (5.351) $$

is the advantage function. This is a direct application of a trust-region method in the policy space, ensuring stable and monotonic policy improvement, which is a form of robust control against overly large policy updates.

The connection is further solidified in model-based reinforcement learning and planning. When a model of the environment dynamics $P(s'|s, a)$ is learned or known, optimal control methods can be directly applied for planning. The learned dynamics model $\hat{P}_{\phi}(s'|s, a)$, often parameterized by a deep neural network, allows the use of trajectory optimization techniques like Model Predictive Control (MPC) and the Linear Quadratic Regulator (LQR). For a finite horizon $H$, the goal is to find an action sequence $a_t, a_{t+1}, \dots, a_{t+H-1}$ that minimizes

$$ \mathbb{E}_{\hat{P}_{\phi}} \left[ \sum_{k=t}^{t+H-1} \gamma^{k-t} R(s_k, a_k) + \gamma^H V(s_{t+H}) \right] \quad (5.352) $$

The iLQR algorithm solves this by iteratively linearizing the dynamics around a nominal trajectory $\bar{s}_k, \bar{a}_k$ and quadratizing the cost, leading to a locally optimal linear feedback policy

$$ a_k = \bar{a}_k + \mathbf{K}_k(s_k - \bar{s}_k) + \mathbf{k}_k \quad (5.353) $$

The dynamics are linearized as

$$ \delta s_{k+1} = \mathbf{A}_k \delta s_k + \mathbf{B}_k \delta a_k \quad (5.354) $$

where

$$ \delta s_k = s_k - \bar{s}_k, \quad \delta a_k = a_k - \bar{a}_k \quad (5.355) $$

$$ \mathbf{A}_k = \nabla_s \hat{f}(\bar{s}_k, \bar{a}_k), \quad \mathbf{B}_k = \nabla_a \hat{f}(\bar{s}_k, \bar{a}_k) \quad (5.356) $$
and the cost is quadratized as

$$C(s_k, a_k) \approx \frac{1}{2} \begin{bmatrix} \delta s_k \\ \delta a_k \end{bmatrix}^\top \begin{bmatrix} C_{ss,k} & C_{sa,k} \\ C_{as,k} & C_{aa,k} \end{bmatrix} \begin{bmatrix} \delta s_k \\ \delta a_k \end{bmatrix} + \begin{bmatrix} c_{s,k} \\ c_{a,k} \end{bmatrix}^\top \begin{bmatrix} \delta s_k \\ \delta a_k \end{bmatrix} + c_k \quad (5.357)$$

The iLQR backward pass then computes the value function

$$V(s_k) = \frac{1}{2} s_k^\top S_k s_k + s_k^\top s_k + \text{const} \quad (5.358)$$

recursively from $k = H$ to $k = 0$ using the Riccati equations

$$Q_{ss,k} = C_{ss,k} + A_k^\top S_{k+1} A_k \quad (5.359)$$

$$Q_{aa,k} = C_{aa,k} + B_k^\top S_{k+1} B_k \quad (5.360)$$

$$Q_{sa,k} = C_{sa,k} + A_k^\top S_{k+1} B_k \quad (5.361)$$

$$q_{s,k} = c_{s,k} + A_k^\top s_{k+1} \quad (5.362)$$

$$q_{a,k} = c_{a,k} + B_k^\top s_{k+1} \quad (5.363)$$

and the optimal feedback gain is

$$K_k = -Q_{aa,k}^{-1} Q_{sa,k} \quad (5.364)$$

$$k_k = -Q_{aa,k}^{-1} q_{a,k} \quad (5.365)$$

with the value function parameters updated as

$$S_k = Q_{ss,k} - Q_{sa,k} Q_{aa,k}^{-1} Q_{sa,k}^\top \quad (5.366)$$

$$s_k = q_{s,k} - Q_{sa,k} Q_{aa,k}^{-1} q_{a,k} \quad (5.367)$$

This direct application of optimal control provides a powerful planning algorithm that can be used to generate targets for policy and value network training, creating a hybrid model-based and model-free RL system where the strengths of both optimal control theory and deep learning are synergistically combined.

## 5.6 Differential Game Theory and Training Dynamics

Differential game theory provides a rigorous mathematical framework for analyzing the continuous-time dynamics of competing or cooperating agents in deep learning systems, such as generative adversarial networks (GANs), multi-agent reinforcement learning, and adversarial training. Consider a game with $K$ players, where each player $k$ controls a parameter vector $\theta_k(t) \in \mathbb{R}^{d_k}$ that evolves in time according to a differential strategy. The collective state of the system is described by the joint trajectory $\Theta(t) = (\theta_1(t), \dots, \theta_K(t))$, governed by coupled ordinary differential equations (ODEs) derived from gradient-based optimization:

$$\frac{d\theta_k(t)}{dt} = -\eta_k \nabla_{\theta_k} \mathcal{L}_k(\Theta(t)), \quad k = 1, \dots, K.$$

Here, $\mathcal{L}_k$ denotes the loss function for player $k$, which may be adversarial (e.g., in GANs) or cooperative (e.g., in distributed optimization). The dynamics are influenced by the interaction matrix $\mathbf{J}(t)$, a block-structured Jacobian encoding pairwise competition or cooperation:

$$\mathbf{J}(t) = \begin{pmatrix} \nabla_{\theta_1} \nabla_{\theta_1} \mathcal{L}_1 & \cdots & \nabla_{\theta_K} \nabla_{\theta_1} \mathcal{L}_1 \\ \vdots & \ddots & \vdots \\ \nabla_{\theta_1} \nabla_{\theta_K} \mathcal{L}_K & \cdots & \nabla_{\theta_K} \nabla_{\theta_K} \mathcal{L}_K \end{pmatrix}.$$
The eigenvalues of $J(t)$ determine stability: if all eigenvalues have negative real parts, the system converges to a Nash equilibrium; if any are positive, instability or limit cycles may emerge, as observed in GAN training.

In the special case of zero-sum games (e.g., GANs), the dynamics can be reformulated as a Hamiltonian system, where the total "energy" $H(\Theta, \Lambda)$ is conserved along trajectories. Let $\Theta = (\theta_G, \theta_D)$ denote the generator and discriminator parameters, respectively. The Hamiltonian dynamics are given by:

$$ \frac{d\theta_G}{dt} = \nabla_{\Lambda} H, \quad \frac{d\theta_D}{dt} = -\nabla_{\Theta} H, \quad \text{where} \quad H(\Theta, \Lambda) = \mathcal{L}_G(\Theta) - \mathcal{L}_D(\Theta). $$

This formulation reveals oscillatory behavior when $\nabla_{\theta_G} \nabla_{\theta_D} \mathcal{L} \neq 0$, explaining cyclic divergence in GANs. Alternatively, training dynamics can be analyzed via Wasserstein gradient flows, where parameter updates approximate the gradient flow of a functional in probability space. For a loss $\mathcal{L}(\rho)$ over a parameter distribution $\rho(t)$, the continuity equation describes its evolution:

$$ \frac{\partial \rho(t)}{\partial t} = \nabla \cdot \left( \rho(t) \nabla \frac{\delta \mathcal{L}}{\delta \rho} \right), $$

where $\frac{\delta \mathcal{L}}{\delta \rho}$ is the functional derivative. This perspective connects deep learning to mean-field theory, explaining phenomena like mode collapse and implicit regularization via the geometry of parameter distributions. The interplay between differential game theory and training dynamics thus provides a unified framework for understanding optimization, stability, and convergence in complex learning systems.
# 6 Optimal transport theory in Deep Neural Networks

Optimal transport (OT) theory provides a rigorous geometric framework for comparing probability distributions and has become a foundational tool in analyzing and training deep neural networks (DNNs). Given two probability measures $\mu$ and $\nu$ defined on metric spaces $\mathcal{X}$ and $\mathcal{Y}$, respectively, the Monge-Kantorovich problem seeks a transport plan $\pi \in \Pi(\mu, \nu)$ that minimizes the expected cost of moving mass from $\mu$ to $\nu$:

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

where $c: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}_+$ is a cost function (e.g., Euclidean distance $c(x, y) = \|x - y\|^2$), and $\Pi(\mu, \nu)$ denotes the set of joint distributions with marginals $\mu$ and $\nu$. The optimal transport distance, or Wasserstein distance, quantifies the minimal effort required to reconfigure $\mu$ into $\nu$:

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

In deep learning, OT is employed to align latent representations, regularize generative models, and define loss functions for distribution matching. For instance, Wasserstein GANs (WGANs) replace the Jensen-Shannon divergence with the 1-Wasserstein distance $W_1$, leading to improved training stability:

$$ \min_G \max_{\|D\|_L \le 1} \mathbb{E}_{x \sim \mu}[D(x)] - \mathbb{E}_{z \sim p_z}[D(G(z))], \qquad (6.3) $$

where $D$ is a 1-Lipschitz discriminator enforcing the Wasserstein constraint.

The computational complexity of OT motivates approximations such as entropy-regularized transport, where the Kantorovich problem is relaxed with a Kullback-Leibler (KL) divergence penalty:

$$ W_{c,\epsilon}(\mu, \nu) = \inf_{\pi \in \Pi(\mu, \nu)} \int c(x, y) d\pi(x, y) + \epsilon \text{KL}(\pi \| \mu \otimes \nu). \qquad (6.4) $$

The Sinkhorn algorithm solves this efficiently via iterative scaling, enabling scalable OT-based losses in DNNs. The Sinkhorn divergence, a debiased version of $W_{c,\epsilon}$, is given by:

$$ 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), \qquad (6.5) $$

which is positive definite and avoids the bias of entropy smoothing. This divergence is used in tasks like domain adaptation, where source and target distributions $\mu_S$ and $\mu_T$ are aligned by minimizing $S_{\epsilon}(\phi \# \mu_S, \mu_T)$ for a feature map $\phi$.

OT theory also informs the design of neural architectures through dynamic transport perspectives. The continuity equation describes the evolution of a parameter distribution $\rho_t$ under gradient flow:

$$ \frac{\partial \rho_t}{\partial t} + \nabla \cdot (\rho_t v_t) = 0, \quad v_t = -\nabla \frac{\delta \mathcal{L}}{\delta \rho_t}, \qquad (6.6) $$
where $v_t$ is the velocity field induced by the loss functional $\mathcal{L}$. This links to the Benamou-Brenier formula, which expresses $W_2$ as an infimum over kinetic energy:

$$W_2^2(\mu, \nu) = \inf_{(\rho, v)} \int_0^1 \int \|v_t(x)\|^2 d\rho_t(x) dt, \qquad (6.7)$$

subject to $\partial_t \rho_t + \nabla \cdot (\rho_t v_t) = 0$. In deep learning, this inspires Neural ODEs and normalizing flows, where layers are viewed as discretized OT maps. For example, a residual block $x_{t+1} = x_t + f_\theta(x_t)$ approximates an infinitesimal transport along $f_\theta$, with the network trained to minimize the OT cost between input and output distributions.

OT theory further provides mechanisms for geometric regularization via barycenters. Given distributions $\mu_1, \dots, \mu_N$, their Wasserstein barycenter $\bar{\mu}$ minimizes:

$$\bar{\mu} = \arg \min_{\nu} \sum_{i=1}^{N} \lambda_i W_p^p(\nu, \mu_i), \qquad (6.8)$$

where $\lambda_i$ are convex weights. In DNNs, this facilitates multi-task learning by interpolating between task-specific distributions in latent space. The barycentric projection of a transport plan $\pi$ also defines a mapping $T(x) = \mathbb{E}_{y \sim \pi(\cdot|x)}[y]$, used in style transfer and cross-domain alignment.

By embedding OT principles into deep learning, models gain robustness to distributional shifts, improved generative fidelity, and theoretically grounded training dynamics, bridging measure theory and neural optimization.

## 6.1 Literature Review of Optimal transport theory in Deep Neural Networks

Optimal transport theory has become a cornerstone in the analysis and design of deep neural networks, providing a rigorous framework for understanding geometric and probabilistic aspects of learning. Peyré and Cuturi (2019) [1515] offer a comprehensive introduction to the mathematical foundations of optimal transport, with a focus on computational aspects that are crucial for scaling these methods to high-dimensional spaces encountered in deep learning. Its contribution lies in bridging theoretical optimal transport with practical algorithms, enabling applications such as Wasserstein barycenters and gradient flows in neural network training. Santambrogio (2015) [1518] provides a rigorous treatment of the analytical aspects of optimal transport, including Kantorovich duality and Monge-Ampère equations, which are essential for understanding the geometric regularization effects of Wasserstein distances in neural networks. This text is particularly valuable for its deep mathematical insights into how transport-based losses can stabilize and improve generalization in deep models. Arjovsky et al. (2017) [1459] revolutionized generative modeling by introducing the Wasserstein distance as a loss function, addressing the vanishing gradients and mode collapse issues in traditional GANs. This work rigorously establishes the connection between optimal transport and adversarial training, showing how the Wasserstein metric provides a meaningful gradient signal even when the supports of the distributions are disjoint. Gulrajani et al. (2017) [1473] further refine this approach by proposing gradient penalty methods to enforce the Lipschitz constraint, a key technical contribution that ensures stable training. These works collectively demonstrate how optimal transport theory can lead to more robust and interpretable generative models. Frogner et al. (2015) [1508] explore the statistical properties of Wasserstein-based learning, providing generalization bounds and approximation guarantees for neural networks trained with optimal transport objectives. This work is foundational for understanding the trade-offs between expressivity and sample complexity in transport-based deep learning.
Villani (2009) [1521] offers a deep dive into the geometric and probabilistic aspects of optimal transport, with implications for understanding the latent space geometry of neural networks. Its contribution lies in connecting abstract measure-theoretic concepts to practical questions about how neural networks transform and interpolate between distributions. Cuturi (2013) [1505] introduces entropy-regularized optimal transport, a breakthrough that enables efficient computation of Wasserstein distances in high dimensions. This work is particularly impactful for deep learning, as it allows scalable training of models with optimal transport losses. Genevay et al. (2018) [1509] extend these ideas to generative modeling, showing how Sinkhorn divergences can be used to train variational autoencoders with improved sample quality. These contributions highlight the interplay between computational optimal transport and deep generative modeling. Taghvaei and Jalali (2019) [1519] propose a neural network-based approach to solving optimal transport problems, effectively learning the transport map as a parameterized function. This represents a significant shift from traditional linear programming methods, enabling scalable solutions to high-dimensional transport problems. Tolstikhin et al. (2018) [1520] introduce a framework for auto-encoder training based on optimal transport, providing a principled way to balance reconstruction accuracy and latent space regularization. This work is notable for its rigorous treatment of the connection between variational inference and Wasserstein distances. Ramdas et al. (2017) [1517] develop statistical tests based on optimal transport, with applications to evaluating generative models and detecting distributional shifts in deep learning. These contributions demonstrate the versatility of optimal transport in addressing fundamental challenges in neural network training and evaluation.

Ambrosio et al. (2008) [1503] provide a rigorous foundation for understanding gradient-based optimization in the space of probability measures, which is crucial for analyzing the dynamics of neural network training through the lens of optimal transport. Its contribution lies in formalizing the connection between Wasserstein gradient flows and optimization algorithms used in deep learning. Mei et al. (2018) [1513] apply optimal transport theory to analyze the convergence of stochastic gradient descent in overparameterized neural networks, showing how the particle dynamics can be described as a Wasserstein gradient flow. This work bridges the gap between theoretical optimal transport and practical deep learning optimization. Janati et al. (2020) [1511] extend optimal transport theory to handle mismatched masses, a common scenario in deep learning where distributions may not be normalized. This technical advancement enables more flexible applications of transport-based methods in neural networks. Evans (2021) [1507] surveys the broad landscape of optimal transport theory, with particular emphasis on its connections to partial differential equations and variational problems. This work is invaluable for understanding how transport-based regularization can influence the solutions learned by neural networks. Courty et al. (2017) [1504] introduce a framework for adapting neural networks across domains using Wasserstein distances, providing rigorous guarantees on target domain performance. This contribution is particularly impactful for transfer learning and unsupervised domain adaptation. Li et al. (2015) [1512] explore the relationship between optimal transport and kernel-based discrepancies, offering insights into alternative distributional metrics for generative modeling. These works collectively expand the toolbox of transport-based methods available for deep learning applications.

Rachev and Rüschendorf (1998) [1516] provide encyclopedic coverage of optimal transport theory, with extensive treatment of its probabilistic and statistical aspects. This reference is essential for understanding the measure-theoretic foundations of Wasserstein distances in neural network analysis. Deshpande et al. (2018) [1506] introduce computationally efficient approximations to optimal transport by projecting distributions onto random directions, enabling scalable training of generative models. This work represents an important practical advancement in making transport-based methods feasible for large-scale deep learning. Ozair et al. (2019) [1514] develop a transport-based dependency metric that can be used to regularize neural network representations, contributing to
the field of disentangled representation learning. These contributions demonstrate the growing influence of optimal transport across various subfields of deep learning. Richards et. al. (2020) [222] provided an intuitive introduction to modeling Rayleigh-Bénard convection within a stochastic framework, emphasizing understanding over mathematical rigor. It outlines how stochastic processes such as Brownian motion and stochastic differential equations can be applied to model convective phenomena. The work serves as a conceptual and computational guide for developing numerical simulations of stochastic fluid dynamics systems. Huang et al. (2021) [1510] propose a method for constructing explicit transport maps in generative models, providing both theoretical guarantees and practical algorithms. This work bridges the gap between normalizing flows and optimal transport theory. The convergence analysis of transport-based optimization in deep learning is further advanced by rigorous studies connecting Wasserstein geometry with stochastic gradient descent dynamics. The interplay between optimal transport and deep learning continues to inspire new theoretical insights and algorithmic innovations, shaping the future of machine learning research and applications.

## 6.2 Analysis of Optimal transport theory in Deep Neural Networks

Optimal transport theory provides a mathematically rigorous framework for measuring distances between probability distributions by solving a constrained optimization problem that minimizes the cost of transporting mass from one distribution to another. Consider two probability measures $\mu$ and $\nu$ defined on a common measurable space $(\Omega, \mathcal{F})$, where $\mu$ has support $\mathcal{X} \subseteq \mathbb{R}^d$ and $\nu$ has support $\mathcal{Y} \subseteq \mathbb{R}^d$. Given a cost function $c: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}_{\ge 0}$, the optimal transport problem seeks to determine a transport plan that minimizes the total cost of transforming $\mu$ into $\nu$. In the Monge formulation, this problem is expressed as

$$ \inf_{T: \mathcal{X} \to \mathcal{Y}} \int_{\mathcal{X}} c(x, T(x)) d\mu(x), \qquad (6.9) $$

subject to the constraint that the transport map $T$ satisfies the push-forward condition

$$ T_{\#}\mu = \nu, \qquad (6.10) $$

which ensures that the measure $\nu$ is obtained by pushing forward $\mu$ through $T$, i.e., for any measurable set $A \subseteq \mathcal{Y}$,

$$ \nu(A) = \mu(T^{-1}(A)). \qquad (6.11) $$

The Monge problem is often ill-posed since a transport map may not exist. To alleviate this issue, the Kantorovich relaxation introduces a coupling $\gamma$ between $\mu$ and $\nu$, where $\gamma$ is a probability measure on $\mathcal{X} \times \mathcal{Y}$ satisfying the marginal constraints

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

The Kantorovich formulation of the optimal transport problem is thus

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

where $\Pi(\mu, \nu)$ is the set of all couplings satisfying the marginal constraints. A commonly used cost function is the $p$-th power of the Euclidean distance, i.e., $c(x, y) = \|x - y\|^p$, which leads to the Wasserstein distance of order $p$, given by

$$ W_p(\mu, \nu) = \left( \inf_{\gamma \in \Pi(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{Y}} \|x - y\|^p d\gamma(x, y) \right)^{\frac{1}{p}}. \qquad (6.14) $$
The Wasserstein distance defines a proper metric on the space of probability measures, provided that $\mu$ and $\nu$ have finite p-th moments, i.e.,

$$ \int_X \|x\|^p d\mu(x) < \infty, \quad \int_Y \|y\|^p d\nu(y) < \infty. \qquad (6.15) $$

In deep neural networks, optimal transport plays a crucial role in generative modeling, where it provides a more stable alternative to divergence-based training objectives. In Wasserstein generative adversarial networks (WGANs), the standard Jensen-Shannon divergence is replaced by the Wasserstein-1 distance, leading to the objective

$$ \sup_{\|f\|_L \le 1} \mathbb{E}_{x \sim P_r}[f(x)] - \mathbb{E}_{z \sim P_z}[f(G(z))], \qquad (6.16) $$

where $f$ is a Lipschitz-1 function, $P_r$ is the real data distribution, $P_z$ is the latent distribution, and $G: Z \to X$ is the generator function. The use of the Wasserstein distance improves gradient behavior and mitigates mode collapse, which is a common issue in standard GAN training. In domain adaptation, optimal transport is used to align source and target distributions by finding a transport plan $\gamma$ that minimizes the Wasserstein distance between them. Given a source distribution $P_s$ and a target distribution $P_t$, the goal is to find a transport map $T$ such that

$$ T_{\#}P_s \approx P_t, \qquad (6.17) $$

which is often achieved by solving

$$ \inf_{\gamma \in \Pi(P_s, P_t)} \int_{X \times Y} \|x - y\|^p d\gamma(x, y). \qquad (6.18) $$

To ensure computational efficiency, entropy-regularized optimal transport introduces a regularization term based on the Kullback-Leibler divergence, leading to the Sinkhorn distance

$$ W_{\lambda}(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y)d\gamma(x, y) + \lambda KL(\gamma \| \mu \otimes \nu). \qquad (6.19) $$

The regularization parameter $\lambda > 0$ controls the trade-off between accuracy and computational efficiency. The Sinkhorn algorithm, which iteratively updates the dual potentials using Bregman projections, provides an efficient means of computing approximate solutions to the OT problem. Optimal transport has also been applied to learning energy-based models, where it is used to define a structured loss function for training probability models. In probabilistic autoencoders, OT distances provide a principled approach for matching latent and data distributions. In Bayesian deep learning, OT-based priors ensure that posterior distributions maintain smooth and stable representations, improving generalization in uncertainty estimation tasks. The OT-based barycenter problem arises in multi-modal learning, where the goal is to compute a central distribution that minimizes the sum of Wasserstein distances to a set of given measures $\mu_1, \mu_2, \dots, \mu_n$. The barycenter $\nu$ is defined as

$$ \arg\min_{\nu} \sum_{i=1}^{n} W_{p}(\nu, \mu_{i})^{p}, \qquad (6.20) $$

which generalizes the notion of averaging to probability distributions. Applications of OT barycenters include meta-learning, ensemble methods, and federated learning, where a global model must aggregate information from multiple local distributions while minimizing divergence. The theoretical properties of optimal transport, including dual formulations and regularity results, provide strong guarantees for its application in deep neural networks. The dual formulation of the Wasserstein distance is given by the Kantorovich-Rubinstein theorem, which states that for $W_1$,

$$ W_1(\mu, \nu) = \sup_{\|f\|_L \le 1} \int_X f(x)d\mu(x) - \int_Y f(y)d\nu(y). \qquad (6.21) $$
This dual representation allows gradient-based optimization methods to efficiently approximate OT distances in neural network training. Furthermore, recent advances in computational optimal transport, including sliced Wasserstein distances and entropic maps, have enabled scalable implementations suitable for high-dimensional deep learning applications.

Optimal transport thus provides a fundamental mathematical framework for distributional matching in deep learning. By leveraging Wasserstein distances and entropy-regularized formulations, OT enables more stable training, improved generalization, and better interpretability in neural networks, making it an essential tool for modern deep learning methodologies.

## 6.3 Jensen-Shannon divergence

The Jensen-Shannon divergence (JSD) is a rigorously defined, symmetric, and bounded measure of dissimilarity between two probability distributions $P$ and $Q$. It is constructed from the Kullback-Leibler divergence (KLD) but incorporates a mixture distribution to ensure symmetry and finiteness. The JSD is defined as:

$$ \mathrm{JSD}(P, Q) = \frac{1}{2} D_{KL}(P \| M) + \frac{1}{2} D_{KL}(Q \| M), \qquad (6.22) $$

where $M$ is the average distribution of $P$ and $Q$, given by:

$$ M = \frac{1}{2}(P + Q). \qquad (6.23) $$

Here, $D_{KL}(P \| M)$ and $D_{KL}(Q \| M)$ are the Kullback-Leibler divergences between $P$ and $M$, and $Q$ and $M$, respectively. The KLD for two discrete probability distributions $P$ and $M$ is defined as:

$$ D_{KL}(P \| M) = \sum_{i} P(i) \log \left( \frac{P(i)}{M(i)} \right), \qquad (6.24) $$

where $P(i)$ and $M(i)$ are the probabilities assigned to the $i$-th event by $P$ and $M$, respectively. For continuous probability distributions, the KLD is expressed as:

$$ D_{KL}(P \| M) = \int_{-\infty}^{\infty} p(x) \log \left( \frac{p(x)}{m(x)} \right) dx, \qquad (6.25) $$

where $p(x)$ and $m(x)$ are the probability density functions of $P$ and $M$, respectively. Substituting the definition of $M$ into the JSD, we obtain:

$$ \mathrm{JSD}(P, Q) = \frac{1}{2} \sum_{i} P(i) \log \left( \frac{P(i)}{\frac{1}{2}(P(i) + Q(i))} \right) + \frac{1}{2} \sum_{i} Q(i) \log \left( \frac{Q(i)}{\frac{1}{2}(P(i) + Q(i))} \right). \qquad (6.26) $$

For continuous distributions, this becomes:

$$ \mathrm{JSD}(P, Q) = \frac{1}{2} \int_{-\infty}^{\infty} p(x) \log \left( \frac{p(x)}{\frac{1}{2}(p(x) + q(x))} \right) dx + \frac{1}{2} \int_{-\infty}^{\infty} q(x) \log \left( \frac{q(x)}{\frac{1}{2}(p(x) + q(x))} \right) dx. \qquad (6.27) $$

The JSD can also be expressed in terms of Shannon entropy. Let $H(P)$ denote the Shannon entropy of $P$, defined for discrete distributions as:

$$ H(P) = - \sum_{i} P(i) \log P(i). \qquad (6.28) $$

For continuous distributions, the differential entropy is:

$$ H(P) = - \int_{-\infty}^{\infty} p(x) \log p(x) dx. \qquad (6.29) $$
Using entropy, the JSD can be rewritten as:

$$JSD(P, Q) = H\left(\frac{P+Q}{2}\right) - \frac{1}{2}H(P) - \frac{1}{2}H(Q). \qquad (6.30)$$

This formulation demonstrates that the JSD measures the difference between the entropy of the average distribution $M$ and the average of the entropies of $P$ and $Q$. The JSD is bounded between 0 and 1 when the logarithm base is 2, as:

$$0 \le JSD(P, Q) \le 1. \qquad (6.31)$$

The lower bound is achieved when $P = Q$, and the upper bound is approached when $P$ and $Q$ are maximally different. The square root of the JSD is a metric, satisfying the triangle inequality:

$$\sqrt{JSD(P, Q)} \le \sqrt{JSD(P, R)} + \sqrt{JSD(R, Q)}, \qquad (6.32)$$

for any probability distribution $R$. This property makes the JSD particularly useful in applications requiring a distance metric, such as clustering and classification. The JSD is also related to the total variation distance $\delta(P, Q)$ by the inequality:

$$\delta(P, Q) \le \sqrt{2 \cdot JSD(P, Q)}. \qquad (6.33)$$

This relationship connects the JSD to other statistical distances, providing a bridge between information-theoretic and measure-theoretic perspectives. The JSD is widely used in machine learning, bioinformatics, and statistics due to its symmetry, boundedness, and interpretability. It is particularly useful for comparing empirical distributions, as it avoids the infinite values that can arise with the KLD. The JSD can also be generalized to compare more than two distributions. For $n$ distributions $P_1, P_2, \dots, P_n$, the generalized JSD is:

$$JSD(P_1, P_2, \dots, P_n) = H\left(\frac{1}{n}\sum_{i=1}^{n} P_i\right) - \frac{1}{n}\sum_{i=1}^{n} H(P_i). \qquad (6.34)$$

This generalization retains the properties of the standard JSD, including symmetry and boundedness. The JSD is also closely related to the concept of mutual information, as it can be interpreted as the mutual information between a random variable and a mixture distribution. Specifically, for two distributions $P$ and $Q$, the JSD can be expressed as:

$$JSD(P, Q) = I(X; Y), \qquad (6.35)$$

where $X$ is a random variable with distribution $P$ or $Q$, and $Y$ is a binary indicator variable selecting between $P$ and $Q$. This interpretation provides a deeper understanding of the JSD as a measure of information overlap. In summary, the Jensen-Shannon divergence is a versatile and mathematically rigorous tool for quantifying the difference between probability distributions, with applications across numerous scientific and engineering disciplines. Its properties make it a preferred choice in many scenarios where symmetric and bounded divergence measures are required.

## 6.4 Matching Latent and Data Distributions in Probabilistic Autoencoders

Optimal Transport (OT) distances provide a mathematically rigorous framework for comparing probability distributions, making them particularly suitable for aligning latent and data distributions in Probabilistic Autoencoders (PAEs). The core idea of OT is to find the most efficient way to transform one distribution into another, where efficiency is quantified by a cost function. In the
context of PAEs, this translates to minimizing the discrepancy between the encoded latent distribution $q_{\phi}(z|x)$ and the prior distribution $p(z)$, as well as between the decoded data distribution $p_{\theta}(x|z)$ and the true data distribution $p_{\text{data}}(x)$.

Let $\mu$ and $\nu$ be two probability measures defined on metric spaces $\mathcal{X}$ and $\mathcal{Y}$, respectively. The OT problem seeks a coupling $\pi \in \Pi(\mu, \nu)$ that minimizes the total transportation cost:

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

where $c(x, y)$ is a cost function, typically the Euclidean distance $c(x, y) = \|x - y\|^p$ for $p \ge 1$. The p-Wasserstein distance is a special case of OT, defined as:

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

In PAEs, the latent distribution $q_{\phi}(z|x)$ is often modeled as a Gaussian $\mathcal{N}(\mu_{\phi}(x), \Sigma_{\phi}(x))$, and the prior $p(z)$ is typically a standard Gaussian $\mathcal{N}(0, I)$. The OT distance between these distributions can be computed using the 2-Wasserstein distance, which has a closed-form expression for Gaussian distributions:

$$W_2^2(\mathcal{N}(\mu_1, \Sigma_1), \mathcal{N}(\mu_2, \Sigma_2)) = \|\mu_1 - \mu_2\|^2 + \text{Tr}(\Sigma_1 + \Sigma_2 - 2(\Sigma_1^{1/2}\Sigma_2\Sigma_1^{1/2})^{1/2}). \qquad (6.38)$$

For the decoded data distribution $p_{\theta}(x|z)$, the OT distance to the true data distribution $p_{\text{data}}(x)$ can be approximated using discrete OT methods. Given samples $\{x_i\}_{i=1}^N$ from $p_{\text{data}}(x)$ and $\{\hat{x}_i\}_{i=1}^N$ from $p_{\theta}(x|z)$, the empirical OT problem becomes:

$$W_c(\hat{\mu}, \nu) = \min_{\pi \in \Pi(\hat{\mu}, \nu)} \sum_{i,j} c(\hat{x}_i, x_j) \pi_{ij}, \qquad (6.39)$$

where $\hat{\mu} = \frac{1}{N} \sum_{i=1}^N \delta_{\hat{x}_i}$ and $\nu = \frac{1}{N} \sum_{i=1}^N \delta_{x_i}$ are empirical measures. The coupling $\pi$ is a doubly stochastic matrix satisfying $\sum_i \pi_{ij} = \frac{1}{N}$ and $\sum_j \pi_{ij} = \frac{1}{N}$. The OT distance can be incorporated into the PAE objective function as a regularization term. The overall loss function $\mathcal{L}$ combines the reconstruction loss and the OT-based regularization:

$$\mathcal{L}(\theta, \phi) = \mathbb{E}_{x \sim p_{\text{data}}} [-\log p_{\theta}(x|z)] + \lambda W_2^2(q_{\phi}(z|x), p(z)) + \gamma W_c(p_{\theta}(x|z), p_{\text{data}}(x)), \qquad (6.40)$$

where $\lambda$ and $\gamma$ are hyperparameters controlling the strength of the regularization terms. The first term ensures accurate reconstruction, while the second and third terms enforce consistency between the latent and prior distributions, and between the decoded and true data distributions, respectively. The OT framework also allows for the use of entropy-regularized OT, which introduces a regularization term to the OT problem to make it computationally tractable. The entropy-regularized OT problem is given by:

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

where $H(\pi) = \int_{\mathcal{X} \times \mathcal{Y}} \pi(x, y) \log \pi(x, y) dx dy$ is the entropy of the coupling $\pi$, and $\epsilon > 0$ is the regularization parameter. This formulation leads to the Sinkhorn algorithm, which provides an efficient way to compute approximate OT distances.

In summary, OT distances offer a principled and mathematically rigorous approach for matching latent and data distributions in PAEs. By minimizing the Wasserstein distance between distributions, PAEs can achieve better alignment between the encoded latent space and the prior, as well as between the decoded data and the true data distribution. This results in more robust and interpretable models, with improved generative and reconstructive capabilities.
## 6.4.1 Probabilistic Autoencoders

Probabilistic autoencoders are generative models that combine the principles of autoencoders with probabilistic frameworks to learn a latent representation of data. Let $\mathbf{x} \in \mathbb{R}^D$ be the observed data, and $\mathbf{z} \in \mathbb{R}^d$ be the latent variable, where $d \ll D$. The goal is to model the joint distribution $p(\mathbf{x}, \mathbf{z})$, which can be factorized as:

$$p(\mathbf{x}, \mathbf{z}) = p(\mathbf{x}|\mathbf{z})p(\mathbf{z}). \qquad (6.42)$$

Here, $p(\mathbf{z})$ is the prior distribution over the latent space, often chosen as a standard Gaussian $N(0, I)$, and $p(\mathbf{x}|\mathbf{z})$ is the likelihood function, which models the data generation process. The encoder in a probabilistic autoencoder approximates the posterior distribution $p(\mathbf{z}|\mathbf{x})$, which is typically intractable. Instead, a variational distribution $q_\phi(\mathbf{z}|\mathbf{x})$ is introduced, parameterized by $\phi$, to approximate the true posterior. This distribution is often chosen as a Gaussian $N(\mathbf{z}; \mu_\phi(\mathbf{x}), \Sigma_\phi(\mathbf{x}))$, where $\mu_\phi(\mathbf{x})$ and $\Sigma_\phi(\mathbf{x})$ are outputs of a neural network. The decoder, parameterized by $\theta$, models the likelihood $p_\theta(\mathbf{x}|\mathbf{z})$, which is also often Gaussian for continuous data:

$$p_{\theta}(\mathbf{x}|\mathbf{z}) = N(\mathbf{x}; \mu_{\theta}(\mathbf{z}), \Sigma_{\theta}(\mathbf{z})). \qquad (6.43)$$

The learning objective is to maximize the marginal likelihood $p(\mathbf{x})$, which can be expressed as:

$$p(\mathbf{x}) = \int p_{\theta}(\mathbf{x}|\mathbf{z})p(\mathbf{z}) d\mathbf{z}. \qquad (6.44)$$

However, this integral is intractable due to the high-dimensional latent space. Instead, the evidence lower bound (ELBO) is maximized, which is derived using variational inference:

$$\log p(\mathbf{x}) \geq \mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})} [\log p_{\theta}(\mathbf{x}|\mathbf{z})] - D_{\text{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z})) = \text{ELBO}. \qquad (6.45)$$

The first term in the ELBO is the reconstruction loss, which encourages the model to accurately reconstruct the input data. The second term is the Kullback-Leibler (KL) divergence between the variational posterior $q_\phi(\mathbf{z}|\mathbf{x})$ and the prior $p(\mathbf{z})$, which regularizes the latent space to match the prior distribution. The ELBO can be rewritten as:

$$\text{ELBO} = \mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})} [\log p_{\theta}(\mathbf{x}|\mathbf{z})] - \frac{1}{2} \sum_{i=1}^{d} (\log \sigma_{\phi,i}^2(\mathbf{x}) - \sigma_{\phi,i}^2(\mathbf{x}) - \mu_{\phi,i}^2(\mathbf{x}) + 1), \qquad (6.46)$$

where $\mu_{\phi,i}(\mathbf{x})$ and $\sigma_{\phi,i}^2(\mathbf{x})$ are the mean and variance of the $i$-th dimension of the latent variable $\mathbf{z}$. The parameters $\phi$ and $\theta$ are optimized using stochastic gradient descent. The expectation term in the ELBO is approximated using Monte Carlo sampling. For a single sample $\mathbf{z}^{(l)} \sim q_\phi(\mathbf{z}|\mathbf{x})$, the ELBO can be approximated as:

$$\text{ELBO} \approx \log p_{\theta}(\mathbf{x}|\mathbf{z}^{(l)}) - D_{\text{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z})). \qquad (6.47)$$

To enable backpropagation through the sampling process, the reparameterization trick is used. Specifically, $\mathbf{z}$ is reparameterized as:

$$\mathbf{z} = \mu_{\phi}(\mathbf{x}) + \epsilon \odot \sigma_{\phi}(\mathbf{x}), \qquad (6.48)$$

where $\epsilon \sim N(0, I)$ and $\odot$ denotes element-wise multiplication. This allows gradients to flow through the sampling process, enabling efficient optimization. The reconstruction loss $\log p_\theta(\mathbf{x}|\mathbf{z})$ depends on the type of data. For continuous data, it is often the negative log-likelihood of a Gaussian distribution:

$$\log p_{\theta}(\mathbf{x}|\mathbf{z}) = -\frac{1}{2}\|\mathbf{x} - \mu_{\theta}(\mathbf{z})\|_{2}^{2} + \text{constant}. \qquad (6.49)$$
For binary data, the Bernoulli distribution is used, and the reconstruction loss becomes:

$$ \log p_{\theta}(\mathbf{x}|\mathbf{z}) = \sum_{i=1}^{D} x_i \log \hat{x}_i + (1-x_i) \log(1-\hat{x}_i), \qquad (6.50) $$

where $\hat{x}_i$ is the i-th element of the reconstructed data $\mu_{\theta}(\mathbf{z})$. The KL divergence term $D_{\text{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))$ can be computed analytically for Gaussian distributions. For a standard Gaussian prior $p(\mathbf{z}) = N(\mathbf{0}, \mathbf{I})$ and a Gaussian variational posterior $q_{\phi}(\mathbf{z}|\mathbf{x}) = N(\mathbf{z}; \mu_{\phi}(\mathbf{x}), \Sigma_{\phi}(\mathbf{x}))$, the KL divergence is:

$$ D_{\text{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z})) = \frac{1}{2} \left( \text{tr}(\Sigma_{\phi}(\mathbf{x})) + \|\mu_{\phi}(\mathbf{x})\|_{2}^{2} - d - \log \det(\Sigma_{\phi}(\mathbf{x})) \right). \qquad (6.51) $$

If the variational posterior is diagonal, i.e., $\Sigma_{\phi}(\mathbf{x}) = \text{diag}(\sigma_{\phi,1}^2(\mathbf{x}), \dots, \sigma_{\phi,d}^2(\mathbf{x}))$, the KL divergence simplifies to:

$$ D_{\text{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z})) = \frac{1}{2} \sum_{i=1}^{d} \left( \sigma_{\phi,i}^{2}(\mathbf{x}) + \mu_{\phi,i}^{2}(\mathbf{x}) - 1 - \log \sigma_{\phi,i}^{2}(\mathbf{x}) \right). \qquad (6.52) $$

The optimization of the ELBO is performed using gradient-based methods. The gradients of the ELBO with respect to $\phi$ and $\theta$ are computed using automatic differentiation. For a mini-batch of data $\{\mathbf{x}^{(1)}, \dots, \mathbf{x}^{(B)}\}$, the stochastic gradient estimate of the ELBO is:

$$ \nabla_{\phi,\theta} \text{ELBO} \approx \frac{1}{B} \sum_{b=1}^{B} \nabla_{\phi,\theta} \left( \log p_{\theta}(\mathbf{x}^{(b)}|\mathbf{z}^{(b)}) - D_{\text{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x}^{(b)})\|p(\mathbf{z})) \right), \qquad (6.53) $$

where $\mathbf{z}^{(b)} \sim q_{\phi}(\mathbf{z}|\mathbf{x}^{(b)})$. The probabilistic autoencoder can be extended to more complex architectures, such as hierarchical models with multiple layers of latent variables. In such cases, the latent variables are structured hierarchically, and the joint distribution is factorized as:

$$ p(\mathbf{x}, \mathbf{z}_1, \dots, \mathbf{z}_L) = p(\mathbf{x}|\mathbf{z}_1) \prod_{l=1}^{L-1} p(\mathbf{z}_l|\mathbf{z}_{l+1})p(\mathbf{z}_L), \qquad (6.54) $$

where $\mathbf{z}_1, \dots, \mathbf{z}_L$ are the latent variables at different levels of the hierarchy. The variational posterior is also factorized hierarchically:

$$ q_{\phi}(\mathbf{z}_1, \dots, \mathbf{z}_L|\mathbf{x}) = q_{\phi}(\mathbf{z}_1|\mathbf{x}) \prod_{l=2}^{L} q_{\phi}(\mathbf{z}_l|\mathbf{z}_{l-1}). \qquad (6.55) $$

The ELBO for hierarchical models includes additional KL divergence terms for each level of the hierarchy:

$$ \text{ELBO} = \mathbb{E}_{q_{\phi}(\mathbf{z}_1, \dots, \mathbf{z}_L|\mathbf{x})} [\log p_{\theta}(\mathbf{x}|\mathbf{z}_1)] - \sum_{l=1}^{L} D_{\text{KL}}(q_{\phi}(\mathbf{z}_l|\mathbf{z}_{l-1})\|p(\mathbf{z}_l|\mathbf{z}_{l+1})). \qquad (6.56) $$

In summary, probabilistic autoencoders provide a principled framework for learning latent representations of data by combining the strengths of autoencoders and probabilistic models. The ELBO serves as the optimization objective, balancing reconstruction accuracy and regularization of the latent space. The reparameterization trick enables efficient gradient-based optimization, and the framework can be extended to hierarchical models for more complex data structures.

## 6.4.2 2-Wasserstein distance

The 2-Wasserstein distance, also known as the Earth Mover's Distance (EMD) in the case of $p = 2$, is a metric used to quantify the difference between two probability distributions $\mu$ and $\nu$ defined on a metric space $(X, d)$. It is a specific instance of the more general $p$-Wasserstein distance, which