{"slug": "probability-of-an-llm-trajectory-starting-from-the-chain-rule", "title": "Probability of an LLM Trajectory: Starting from the Chain Rule", "summary": "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.", "body_md": "# Probability of an LLM Trajectory: Starting from the Chain Rule\n\n          In the [REINFORCE\n          article](llm-token-generation.html#reinforcement-learning), we jumped into the gradient. But I wanted to take a step\n          back and start by factorizing the probability of a trajectory from\n          first principles.\n        \n\nLet us consider a short trajectory:\n\n$$\n\\tau=(s_0,a_0,s_1,a_1,s_2).\n$$\nWe start in state $s_0$, take action $a_0$, move to state $s_1$, take action $a_1$, and arrive at $s_2$.\n\nWhat is the probability of this entire trajectory under a policy with parameters $\\theta$?\n\nFor 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.\n\n## Start with the Chain Rule\n\nWe can begin with the chain rule. For three random variables, the chain rule says\n\n$$\np(x,y,z)=p(x)p(y\\mid x)p(z\\mid x,y).\n$$\nApplying the same rule to our trajectory gives\n\n$$\n\\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}\n$$\nSo far, we have only used the chain rule. We have not made any assumptions about the policy or the environment.\n\n## The Initial State\n\nThe 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\n\n$$\np(s_0\\mid\\theta)=p(s_0).\n$$\nFor a language model, $s_0$ is the initial prompt. The prompt is given to the model before generation begins.\n\n## The Environment\n\nThe next state is produced by the environment after the policy chooses an action. We assume that the environment is not parameterized by $\\theta$. Therefore,\n\n$$\np(s_1\\mid s_0,a_0,\\theta)=p(s_1\\mid s_0,a_0),\n$$\nand\n\n$$\np(s_2\\mid s_0,a_0,s_1,a_1,\\theta) =p(s_2\\mid s_0,a_0,s_1,a_1).\n$$\nThis 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$.\n\nWe 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,\n\n$$\np(s_2\\mid s_0,a_0,s_1,a_1)=p(s_2\\mid s_1,a_1).\n$$\nAt an arbitrary time $t$, this is\n\n$$\np(s_{t+1}\\mid s_0,a_0,\\ldots,s_t,a_t) =p(s_{t+1}\\mid s_t,a_t).\n$$\nFor 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.\n\n## The Policy\n\nThe 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\n\n$$\np(a_t\\mid s_t,\\theta) \\overset{\\mathrm{def}}{=} \\pi_\\theta(a_t\\mid s_t).\n$$\nThis is just notation, but it helps us distinguish the probabilities produced by the policy from the probabilities produced by the environment.\n\nWe also use the Markov assumption for the policy. If the current state contains the relevant history, then\n\n$$\np(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).\n$$\nFor the two actions in our trajectory, this gives\n\n$$\np(a_0\\mid s_0,\\theta)=\\pi_\\theta(a_0\\mid s_0),\n$$\nand\n\n$$\np(a_1\\mid s_0,a_0,s_1,\\theta)=\\pi_\\theta(a_1\\mid s_1).\n$$\n## Putting It Together\n\nWe can now substitute these assumptions into the chain rule factorization:\n\n$$\n\\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}\n$$\nAt each time step, the policy chooses an action and the environment produces the next state. The probability contributed by time step $t$ is therefore\n\n$$\n\\pi_\\theta(a_t\\mid s_t)\\, p(s_{t+1}\\mid s_t,a_t).\n$$\nA trajectory with $T$ actions starts at $s_0$ and ends at $s_T$:\n\n$$\n\\tau=(s_0,a_0,s_1,a_1,\\ldots,s_{T-1},a_{T-1},s_T).\n$$\nIts probability is\n\n$$\np(\\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).\n$$\n## Token Generation\n\nNow let us apply this factorization to a language model.\n\nSuppose the initial prompt is $x_{1:n}=(x_1,\\ldots,x_n)$. The initial state is\n\n$$\ns_0=(x_1,\\ldots,x_n).\n$$\nThe first action $a_0$ is the first token sampled from the model. Once it is generated, we append it to the prompt:\n\n$$\ns_1=(x_1,\\ldots,x_n,a_0).\n$$\nAfter the model generates $a_1$,\n\n$$\ns_2=(x_1,\\ldots,x_n,a_0,a_1).\n$$\nAt time $t$, the state is\n\n$$\ns_t=(x_1,\\ldots,x_n,a_0,\\ldots,a_{t-1}),\n$$\nand the next state is obtained by appending the chosen token:\n\n$$\ns_{t+1}=\\operatorname{append}(s_t,a_t).\n$$\nOnce $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\n\n$$\np(s_{t+1}\\mid s_t,a_t) =p(s_{t+1}\\mid s_{t+1}) =1.\n$$\nEvery environment transition term is equal to $1$ for a valid token trajectory. For our short example, the trajectory probability becomes\n\n$$\np(\\tau\\mid\\theta) =p(s_0)\\, \\pi_\\theta(a_0\\mid s_0)\\, \\pi_\\theta(a_1\\mid s_1).\n$$\nFor $T$ generated tokens,\n\n$$\np(\\tau\\mid\\theta) =p(s_0) \\prod_{t=0}^{T-1}\\pi_\\theta(a_t\\mid s_t).\n$$\nUsually, the prompt is already given. We condition on $s_0$, so we only need the probability of the generated tokens:\n\n$$\np(a_0,\\ldots,a_{T-1}\\mid s_0,\\theta) =\\prod_{t=0}^{T-1}\\pi_\\theta(a_t\\mid s_t).\n$$\nWriting out the contents of each state gives\n\n$$\np(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}).\n$$\nThis 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.", "url": "https://wpnews.pro/news/probability-of-an-llm-trajectory-starting-from-the-chain-rule", "canonical_source": "https://www.sithankanna.com/posts/trajectory-probability.html", "published_at": "2026-09-13 10:56:04+00:00", "updated_at": "2026-09-13 11:05:01.478154+00:00", "lang": "en", "topics": ["large-language-models", "machine-learning", "artificial-intelligence"], "entities": ["REINFORCE"], "alternates": {"html": "https://wpnews.pro/news/probability-of-an-llm-trajectory-starting-from-the-chain-rule", "markdown": "https://wpnews.pro/news/probability-of-an-llm-trajectory-starting-from-the-chain-rule.md", "text": "https://wpnews.pro/news/probability-of-an-llm-trajectory-starting-from-the-chain-rule.txt", "jsonld": "https://wpnews.pro/news/probability-of-an-llm-trajectory-starting-from-the-chain-rule.jsonld"}}