Fork just before the model changes its mind

Why fork placement decides step level credit in tree structured reinforcement learning, and how belief shift branching offers a practical answer.

Read the original paper
Cover image for 'Fork just before the model changes its mind'

Fork Where the Model Changes Its Mind: Belief-Shift Branching for Tree-Structured Reinforcement Learning is a paper published on arXiv under identifier 2609.11061. The paper studies where to place forks in tree structured rollouts and proposes branching just before the step where the model's belief about the answer diverges most. A link to the original paper appears directly under the title of this post.

What interests me more than the headline result is the placement of the question. Most reinforcement learning papers ask how to give better rewards, while this one asks where to spend a fixed branching budget for the largest learning gain. So this post follows that question from the start. I will walk through why step level scores matter, why branching can stand in for those scores, and why the moment just before the mind changes turns out to be the answer.

Why I picked up this paper

Hello, this is Patrick.

I feel the same frustration whenever I look at training logs from reasoning models. A model writes a long solution, gets a reward of 1 when the final answer is right and 0 when it is wrong, and every one of the dozens of steps in between receives the same treatment. Whether there was a good idea in the middle or where the path went wrong leaves no trace in the record. So how can we tell which step in a long solution actually helped?

That question has been an old homework problem in reinforcement learning with verifiable rewards, usually shortened to RLVR. Learning from only the final answer gives a coarse signal. Some researchers proposed adding a separate critic that scores each step, and others proposed collecting human step labels. Both paths cost money. A critic makes the training system heavy, and human labels are hard to scale.

That is why I read this paper. The authors say step level signals are possible without a critic and without human labels. If you stop an existing chain in the middle, branch from there, and compare the sibling solutions, you can work backward to which step changed the result. People review their own work in a similar way. After failing an exam, we do not redo every page from the start when we write our error notes. We find where our thinking started to slip and redo the neighborhood of that step. This paper puts that review process inside the learning algorithm.

What final answers cannot tell you

Picture a model solving a math problem while you watch from the side. It expands expressions for ten steps, then copies a sign incorrectly at step seven. Everything after that is careful, but the result is wrong. Judged by the final reward alone, this solution is simply a failed zero. Yet fixing only step seven might have produced the right answer. The final score says nothing about that difference.

Researchers call this step level credit assignment. It means asking how much each step in a long chain contributed to the final result. With those scores, learning becomes much smarter. You can update the model toward fixing only the wrong step while keeping the good steps. Without those scores, the model learns in broad strokes, pushing down the whole failed solution and pushing up the whole successful one.

Here an interesting question appears. Would step scores not require step labels? At first glance, that seems necessary. Tree structured rollouts show another path. You stop at an intermediate step, sample several continuations from that point, and compare the success rates of those siblings. If siblings after some point solve the problem unusually often, the progress up to that point was probably pointing in a good direction. If siblings after some point keep getting stuck, the path had likely already gone wrong before that. I liked this move. It uses sampling itself as the measuring tool instead of adding a separate judge.

Why branching gives you step level credit

Let me walk through the tree rollout in more concrete terms. The model reads a problem and writes its first solution line by line. The training system picks a spot in that chain and marks a stop. From that spot, it draws several continuations at the same temperature. The sibling solutions share the early part and differ only in the later part. At the end, each sibling is checked for correctness, and the difference in success rates is used to estimate the value of that branching point.

The reason this estimate can stand in for a step score is surprisingly intuitive. Because everything before the fork is shared by all siblings, differences in success must come from choices made after the fork. Repeating this at different points builds a sense of which step moves the outcome most. It resembles sending several hikers down different paths at a trail junction and then learning which junction decided who reached the summit.

But the next problem follows immediately. Each fork costs sampling budget. Every extra continuation burns more tokens, runs more verification code, and holds more memory. Under a realistic budget, each chain gets only a few forks. So the authors put it this way. Once tree size is fixed, what decides the gain is not how many forks you place but where you place them. I nodded at that sentence. When my team looks at a new AI architecture, we do not look at benchmark numbers alone. We look at what actually changed and what cost structure that change creates in the system. If forks are expensive, the rule for placing them is the performance.

What happens when you fork where things have settled

So what does a bad fork look like? The answer is simpler than it sounds. When you fork where the outcome has already settled, the siblings mostly reach the same ending. If the answer was effectively decided by step eight, branching at step ten gives siblings that all succeed together or all fail together. With no difference in success, there is almost no signal about step value. You paid for expensive sampling and learned nothing.

Branching at a point that still wobbles tells a different story. Some siblings reach the right answer and others go wrong. That difference is the signal. It shows that the choice at that step changed the result. Fork placement therefore becomes the problem of finding where the signal is largest. Where should you branch so that sibling fates split most sharply?

This is where I would be careful not to over-read the experiment. A split among siblings does not mean that step was wrong. A good choice can split siblings too. Whether the split came from a bad move or a good move becomes clear only when you look at the average success rate of the siblings together. What the paper claims is about location, not direction. Find where the verdict splits first, then read the value of that place from sibling success rates. I think that restraint shows editorial care in the paper. It keeps the job of finding the turning point separate from the job of judging whether the turning point was good or bad.

Fork just before the mind changes

Now let me describe the proposed answer. The name is belief shift branching. The rule fits in one sentence. Place the fork just before the step where beliefs about the answer diverge most between consecutive steps. A large shift in belief suggests the internal judgment turned somewhere in between, and that neighborhood is likely the pivot that decides the result.

The authors phrase this intuition as locating pivots of the value curve. Imagine assigning each step the probability that progress so far leads to final success, then connecting those probabilities in order. The curve starts near fifty fifty, then at some step it jumps upward or bends downward. That bend is the pivot. The paper says to fork just before it. Spread several futures from the state before the bend, and the cause of the bend shows most clearly.

Teachers grading exam papers look at a similar place. They find the line where the answer suddenly changes, then read the expression just before it with care. They check whether a sign changed on that line, whether an assumption changed, or whether a different theorem was brought in. Measuring a model's belief change works in a similar way. The difference is that a person reads with eyes while the method turns the change into numbers with an internal signal or an outside probe.

So how do we actually turn a word like belief into numbers? I paused here because the phrase sounds plausible, yet measuring it felt slippery at first. That question unfolds across the next two parts, which makes sense because there are three methods and each carries a different cost and character.

Three ways to measure belief change without step labels

The paper offers three ways to measure belief change. All three share one property. None of them needs human step labels. The first is a black box probe. Without opening the model, it shows the partial solution up to some step and repeatedly asks how likely completion looks from there. It then measures how much that checkup answer moves as steps advance. Because the model stays a black box, this works even when the structure is unknown.

The second is a logit lens depth profile. It takes hidden states from middle layers, projects them onto the output distribution, and watches how the pull toward the right answer changes as steps advance. It reads when the answer related direction becomes sharp across stacked representations. The third is a learned activation direction. Before reinforcement learning starts, one direction vector is fitted offline, and during training the method reads how much the projection onto that direction moves between steps. It needs preparation cost up front, but reads cheaply during training.

I read this set as a sign of a practical team. Instead of insisting on one right answer, it offers choices for different settings. Teams that cannot open the model can use the first probe, and teams that can open it can use the second or third. (The cost numbers quoted in the paper use the first probe. As I will explain later, that signal only places forks and does not change the training objective itself.)

One confusion is worth removing here. A good probe signal does not itself solve the answer. A probe resembles a metal detector that tells you where to dig. A metal detector does not dig the gold for you, and likewise the probe only places forks while the actual reward still comes from verifiable correctness. Because of that division of labor, the whole training run does not collapse even when the probe is slightly noisy.

Why a one percent signal beats an expensive fork

This part raises a natural systems question. How would you actually serve this? Running a probe at every step could become its own burden. The paper answers with numbers. When the probe runs inside the rollout engine, placing forks costs around one percent of step compute on mathematics and under five percent on code. Compared with the sampling cost of a full fork, that is a small fraction.

That difference matters more than it looks. A fork generates whole continuations, so token costs are large. A probe is a light read that only picks a location. Digging the ground and scanning the ground before digging are different jobs. Spending a little on scanning to save one round of digging is a good trade. In settings like code tasks, where running verification is heavy, placing one fork badly costs even more. So the report that probe cost stays under five percent on code reads as welcome news for operations.

From a systems perspective, the next question is obvious. Should the probe sit on the same GPUs as the training loop, or should it attach to a rollout engine? The paper reports costs for the setup where the probe runs together with the rollout engine. The design reads belief change on the side without stopping generation. I think that choice shows good judgment. Keeping the policy being trained apart from the measuring probe stops small probe noise from polluting policy updates directly. The probe draws the map, while sibling correctness still does the actual surveying.

So what happens when the probe is wrong? I worried about this first because a fork that misses the pivot might seem to waste the signal entirely. But that worry is only half right. The neighborhood of a pivot already wobbles, so siblings still differ enough even when the fork lands one or two steps away. Of course repeated wild misses would be another story, which is why the ranking test in the next part matters to me. In this test we need to check whether this signal actually finds pivots better than its rivals.

A signal that ranked first in all eight panels

This is why I think the more important result is the structure behind the numbers. The paper tests how well the belief shift signal finds pivots of the value curve against Monte Carlo value curves. A Monte Carlo value curve means sampling many continuations at each intermediate step and directly measuring success rates. It is costly but trustworthy, so the test pits cheap signals against an expensive reference. The rivals include uncertainty signals like entropy, structural signals, and large model judges.

The result was clean. Across eight model by benchmark panels, the belief shift signal ranked first in every panel. That makes sense once you think about it. High entropy does not always mean the outcome splits, and unusual length or structure does not always mark a pivot. Model judges sound plausible, but for the specific job of placing forks, directly measuring movement in belief about the final answer proved more accurate.

I think the more instructive point is the distinction underneath. Entropy measures hesitation about the next token, while belief shift measures movement in judgment about the final answer. A shaky word habit and a changed mind are different events. Forks need the second one, so a signal aimed at the final answer wins. That may sound obvious now. Yet moving the fork placement discussion from a broad word like uncertainty to a concrete target like belief about the answer is a genuine clarification.

Still, hype needs a check here. First place in eight panels ranks signals, but it does not guarantee training gains. Finding pivots well and scoring higher after training on those pivots are claims on different levels. So the next part looks at what happened when the authors actually ran reinforcement learning.

What changes when you learn on the same budget

Choosing where to spend a fixed budget is the point of the reinforcement learning test. The paper compares belief shift forks against the strongest baselines at the same tree size. On mathematics aggregates with OLMo-3-7B as the base, the aggregate rose by 2.6 points and AIME 2026 rose by 2.9 points. On code, LiveCodeBench medium rose by 6.5 points on the OLMo code setup. Every gain came from moving fork locations under the same budget rather than adding more sampling.

These numbers need their conditions stated before anyone trusts them. The comparisons hold tree size, sampling budget, and verification rewards fixed. With unlimited budget, even random forks would eventually hit pivots. In practice, each chain allows only a few forks, so placement decides the score. It resembles drilling for oil with a limited number of wells. The number of holes matters less than where you drill.

What stayed with me is that the direction held in both mathematics and code. Methods that work in math often lose force in code. Verification works differently and solution shapes differ. Yet the belief shift yardstick held in both areas. The intuition that points where judgment wobbles are points with the most to learn carried across task types. Still, I would not stretch this to every reasoning task yet. The paper shows results on specific model and task groups, and settings outside that range still need checking.

So where could this trend break? I want to hold that question for the end because first we should clear up what this method is not, which makes the contribution easier to see. That is why the next part sorts out what this method is not: not a reward, and not a larger tree.

The remaining question is where not to fork

Saying what it is not explains it faster. Belief shift branching is not a smarter reward function. Rewards still come from verifiable correctness. It is a placement rule for forks. It is also not a larger tree. Larger trees may score higher, but they charge separately. This paper studies how to pick locations once size is fixed. It is not a substitute for human step labels either. Human labels could give richer signals, but this method shows a path that works without them.

Viewed from real systems, there are gaps too. A large belief shift does not always mean rich lessons. A model can move from one wrong conviction to another wrong conviction. Forks there will split siblings, but both sides may split toward wrong answers. The signal is large while the learning value is small. The paper eases this worry by checking agreement with Monte Carlo curves, but it does not remove the worry entirely. Some pivots picked by probes will always contain noise.

What I would watch next is whether the rule keeps its shape on longer chains and larger models. Longer chains can hold several pivots in one belief curve, which raises the new problem of which pivot to drill first. Larger models may move beliefs in subtler ways, which forces another look at probe resolution. As I always ask about structures, I also want to see what cost structure this creates across memory and bandwidth, latency and throughput. The rule for placing forks is light, but running whole trees is still not light.

So where can this trend be stopped in a good sense? I think it stops at knowing where not to fork. Skipping settled stretches in advance and saving sampling there, then spending budget only where things wobble, is the practical message I take from this paper. The next time I see a learning curve from a long solution, I will look for where the curve bends before I look at the score. Most of the learning happens just before the model changes its mind.

References

  1. Fork Where the Model Changes Its Mind: Belief-Shift Branching for Tree-Structured Reinforcement Learning · arxiv.org

    Reviewed source