An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning

arXiv:2504.18587 · cs.LG, cs.AI, stat.ML · Submitted 2026-08-09 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning".

Jane: The paper was written by Tianbing Xu from.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Alright, welcome back to the show, everybody. Today we are digging into a fresh arXiv paper that's been getting a lot of buzz in the AI world. It's called "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning." Jane, what's the first thing that jumps out at you from that title?

Jane: Honestly, Tom, the title is a mouthful, but it's actually a really clever framing. It's taking this old, well-established statistical technique called Expectation-Maximization, or EM, and applying it to a very modern problem: how we teach large language models to think better.

Tom: Right, and that's the key, right? We're not just talking about a new trick. We're talking about a whole new way of looking at the problem. Usually, we think about training these models with algorithms like PPO or GRPO, which are these complex, engineered beasts.

Jane: Exactly. And those methods work, but they're messy. They rely on things like importance weighting and clipping to keep the training stable. It's like trying to steer a car by constantly jerking the wheel back and forth instead of just turning it smoothly.

Tom: So this paper is saying, "Hey, let's step back. Let's use a cleaner mathematical lens." And that lens is EM. The idea is that the reasoning process, the chain of thought the model generates, is a hidden variable. We can't see the "right" reasoning, we only see the final answer and the reward it gets.

Jane: And EM is perfect for that. It's an algorithm designed to find hidden structure in data. You alternate between guessing what the hidden structure is and then improving your model based on that guess. It's a beautiful, principled way to think about it.

Tom: And the implications are huge. If this works, it could mean simpler, more stable, and more efficient training for reasoning models. That's not just a small tweak; that could change how everyone does this.

Jane: For sure. And it's not just about making the training simpler. The paper shows that this approach can actually make the models reason better, and even produce shorter, more concise chains of thought. Who doesn't want a model that's both smarter and faster?

Tom: I'm already excited to see the actual math and the results. But before we get into the nitty-gritty, I want to bring in our resident engineer, Meng. Meng, from a practical standpoint, what does a title like this promise? Does it just sound good, or does it actually mean something for people building these systems?

Meng: It promises a lot, Tom. If it delivers, it means we can throw away a lot of the fragile engineering that goes into stabilizing these RL training runs. That's a huge win for anyone trying to train a model on a limited budget. Less code to debug, less hyperparameter tuning, more time actually building the product.

Tom: So it's not just academic. It's potentially a real engineering win.

Meng: Potentially, yeah. That's the promise. Whether it holds up is what we're here to find out.

Jane: And that's the perfect segue. Let's not just take the title's word for it. Let's look at what the paper actually claims to have done.

Tom: Let's do it. We're just getting started with "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning." Stick around.

Summary: Tom: So we're back, and we're still talking about "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning." Jane, we've set the stage. Now, what's the one-sentence summary of what these authors actually pulled off?

Jane: They've taken the core idea of EM and built a new, simpler reinforcement learning algorithm called EM Policy Gradient, or EMPG, that trains reasoning models by directly optimizing a reward-weighted log-likelihood, without all the extra machinery that PPO and GRPO need.

Tom: And that "extra machinery" is the importance weighting and clipping we mentioned. The paper makes a really strong point that those things add variance and bias, especially when you're dealing with long reasoning traces. It's like adding noise to your steering wheel.

Meng: Right, and the variance issue is real. When you have a long chain of thought, the probability of that entire sequence under the new policy can become astronomically small, making the importance weight explode. Clipping is a band-aid to stop that explosion, but it also biases the gradient.

Jane: Exactly. And EMPG sidesteps the whole problem. The E-step is like a smart sampling phase. You generate a bunch of reasoning paths with the current model, and you score them with the reward function. The M-step is then just a simple weighted fine-tuning step, where you push the model to make the high-reward paths more likely.

Tom: So instead of trying to correct for the difference between the old and new policies, you just... generate, score, and learn. That's so clean.

Meng: It is clean, but my engineer brain immediately asks, "How is that stable?" If you just do multiple gradient updates on the same batch of trajectories, aren't you going to overfit to that batch and wreck the policy?

Jane: That's a great question, and the paper addresses it. They call it a "one-step off-policy" method. The key is that the buffer of trajectories is only from the immediately preceding iteration. So the distribution mismatch is small. You're not reusing old data from many steps ago.

Tom: And the results seem to back that up. They tested it on GSM8K and a hard subset of MATH, and EMPG matched or beat GRPO across a range of Qwen2 point 5 models, from 1 point 5B all the way up to 14B.

Meng: And the fact that it beat GRPO on some of the harder models, like the Math-7B, where GRPO just failed to learn under the same setup, is a pretty strong statement. It suggests EMPG is more robust, not just simpler.

Tom: Robust and simple. That's a killer combo. But Jane, you mentioned it produces shorter reasoning traces. Why is that such a big deal?

Jane: Because inference cost is real money and real time. If a model can get the same answer in three hundred tokens instead of seven hundred that's a massive speedup and cost saving in production. The paper shows EMPG consistently generates more concise reasoning, which is a huge practical advantage.

Tom: So we have a simpler algorithm, better or equal performance, and more efficient outputs. This is sounding almost too good. I need to hear more about the theoretical guarantees.

Jane: That's exactly where we're headed. The paper doesn't just show it works in practice; it proves it should work in theory. Let's get into that next.

Improvements: Tom: Welcome back. We're still on "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning." So Jane, we've covered the "what" and the "how." Now, what's the big theoretical improvement that this paper brings to the table?

Jane: The big one is the monotonic improvement guarantee. In plain English, the paper proves that if you do the E-step and M-step correctly, every single iteration of the algorithm is guaranteed to make the expected reward go up. You never take a step backward.

Meng: That's a huge deal. In standard policy gradient methods, you're always hoping the update improves things, but there's no guarantee. You can easily take a bad step and ruin a perfectly good model. A monotonic guarantee is like having a safety net.

Tom: And it's not just a hand-wavy argument. The proof is in the appendix. It leverages the classic EM property where the lower bound on the objective is tight after the E-step, and the M-step is guaranteed to increase that lower bound.

Jane: Right. It's a beautiful, elegant result. The math just works out because they've formulated the problem correctly. The reward-weighted posterior distribution is the key. The E-step finds that distribution, and the M-step fits the model to it.

Lu: And this is where the perspective shift is so powerful. It's not just about making RL more stable. It's about re-framing the entire goal. Instead of viewing it as a policy optimization problem with all its pitfalls, you view it as a latent variable estimation problem, which has a much richer theoretical toolbox.

Tom: So Lu, you're saying this opens the door to using other latent variable techniques?

Lu: Absolutely. Think about it. Once you have this EM framing, you can start borrowing ideas from variational inference, from Bayesian methods. You could have a richer posterior, not just a single point estimate. You could think about structured latent variables, not just a flat sequence of tokens. The possibilities are enormous.

Meng: Okay, that's the dream. But let's get back to the practical improvement for a second. The paper also mentions a reward transformation and a baseline for variance reduction. Can you break that down?

Jane: Sure. The raw reward is often just a binary zero or one which is a very sparse signal. So they standardize the rewards within a mini-batch and then push them through a sigmoid function. This smooths out the signal and makes it easier for the model to learn the relative quality of different trajectories.

Meng: And the baseline is a classic REINFORCE trick. You subtract the mean reward from each sample. This reduces the variance of the gradient estimate without introducing bias. It's a standard, reliable technique.

Tom: So they're not just throwing out a new algorithm; they're also incorporating best practices to make it work in the real world. That's the kind of rigor I like to see. And it clearly pays off, given the results we saw on GSM8K and MATH Hard.

Jane: It does. And the results are even more interesting when you look at the emergent behaviors. The models trained with EMPG aren't just getting the right answer; they're showing signs of real reasoning.

Tom: Like what? What kind of behaviors are we talking about?

Jane: We're talking about the model decomposing a problem into subproblems, writing Python code to verify its own intermediate steps, and even backtracking when it realizes it's on the wrong path. It's like watching a student learn to think.

Tom: That's wild. We're not just teaching it to parrot answers; we're seeing it develop problem-solving strategies. I can't wait to see those examples. Let's take a closer look at the first page of the paper to see how they set all this up.

First Page: Tom: We're back on the show, and we're taking a closer look at the opening of "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning." Jane, the first page really sets the tone. What's the core problem they're identifying right out of the gate?

Jane: They're pointing out that the current state-of-the-art algorithms, like PPO and GRPO, are really just general-purpose policy optimization methods that have been adapted for LLMs. They're not designed for reasoning. And because of that, they carry all this baggage, like the importance weighting and clipping we talked about.

Tom: And they frame it as a question, right? They ask, "Can reinforcement learning for reasoning be formulated from a different optimization perspective that admits a simpler learning algorithm?" It's a direct challenge to the status quo.

Lu: And that's the right question to ask. The field has been so focused on scaling up these existing algorithms that we haven't stopped to ask if we're using the right tool for the job. This paper is essentially saying, "Maybe there's a more natural way."

Meng: And from my perspective, the first page also makes a promise about what that simpler algorithm would look like. It says it would eliminate the need for trajectory-level importance weighting and heuristic clipping. That's a very concrete promise that I can get behind.

Tom: It's a bold claim. And they back it up by saying their algorithm, EMPG, directly optimizes the EM surrogate objective. It's not just a hack to make things work; it's a principled derivation.

Jane: Exactly. And they also set up the key insight on that first page: the reasoning process is a latent variable. We don't observe the "correct" chain of thought, we only see the final answer and the reward. So they're treating the reasoning trace as a hidden variable to be inferred.

Lu: And that's the fundamental shift. It moves the problem from "how do I optimize this policy" to "how do I infer the hidden reasoning process that leads to good outcomes." And that's a much more tractable problem.

Meng: It also means the algorithm is naturally off-policy. You generate data with the old policy, then you learn from it. You don't have to worry about the data being "stale" because you're not trying to correct for the distribution shift. You're just fitting a model to a weighted dataset.

Tom: So on the first page alone, they've identified the problem, proposed a new perspective, and hinted at a simpler algorithm. It's a really strong opening. And it makes me want to see if they can actually deliver on that promise.

Jane: And they do. The experiments we discussed show that it works. But more than that, the first page also hints at the bigger picture. It's not just about a new algorithm; it's about a new way of thinking about how we teach models to reason.

Tom: And that new way of thinking could have implications far beyond just math problems. If we can teach a model to structure its own reasoning, to verify its own work, to backtrack when it's wrong, that's a step towards more general intelligence.

Lu: It's a step towards models that don't just generate text, but actually think. And that's a profound idea.

Tom: It really is. But before we get too philosophical, let's wrap this up. We've covered a lot of ground today. Let's bring it all together in our conclusion.

Conclusion: Tom: Alright, we're in the final stretch of our discussion on "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning." Jane, can you give us the big picture one more time?

Jane: Sure, Tom. This paper takes a well-known statistical method, EM, and applies it to the problem of training reasoning models. The result is a new algorithm, EMPG, that's simpler, more stable, and often more effective than the current standard, GRPO.

Meng: And from my side, the practical impact is clear. It's less engineering complexity, less hyperparameter tuning, and it produces models that are both more accurate and faster at inference because they generate shorter reasoning traces.

Lu: And the theoretical impact is just as significant. The monotonic improvement guarantee is a powerful property, and the EM framing opens up a whole new avenue of research. We can start borrowing ideas from the vast literature on latent variable models.

Tom: And we can't forget the emergent behaviors. The models aren't just getting the right answer; they're learning to decompose problems, verify their work, and backtrack from mistakes. That's a glimpse of something much bigger.

Jane: It's a glimpse of models that are learning to think, not just predict. And that's why this paper feels so important. It's not just a tweak; it's a new foundation.

Tom: So, to the authors of "An Expectation-Maximization Perspective on Reinforcement Learning for LLM Reasoning," we say thank you. You've given us a lot to think about, and we're excited to see where this line of research goes.

Meng: And I'm excited to try implementing it. The simplicity is going to be a breath of fresh air.

Tom: That's the spirit. Alright, that's our show for today. We've said goodbye to this paper, and we're ready to dive into the next one. Thanks for listening, everyone.

Tianbing Xu

cs.LG, cs.AI, stat.ML

Submitted: 2026-08-09

Updated: 2026-08-11

Comments: 20 pages

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 57/100

The gist: "We present EM Policy Gradient (EMPG), a simple and principled framework that formulates reinforcement learning over reasoning trajectories as an Expectation-Maximization (EM) problem [5, 6]." The

Key concepts

Expectation-Maximization (EM): EM is a statistical algorithm used to find hidden structure within data. It operates by alternating steps: first, guessing the underlying hidden structure (the E-step), and then updating the model parameters based on that guess (the M-step). This process helps infer patterns when parts of the data are unobserved.
Reinforcement Learning for LLMs: This is a training method designed to teach large language models to think better. Instead of just predicting the next word, the models are trained using rewards for their generated 'chain of thought,' enabling them to learn problem-solving strategies and complex reasoning paths.
EM Policy Gradient (EMPG): EMPG is a specific, simpler algorithm derived from EM principles. It trains reasoning models by generating potential reasoning paths, scoring them with a reward function (E-step), and then performing a weighted fine-tuning step (M-step) to make high-reward paths more likely.
Monotonic Improvement Guarantee: This is a theoretical proof that guarantees the algorithm's success. It ensures that every single iteration of the training process, when performed correctly, will increase the model's expected reward. This prevents the model from ever taking a step backward in its learning process.

Terminology

GRPO — 55× in this paper · explained in 6 episodes
Reinforcement Learning — 47× in this paper · explained in 4 episodes
Benchmark — 11× in this paper · explained in 2 episodes
Verification — 11× in this paper · explained in 2 episodes
Trace — 9× in this paper · explained in 2 episodes
Reasoning Traces — 8× in this paper · explained in 3 episodes
Regularization — 8× in this paper · explained in 3 episodes
Fine-Tuning — 7× in this paper · explained in 5 episodes
KL Regularization — 7× in this paper · explained in 2 episodes
Accuracy — 6× in this paper · explained in 1 episode
Structured Reasoning — 4× in this paper · explained in 3 episodes
Core — 3× in this paper · explained in 2 episodes
Reward Function — 3× in this paper · explained in 2 episodes
Chain-of-Thought (CoT) — 2× in this paper · explained in 4 episodes
Latent Reasoning — 2× in this paper · explained in 2 episodes
Verifier — 2× in this paper · explained in 2 episodes
Adaptation — 1× in this paper · explained in 1 episode
Chain-of-Thought (CoT) Prompting — 1× in this paper · explained in 3 episodes
Chain-of-Thought (CoT) Reasoning — 1× in this paper · explained in 2 episodes
Chain-of-Thought Prompting — 1× in this paper · explained in 2 episodes
Distribution Shift — 1× in this paper · explained in 6 episodes
Group Relative Policy Optimization — 1× in this paper
International Conference on Machine Learning — 1× in this paper
Iterative Refinement — 1× in this paper · explained in 3 episodes
Large Language Models (LLMs) — 1× in this paper · explained in 6 episodes
Steering — 1× in this paper · explained in 2 episodes
Test-time Scaling — 1× in this paper · explained in 2 episodes
Theoretical Guarantees — 1× in this paper · explained in 2 episodes
Verifiable Rewards — 1× in this paper · explained in 2 episodes

Summary

Summary

The paper introduces EM Policy Gradient (EMPG), a reinforcement learning framework for training reasoning language models, formulated as an Expectation-Maximization (EM) optimization problem. The authors state: We present EM Policy Gradient (EMPG), a simple and principled framework that formulates reinforcement learning over reasoning trajectories as an Expectation-Maximization (EM) problem [5, 6].

The core formulation models the reasoning process as a latent rationale sequence generated autoregressively by a policy, with the objective of maximizing expected reward over complete reasoning trajectories. The authors derive a variational lower bound on the log objective, showing that the lower bound is tight if and only if q = qθ⋆, where qθ⋆ is the reward-weighted posterior trajectory distribution. The EM algorithm alternates between an E-step, which computes qt = arg max L(θt, q) = qθ⋆t, reweighting trajectories by their rewards, and an M-step, which updates parameters by θt+1 = arg max Eτ∼qt [log Pθ(τ)].

The paper establishes a monotonic improvement guarantee: "Theorem 1 (Monotonic Improvement of EM Policy Gradient). Let J(θ) = Eτ∼Pθ [R(τ)], where R(τ) > 0. Suppose that the E-step computes the exact posterior qt = qθ⋆t, and the subsequent M-step satisfies L(θt+1, qt) ≥ L(θt, qt). Then J(θt+1) ≥ J(θt)."

The practical EMPG algorithm approximates the E-step by sampling multiple trajectories per query using stochastic decoding and storing reward-annotated trajectories in a replay buffer. The M-step performs reward-weighted empirical risk minimization with gradient estimator ∇θ Lt(θ) = Eτ∼Dt [∇θ log Pθ(τ)R(τ)]. The authors emphasize that The only difference between Eq. (8) and the proposed EMPG gradient in Eq. (7) is the absence of the importance weight, and that EMPG avoids both trajectory-level importance weighting and heuristic clipping, yielding a simpler and lower-variance optimization procedure.

The algorithm includes a reward transformation that standardizes rewards within mini-batches and applies a sigmoid function, plus variance reduction via a mini-batch mean baseline. The authors note: "Unlike PPO and GRPO, EMPG does not require likelihood-ratio clipping, trajectory-level importance weighting, a separate value model, or generalized advantage estimation. Its optimization consists only of trajectory generation, reward evaluation, and reward-weighted language-model fine-tuning."

Experiments were conducted on GSM8K and MATH Hard benchmarks using Qwen2.5 base models ranging from 1.5B to 14B parameters. The results show EMPG achieves performance comparable to or better than GRPO. Specifically, on GSM8K, EMPG scores were 0.744, 0.888, 0.918, and 0.952 for Qwen2.5-1.5B, Qwen2.5-7B, Qwen2.5-Math-7B, and Qwen2.5-14B-Instruct, respectively, compared to GRPO scores of 0.739, 0.867, 0.004 (where GRPO failed to learn in this setting), and 0.940. On MATH Hard, EMPG scores were 0.575, 0.742, 0.794, and 0.808, compared to GRPO scores of 0.568, 0.730, 0.755, and 0.816.

The paper also examines reasoning length, finding that Despite achieving comparable or better reasoning performance, EMPG consistently generates shorter reasoning trajectories than GRPO. For example, with Qwen2.5-1.5B on GSM8K, EMPG converges to an average response length of approximately 120–180 tokens, whereas GRPO typically generates 200–500 tokens. On MATH Hard with Qwen2.5-Math-7B, EMPG stabilizes around 400–600 tokens, compared with roughly 700–900 tokens for GRPO.

Qualitative analysis of reasoning trajectories reveals emergent behaviors including subproblem decomposition, self-verification, reflection, and backtracking. The authors note that "the models exhibit self-reflection and backtracking: after detecting inconsistencies in an initial solution, they revisit earlier reasoning steps, correct errors, and pursue alternative strategies before producing the final answer."

The paper also discusses KL regularization toward a reference model, finding that the effectiveness of KL regularization depends on the strength of the base model. For weaker models such as Qwen2.5-1.5B, we find that setting β = 0 yields strong performance, while stronger base models such as Qwen2.5-14B-Instruct benefit from KL regularization.

The authors conclude: "More broadly, this work suggests that reinforcement learning for reasoning can be viewed through the lens of latent-variable optimization rather than solely as policy optimization. We hope that the proposed EM formulation provides a foundation for developing new optimization algorithms, richer variational objectives, and more effective reinforcement learning methods for reasoning language models."

Improvements for AI systems

Based on the paper, here are the specific improvements I can implement in an AI system:

Implementation: Replace the policy optimization core with EMPG's two-step process:

E-step: Sample N trajectories per query (temperature 1.0), store all with rewards in a replay buffer (no importance weighting, no clipping)
M-step: Optimize reward-weighted log-likelihood with mini-batch standardization (sigmoid transform) and mean baseline subtraction

What the improved system can do:

Achieve comparable or better accuracy on math reasoning (GSM8K: 0.918 vs 0.004 for GRPO on Qwen2.5-Math-7B; MATH Hard: 0.794 vs 0.755)
Eliminate the need for trajectory-level importance weights and likelihood-ratio clipping, reducing variance and optimization bias
Guarantee monotonic improvement of expected reward under exact E/M steps (Theorem 1)

These improvements collectively enable a reasoning system that is simpler, more efficient, and more robust than current PPO/GRPO-based approaches, with demonstrated gains on mathematical reasoning benchmarks.