In the [REINFORCE
article](llm-token-generation.html#reinforcement-learning), we jumped into the gradient. But I wanted to take a step
back and start by factorizing the probability of a trajectory from
first principles.
Let us consider a short trajectory:
$$ \tau=(s_0,a_0,s_1,a_1,s_2). $$ We start in state $s_0$, take action $a_0$, move to state $s_1$, take action $a_1$, and arrive at $s_2$.
What is the probability of this entire trajectory under a policy with parameters $\theta$?
For people coming from supervised learning, think of the policy as a model with parameters $\theta$. A policy gives a probability for each action it can take. In token generation, the actions are the possible next tokens, so this is the same as the model giving a probability for each possible next-token class.
Start with the Chain Rule #
We can begin with the chain rule. For three random variables, the chain rule says
$$ p(x,y,z)=p(x)p(y\mid x)p(z\mid x,y). $$ Applying the same rule to our trajectory gives
$$ \begin{aligned} p(\tau\mid\theta) ={}&p(s_0,a_0,s_1,a_1,s_2\mid\theta)\ ={}&p(s_0\mid\theta)\ &\times p(a_0\mid s_0,\theta)\ &\times p(s_1\mid s_0,a_0,\theta)\ &\times p(a_1\mid s_0,a_0,s_1,\theta)\ &\times p(s_2\mid s_0,a_0,s_1,a_1,\theta). \end{aligned} $$ So far, we have only used the chain rule. We have not made any assumptions about the policy or the environment.
The Initial State #
The policy, i.e. the model, does not choose the initial state. The initial state is sampled from an initial state distribution. We assume this distribution does not depend on the policy parameters, so
$$ p(s_0\mid\theta)=p(s_0). $$
For a language model, $s_0$ is the initial prompt. The prompt is given to the model before generation begins.
The Environment #
The next state is produced by the environment after the policy chooses an action. We assume that the environment is not parameterized by $\theta$. Therefore,
$$ p(s_1\mid s_0,a_0,\theta)=p(s_1\mid s_0,a_0), $$ and
$$ p(s_2\mid s_0,a_0,s_1,a_1,\theta) =p(s_2\mid s_0,a_0,s_1,a_1). $$ This does not mean that the policy has no influence on the states we visit. The policy chooses the actions, and those actions affect the next states. It only means that once the state and action are given, the environment uses the same transition rule for every value of $\theta$.
We now use the Markov assumption. Given the current state and action, the environment does not need the earlier states and actions to produce the next state. In our example,
$$ p(s_2\mid s_0,a_0,s_1,a_1)=p(s_2\mid s_1,a_1). $$ At an arbitrary time $t$, this is
$$ p(s_{t+1}\mid s_0,a_0,\ldots,s_t,a_t) =p(s_{t+1}\mid s_t,a_t). $$
For this equality to hold, the state $s_t$ must contain the information from the past that is needed to determine the distribution of the next state.
The Policy #
The actions are generated by the policy. Instead of writing the action probability as $p(a_t\mid s_t,\theta)$, we define the policy notation
$$ p(a_t\mid s_t,\theta) \overset{\mathrm{def}}{=} \pi_\theta(a_t\mid s_t). $$ This is just notation, but it helps us distinguish the probabilities produced by the policy from the probabilities produced by the environment.
We also use the Markov assumption for the policy. If the current state contains the relevant history, then
$$ p(a_t\mid s_0,a_0,\ldots,s_t,\theta) =p(a_t\mid s_t,\theta) =\pi_\theta(a_t\mid s_t). $$
For the two actions in our trajectory, this gives $$ p(a_0\mid s_0,\theta)=\pi_\theta(a_0\mid s_0), $$ and
$$ p(a_1\mid s_0,a_0,s_1,\theta)=\pi_\theta(a_1\mid s_1). $$
Putting It Together #
We can now substitute these assumptions into the chain rule factorization:
$$ \begin{aligned} p(\tau\mid\theta) ={}&p(s_0), \pi_\theta(a_0\mid s_0), p(s_1\mid s_0,a_0)\ &\times\pi_\theta(a_1\mid s_1), p(s_2\mid s_1,a_1). \end{aligned} $$ At each time step, the policy chooses an action and the environment produces the next state. The probability contributed by time step $t$ is therefore
$$ \pi_\theta(a_t\mid s_t), p(s_{t+1}\mid s_t,a_t). $$ A trajectory with $T$ actions starts at $s_0$ and ends at $s_T$:
$$ \tau=(s_0,a_0,s_1,a_1,\ldots,s_{T-1},a_{T-1},s_T). $$ Its probability is
$$ p(\tau\mid\theta) =p(s_0) \prod_{t=0}^{T-1} \pi_\theta(a_t\mid s_t) p(s_{t+1}\mid s_t,a_t). $$
Token Generation #
Now let us apply this factorization to a language model.
Suppose the initial prompt is $x_{1:n}=(x_1,\ldots,x_n)$. The initial state is $$ s_0=(x_1,\ldots,x_n). $$ The first action $a_0$ is the first token sampled from the model. Once it is generated, we append it to the prompt:
$$ s_1=(x_1,\ldots,x_n,a_0). $$ After the model generates $a_1$,
$$ s_2=(x_1,\ldots,x_n,a_0,a_1). $$ At time $t$, the state is
$$ s_t=(x_1,\ldots,x_n,a_0,\ldots,a_{t-1}), $$ and the next state is obtained by appending the chosen token:
$$ s_{t+1}=\operatorname{append}(s_t,a_t). $$ Once $s_t$ and $a_t$ are known, there is only one possible value of $s_{t+1}$. By definition, $s_{t+1}=(s_t,a_t)$. We can therefore read the transition probability as
$$ p(s_{t+1}\mid s_t,a_t) =p(s_{t+1}\mid s_{t+1}) =1. $$ Every environment transition term is equal to $1$ for a valid token trajectory. For our short example, the trajectory probability becomes
$$ p(\tau\mid\theta) =p(s_0), \pi_\theta(a_0\mid s_0), \pi_\theta(a_1\mid s_1). $$
For $T$ generated tokens, $$ p(\tau\mid\theta) =p(s_0) \prod_{t=0}^{T-1}\pi_\theta(a_t\mid s_t). $$ Usually, the prompt is already given. We condition on $s_0$, so we only need the probability of the generated tokens:
$$ p(a_0,\ldots,a_{T-1}\mid s_0,\theta) =\prod_{t=0}^{T-1}\pi_\theta(a_t\mid s_t). $$ Writing out the contents of each state gives
$$ p(a_0,\ldots,a_{T-1}\mid x_{1:n},\theta) =\prod_{t=0}^{T-1} \pi_\theta(a_t\mid x_{1:n},a_0,\ldots,a_{t-1}). $$ This is the trajectory probability used in the REINFORCE article. The policy gradient is derived there. Here, I wanted to derive the factorization itself and explain why the environment terms become $1$ for token generation.