Reasoning Trees

[diatkinson]

Research Question

Do reasoning models represent their own reasoning tree? In particular, can we identify representations of "success probability" at reasoning nodes?

Owners

David Atkinson (@diatkinson), Chris Wendler (@wendlerc), Harsh Raj (@harshraj172)

Project Status

(very early) work in progress

Rough Plan

1. Begin by generating a "syntax tree" via output entropy: generate until the entropy reaches a certain threshold. Then, add branches until the remaining entropy falls below the threshold. That gives you a tree of possible-generations.
2. Annotate the nodes of the syntax tree with their probability of success. At each node, rollout along the syntax tree multiple times, and record the percentage of successful rollouts.
3. Convert the annotated syntax tree to a "reasoning tree" by removing nodes which have similar success probabilities to their parents.
4. Armed with the reasoning tree we can ask questions like:
   • Does the model (linearly?) represent its chance of success? (If not, does a bigger model?)
   • At a reasoning node, is it possible to probe for candidate next nodes? What about their chances of success? Can we use this to edit the reasoning tree?
   • To what extent does the reasoning tree align with our intuitive understanding of the CoT? Do nodes align with "wait" moments?
   • Given some measure of node similarity, do models implement DFS? BFS? Something else?

Progress

So far, all I (@diatkinson) have done is curate a small dataset of GSM8K questions that DeepSeek-R1-Distill-Qwen-1.5B finds difficult (gets between 25% and 75% correct). Right now, I'm tweaking the hyperparameters for syntax tree generation.

[wattenberg]

Can you say more about step 1? I have a vague idea of what this might mean, but I'm not exactly sure. It sounds interesting!

[diatkinson]

Thanks! Yeah, I need a better name than "syntax tree". Really, I'm just trying to generate a tree of plausible completions for a single prompt. In the current approach (which I've tweaked a bit from the description above), I generate the the tree by, at each step, selecting max(1, min(max_branches, round(alpha * exp(beta * entropy(softmax(logits)))))) tokens to branch off of. So when logits are sufficiently low entropy, I just continue (greedily) on the current branch. When they're high I create a bunch of branches, and proceed through them as before until I sample a EOS.

There are plenty of other ways to do this of course, so if people have suggestions, I'd love to hear them.

[firstuserhere]

Convert the annotated syntax tree to a "reasoning tree" by removing nodes which have similar success probabilities to their parents.

There is a risk that important intermediate reasoning steps are removed if they are too similar to their parent.

[ARBORproject]

Hello - just checking in if this project is still active. To keep our project statuses accurate, otherwise we would like to switch the project status to Inactive until there is activity again!