LLM Learning Dynamics follows Evolutionary Game Theory

Let $V := |\mathcal{Y}|$ be the number of labels. A classification model $\theta$ processes input $x \in \R^d$ and produce output logits $z := f(x, \theta) \in \R^{V}$. We use $y \in \R^{V}$ to denote the label vector and $y_i \in \mathcal{Y}$ a particular label. Let $\pi_\theta(y_i|x)$ be the probability mass assigned to label $y_i$ by model state $\theta$, obtained through $\pi_\theta(y|x) = \softmax(z)$. For a training data point $(x_u,y_u)$, we focus on cross-entropy loss given by $ \loss(\theta,x_u,y_u) = - \log \pi_\theta(y_u|x_u) $.

Following the empirical neural tangent kernel (eNTK) framework, we roll out the gradient of cross-entropy loss $\nabla_{\theta} \log \pi_{\theta_{t}}(y_u|x)$ into an output gradient term and a network gradient term: \begin{equation*} \begin{aligned} & \nabla_{\theta_{t}} \log \pi_{\theta_{t}} (y_i | x) = \underbrace{\nabla_{z}\log \pi_{\theta_{t}}(y_i | \mathbf{x})^\top }_{\text{Output Gradient}} \underbrace{\nabla_{\theta_{t}} f(x,\theta_t)}_{\text{Network Gradient}} \end{aligned} \end{equation*}

Let $x_o$ be such test-time input. We want to trace the change in $\pi_\theta(y_i|x_o)$ for all labels $y_i\in \mathcal{Y}$. Formally, one-step learning dynamics in this work is defined as: \begin{equation*} \begin{aligned} & \Delta \log \pi_{\theta_{t}}(y | x_o) \delequal \log \pi_{\theta_{t+1}} (y | x_o) - \log \pi_{\theta_{t}}(y | x_o) \end{aligned} \end{equation*}

Let $J(x_u,\theta_t) := \nabla_{\theta} f(x_u,\theta_t) \in \R^{V\times d}$ be the Jacobian matrix $J_{ij} = \frac{\partial f_i(x_u,\theta_t)}{\partial \theta_t^{(j)}}$. Define matrix-valued function $K(\cdot) : \R^{d} \times \R^d \rightarrow \R^{V \times V}$: \[K_t(x_o,x_u) \delequal J_t(x_o) J_t(x_u)^\top\] Let $e_i \in \R^d$ be the standard basis vector for label $y_i$, which has $0$ for all entries except the $y_i$-th entry. Define probability vector $\pi_t(x) := \pi_{\theta_{t}} (y|x) \in \R^{V}$ for simplicity. Apply a simple first-order Taylor expansion around $\pi_{\theta_{t+1}}$ to obtain the approximated single-step learning dynamics: \begin{equation}\label{eq: single-step-update} \begin{aligned} \Delta\log \pi_{\theta_{t}}(y_i|x_o) & \approx \eta \big[ e_i - \pi_t(x_o) \big]^\top K_t(x_o,x_u) \big[ e_u - \pi_t(x_u)\big] \end{aligned}\tag{1} \end{equation} We'll come back to this equation later.

Now let's switch gears to evolutionary game theory. For a more detailed introduction, see this excellent post.

Imagine a large population of individuals, where each individual can adopt one of $m$ different strategies. Let $x_k(t)$ be the proportion of the population using strategy $k$ at time $t$. The vector $x(t) = (x_1(t), \dots, x_m(t))$ represents the state of the population.

The success of a strategy is determined by its payoff or fitness. In a simple symmetric game, we can define a payoff matrix $A$, where $A_{ij}$ is the payoff for an individual using strategy $i$ when interacting with an individual using strategy $j$. The expected fitness of an individual using strategy $k$ in a population with state $x$ is $f_k(x) = (Ax)_k$. The average fitness of the entire population is $\bar{f}(x) = x^\top A x$.

Replicator dynamics provides a simple, intuitive model for how the population state $x(t)$ evolves. The core idea is that strategies with above-average fitness will become more common, while those with below-average fitness will decline. This is captured by the replicator differential equation:

The rate of change of strategy $k$'s proportion, $\dot{x}_k$, is proportional to its current proportion $x_k$ and its "fitness advantage"—the difference between its own fitness and the population average. If a strategy's fitness is higher than the average, its share of the population grows. If it's lower, its share shrinks. This simple rule can produce rich and complex dynamics.

Let's start with the simplest case: training repeatedly on a single data point $(x_u, y_u)$. Using our learning dynamics equation (1), we can derive a discrete update rule for the model's predicted probabilities. After each step, we re-normalize the probabilities to ensure they sum to one. This gives the following system, which describes how predictions at the test point $x_o$ and training point $x_u$ co-evolve. Note that our framework handles both gradient descent ($\eta > 0$) and gradient ascent ($\eta < 0$).

The following box says that the Eq. (2) is a well-behaved discretization of some ODE system. As the learning rate $\eta$ gets smaller, the trajectory of the discrete updates gets closer and closer to the trajectory of the ODE.

While training on a single data point is instructive, a more realistic setup involves a full dataset $D = \{(x_u, y_u)\}_{u=1}^N$. With full-batch gradient descent, the model's parameters are updated based on the average gradient across all training examples: \[ \theta_{t+1} \leftarrow \theta_{t} + \eta \sum_{u=1}^N \nabla_\theta \log \pi_\theta(y_u | x_u) \Big|_{\theta = \theta_t} \] The change in log-probabilities at our test point $x_o$ is now an average of the influences from all training points: \[ \Delta\log \pi_{\theta_t}(y_i|x_o)\approx \frac{\eta}{N}\sum_{u=1}^N\big[e_i-\pi_t(x_o)\big]^\top K(o,u)\big[e_u-\pi_t(x_u)\big]. \]

This leads to a more complex, but analogous, discrete update rule for the probabilities:

The following box is the main result. The state of our system, $Z$, now includes the model's probability distribution for the test point ($a$) and for every single training point ($b^{(v)}$). The theorem shows that this high-dimensional discrete process also converges to a coupled system of replicator ODEs. Each data point's prediction evolves according to its own replicator equation, but the fitness terms ($\phi$) are coupled, meaning the evolution of one depends on the current state of all the others, as one may expect.

What does "fitness" mean in the context of LLM learning? The fitness term for a given label $i$ at a data point $\alpha$ is a sum of interactions with all training data points $\beta \in \{1, ..., N\}$. Each interaction, governed by the kernel $K(\alpha, \beta)$, creates a "push-pull" dynamic.

Let's break down the fitness term for label $i$ at data point $\alpha$:

Here, we've expanded the kernel $K_{ij}(\alpha,\beta)$ as the inner product of Jacobians $\langle\nabla f_i(\alpha),\nabla f_j(\beta)\rangle$.

• The "Pull" Term measures the alignment (or synergy) between the gradient for label $i$ at point $\alpha$ and the gradient for the true label $y_\beta$ at point $\beta$. High alignment means that increasing the probability of label $i$ for $\alpha$ is "helpful" for correctly classifying $\beta$. This term pulls the dynamics towards the ground truth data.
• The "Push" Term measures the average alignment between the gradient for label $i$ at $\alpha$ and the gradients for all labels at $\beta$, weighted by the model's current prediction $\pi(\cdot|\beta)$. This term represents the influence of the model's own beliefs. If the model is currently confident in a wrong answer for $\beta$, this term might push the dynamics in the wrong direction.

Learning is thus a dynamic competition. The probability of a label grows if, on average, it aligns better with the true labels across the dataset than it does with the model's current (and potentially flawed) predictions.