Probability of an LLM Trajectory: Starting from the Chain Rule

A technical explainer derives the probability of a large language model trajectory from first principles using the chain rule, factorizing p(τ|θ) into the initial state distribution p(s0), policy terms πθ(at|st), and environment transition probabilities p(st+1|st,at). The derivation assumes the initial state distribution and environment are independent of the policy parameters θ and applies the Markov assumption to both the environment and the policy. The piece frames token generation as a policy that assigns a probability to each possible next-token action.

Probability of an LLM Trajectory: Starting from the Chain Rule 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.