Reasoning & Test-Time Compute · 2025

Atom of Thoughts for Markov LLM Test-Time Scaling

Reasoning & Test-Time Compute Atom of Thoughts for Markov LLM Test-Time Scaling 2025 · arXiv 2502.12018
Topic
Reasoning & Test-Time Compute
Venue
DeepWisdom · Renmin University) · Feb 2025
Read
16 min
Source
arXiv:2502.12018

In one line

Instead of dragging the entire reasoning history forward at every step, AOT repeatedly rewrites a hard question into a smaller, self-contained "atomic" question — so the model spends all its compute on the problem in front of it, not on re-reading its own past work.

The breakdown

TL;DR

Test-time scaling (making the model think longer at inference) works, but every popular method — Chain-of-Thought, Tree-of-Thoughts, Graph-of-Thoughts — carries the whole reasoning trace along as it goes. As the trace grows, the model burns tokens re-processing stale history and gets distracted by it. AOT (Atom of Thoughts) reframes reasoning as a Markov process: each step produces a brand-new question that is equivalent to the original but already has the easy parts baked in as known facts, so the next step depends only on the current question, not the history. The trick is a two-phase loop — decompose the current question into a dependency DAG of subquestions, then contract that DAG back into one simpler standalone question. Run it on gpt-4o-mini and you beat heavyweight reasoning models: 80.6% F1 on HotpotQA, +3.4% over o3-mini and +10.6% over DeepSeek-R1, at a fraction of the cost.

Problem & Motivation

The pain is concrete and it’s about context bloat in reasoning.

When you ask an LLM to “think step by step,” every new step is conditioned on all the previous steps plus the original question. Chain-of-Thought keeps the entire chain. Tree-of-Thoughts keeps ancestors and siblings so it can pick branches. Graph-of-Thoughts allows arbitrary node-to-node dependencies, which is even worse. The longer the reasoning, the more of the context window is occupied by stuff the model has already figured out.

This causes two distinct harms:

  1. Wasted compute. Re-attending over a growing history costs tokens (and dollars) on every step, and most of that history is solved and no longer load-bearing.
  2. Active interference. This is the subtler point. A bloated context doesn’t just cost money — it degrades reasoning. The model has to re-derive what matters from a pile of resolved detail, and the noise pulls it off track. The authors’ ablation makes a sharp version of this claim: imperfect structural guidance can be more harmful than no guidance at all.

The motivating analogy is human problem-solving (Simon, Polya). When you solve a multi-step problem, you don’t keep re-reading your scratch work. You solve a sub-piece, fold its result in as a known fact, and now you’re staring at a simpler problem — you’ve thrown away the derivation and kept only the conclusion. That “fold it in and forget how you got there” move is exactly a Markov transition: the next state depends only on the current state, not the path that led there.

What’s New (Core Contribution)

Three things, and only the first is genuinely novel mechanism:

  • Atom of Thoughts (the Markov reformulation). Before: reasoning frameworks model p(answer | all previous thoughts, original question) — history-dependent. Now: AOT engineers an explicit Markov chain of question-states Q0 → Q1 → ... → QD where each Q_{i+1} is a new, self-contained question equivalent to the original and depends only on Q_i. The reasoning history is structurally discarded at each step instead of accumulated. This is the heart of the paper.

  • The decompose→contract state-transition operator. Before: decomposition methods (Least-to-Most, Plan-and-Solve) split a problem into subquestions and then solve them in sequence, keeping all of them around. Now: AOT decomposes into a dependency DAG, then contracts — it absorbs the answerable “independent” subquestions into known conditions and rewrites the still-open “dependent” subquestions into a single new standalone question. Decomposition is temporary scaffolding, thrown away after each step. That contraction step is the new primitive.

  • Plug-in compatibility via “answer equivalence.” Before: test-time scaling methods (Self-Consistency, ToT, Forest-of-Thoughts) each operate on the raw original question. Now: because every Q_i is provably equivalent to Q0, you can stop AOT after a single decompose-contract cycle and hand the simplified question to any of those methods as a drop-in preprocessor. Result: AOT(d=1) + FoT(n=2) ≈ FoT(n=8) at much lower cost.

The “extensive evaluation” contribution is real but standard — six benchmarks, three runs each.

How It Works (Technically)

AOT is a loop. Each iteration takes the current question Q_i and produces a simpler equivalent question Q_{i+1}. The loop runs until the question is directly solvable, then a final LLM call answers it.

Demystifying the math

The paper frames everything probabilistically. Don’t let the notation scare you — it’s just describing what each method conditions on.

Chain-of-Thought (Eq. 1):

A ~ p(A | T, Q0) · ∏_i p(T_i | T_{<i}, Q0)

Plain English: each thought T_i is generated conditioned on all earlier thoughts T_{<i} and the original question. The T_{<i} term is the villain — it’s the accumulating history. Read it as “to write step i, re-read steps 0..i-1.”

The ideal Markov chain they want (Eq. 3):

A ~ p(A | Q_N) · ∏_i p(Q_{i+1} | Q_i)

Plain English: each new question depends only on the immediately previous question Q_i — the T_{<i} history term is gone. The final answer comes from the last (simplest) question alone. This is the goal. The whole paper is machinery to make a real LLM behave like this.

What AOT actually computes (Eq. 7):

A ~ p(A | Q_D) · ∏_i p(Q_{i+1} | G_i) · p(G_i | Q_i)

Plain English, two factors per step:

  • p(G_i | Q_i): from the current question, generate a dependency DAG G_i (the decompose call).
  • p(Q_{i+1} | G_i): from that DAG, generate the next simplified question (the contract call).

The DAG G_i is temporary — notice it never appears as a conditioning term in the next step. It’s scaffolding: built, used to produce Q_{i+1}, discarded. That’s the mechanism that buys the Markov property.

Phase 1 — Decomposition (build the DAG)

A single JSON-mode LLM call takes Q_i and emits a directed acyclic graph G = (Q, E):

  • Nodes = granular subquestions Q_i.
  • Edge (Q_j, Q_i) = “subquestion Q_j produces information needed to answer Q_i.”

Acyclicity is free: subquestions are emitted in natural-language order, so any subquestion can only depend on earlier ones. You can’t form a cycle without pointing forward, which the ordering forbids.

Two node categories fall out of the edge structure:

  • Independent subquestions Q_ind — no incoming edges. These are self-evident, answerable right now, with no prerequisites.
  • Dependent subquestions Q_dep — have incoming edges. These still need other pieces resolved first.

Phase 2 — Contraction (collapse the DAG into one new question)

A second single LLM call rewrites the DAG into one standalone question Q_{i+1}:

  • Treat the independent subquestions as solved — fold their results in as given conditions (or discard them if they were dead-end explorations).
  • Weave the dependent subquestions into the body of the new question.

The output is a question that (a) is answer-equivalent to Q_i, and (b) is simpler, because the easy parts are now stated as facts rather than left to be derived. Crucially, the derivation of the independent parts is thrown away — only their conclusions survive. That discarding is what keeps each state “atomic” and preserves the Markov property.

Termination and depth

  • The max depth D is set automatically to the longest path length in the very first DAG G_0 (a rule-based graph computation, no LLM). This prevents infinite decomposition.
  • There’s also an LLM-based early-stop: after each contraction, an LLM synthesizes an answer for Q_i from {original Q_i result, DAG G_i, the contracted-question Q_{i+1} result}. If that synthesized answer agrees with solving Q_{i+1} directly, keep going; the consistency check governs when to stop.
  • On termination, AOT assembles the final answer by combining the last contracted question with the union of all independent subquestions accumulated across iterations — giving a complete solution made entirely of independent pieces.

Architecture & data flow

flowchart TD
  Q0[Q0: original question] --> DEC[decomposeLLM:<br/>build dependency DAG G_i]
  DEC --> DAG{DAG G_i}
  DAG -->|nodes with no<br/>incoming edges| IND[Q_ind:<br/>independent / solvable now]
  DAG -->|nodes with<br/>incoming edges| DEP[Q_dep:<br/>still dependent]
  IND --> CON[contractLLM:<br/>fold Q_ind in as known facts,<br/>rewrite Q_dep into one question]
  DEP --> CON
  CON --> Q1[Q_i+1: simpler, equivalent,<br/>self-contained question]
  Q1 -->|loop: depends only on Q_i| DEC
  Q1 -->|terminate| SOLVE[solveLLM: answer Q_D]
  SOLVE --> A[Final answer A]
  Q1 -. plug-in entry point .-> EXT[Hand simplified question<br/>to ToT / FoT / Self-Consistency]

Schematic: how the model's context window fills up over reasoning steps. Toggle between a history-carrying method (CoT/ToT) and AOT. Watch the blue "live problem" bar stay constant under AOT while the gray "stale history" bar grows under the others — this is the wasted-compute / interference story from Figure 1, made concrete.

The algorithm, simplified

# Atom of Thoughts: reasoning as a Markov chain of self-contained questions.
# llm(...) is a single model call. Three roles, three prompts.

def aot(q0):
    q = q0
    max_depth = None
    accumulated_independent = []          # we keep CONCLUSIONS, not derivations

    i = 0
    while max_depth is None or i < max_depth:
        dag = decompose_llm(q)            # JSON: nodes=subquestions, edges=dependencies
        if max_depth is None:
            max_depth = longest_path(dag) # rule-based; caps infinite decomposition

        # split nodes purely by edge structure
        q_ind = [n for n in dag.nodes if dag.in_degree(n) == 0]   # solvable now
        q_dep = [n for n in dag.nodes if dag.in_degree(n) > 0]    # still blocked

        # CONTRACT: fold solved parts in as facts, rewrite the rest into ONE question.
        # The new question is equivalent to q but simpler. q's history is discarded here.
        q = contract_llm(known=q_ind, open=q_dep)
        accumulated_independent += q_ind

        if answers_agree(q):              # LLM consistency check -> early stop
            break
        i += 1

    return solve_llm(q, context=accumulated_independent)   # answer the atomic question

The Markov property lives in one line: q = contract_llm(...). After it runs, the old q and its derivation are gone — the loop body only ever reads the current q. That’s the whole trick.

Built on Prior Work

Prior ideaWhat it gaveWhat this paper changes
Chain-of-Thought (Wei 2022)Step-by-step reasoning in the promptStops carrying the full chain; each step is a fresh equivalent question
Least-to-Most / Plan-and-Solve (Zhou 2023, Wang 2023a)Decompose into ordered subquestionsAdds a DAG (not just a list) and a contraction step that discards solved subquestions
Tree / Graph of Thoughts (Yao 2023, Besta 2024)Explore multiple paths; richer structureUses structure (the DAG) as temporary scaffolding per step, not a persistent global object
Self-Consistency / Forest-of-Thoughts (Wang 2023b)Sample many traces, voteAOT becomes a preprocessor that hands these methods a simplified, equivalent question
Markov processes (Markov 1906)Memoryless state transitionsThe conceptual frame: reasoning state = a question; transition = decompose+contract

The lineage is “decomposition methods” — but the delta is treating decomposition as disposable and inventing the contraction operator that resets the state to be atomic each step.

Results & Evidence

Headline numbers (backbone = gpt-4o-mini-0718, 6 benchmarks, averaged over 3 runs):

  • HotpotQA (multi-hop QA): 80.6% F1 for base AOT, +7.1% over AFlow (the strongest framework baseline). With LLM answer-selection over 3 runs (AOT*), 81.0%.
  • LongBench: 68.5% F1, +7.5% over AFlow. Multi-hop is clearly where AOT shines — exactly the regime where history bloat hurts most.
  • MATH: 83.6% (AOT*, 84.9%); GSM8K: 95.0% (AOT*, 95.1%). Solid but smaller margins on math.
  • BBH: 86.0%, a big jump over CoT-SC (83.4%) and FoT (82.4%).
  • Average across all six: 80.8% for AOT vs 76.1% for the best baseline (AFlow).

The flashy claim: on HotpotQA, gpt-4o-mini + AOT (80.6%) beats o3-mini (77.2%) by 3.4% and DeepSeek-R1 by 10.6%. Putting o3-mini inside AOT pushes it to 81.4% F1 / 91.4% Hit — claimed new SOTA.

Scaling behavior (Figure 3): force up to 5 iterations; accuracy on MATH climbs 83.2% → 92.7% with depth, with diminishing returns. Most problems are solved shallow (1000 samples reach depth 1, only 207 reach depth 5) — a natural cost/quality knob.

Cost (Figure 4): AOT has the steepest performance-per-dollar slope of all methods compared; AOT(d=1)+FoT(n=2) matches FoT(n=8) for much less compute.

Ablations (Table 3): remove the contraction-feeding decomposition → drop; remove the DAG but keep decomposition → bigger drop. Insight: half-baked structure (single subproblem, no dependency graph) is worse than nothing because it breaks parallel relationships between subproblems.

What the evidence does NOT establish — read this part before you sell it:

  • Single backbone for main results. Everything rides on gpt-4o-mini. The cross-model comparison (Table 2) is only on the two multi-hop QA datasets, not the full suite.
  • The o3-mini comparison is a bit apples-to-oranges. A scaffolded gpt-4o-mini beating a bare-prompt reasoning model is a framework-vs-raw comparison, not model-vs-model. o3-mini’s training already bakes in long reasoning; AOT adds orchestration on top of a cheaper model.
  • Cost figures are token/dollar trend lines on MATH only, not wall-clock latency. AOT issues several sequential LLM calls per step (decompose, contract, consistency-check, solve), which is more round-trips even if total tokens drop.
  • The authors’ own stated limitation: no reflection/repair. A bad initial DAG poisons every downstream contraction with no recovery — and they say this “occurs frequently in practice.”
  • First-1000-examples subsets and a 100-example LongBench subset for the model comparison; small-sample noise is plausible on the multi-hop deltas.

How You’d Use It

For someone running an AI services shop, AOT is most interesting as a cheap-model uplift and a context-discipline pattern, not as a research artifact.

  • Lower your inference bill on multi-hop / RAG-heavy tasks. The whole value prop is “use a small model + structure to beat a big model.” If you’re paying for o-series or R1 calls on multi-hop QA or document-synthesis pipelines, an AOT wrapper around a mini model is a direct margin lever. Multi-hop QA is exactly the AOT sweet spot.

  • As a preprocessing stage in an existing agent. This is the most pragmatic move. Run one decompose-contract cycle to simplify a gnarly user question, then feed the cleaned-up question to whatever you already run (ToT, self-consistency voting, a RAG chain). You get most of the benefit without rebuilding your stack — that’s the AOT(d=1)+FoT pattern.

  • Context hygiene for long-running agents (MAS-relevant). If you’ve built multi-agent systems, you know the killer is context drift — agents re-reading bloated shared scratchpads. AOT’s core discipline (“after solving a sub-piece, replace the running state with a clean restatement that bakes the result in as a fact”) is a message-passing pattern you can adopt even without the full DAG machinery: have a “contractor” agent periodically rewrite the working state into a minimal self-contained brief.

  • A tangible client offering: “reasoning optimization audit” — take a client’s expensive reasoning pipeline, insert AOT-style contraction, and demonstrate equal-or-better accuracy at 2–5x lower token cost. The plug-in design means low integration risk.

Build Your Own (Minimal Recipe)

You can get ~80% of the value in an afternoon. Skip the DAG-path-length auto-depth and the consistency early-stop at first; just do the decompose→contract loop with a fixed small depth.

Components (all are prompts to one model):

  1. decompose(q) → JSON list of subquestions, each with a depends_on: [indices] field. (This is your DAG; no graph library needed — adjacency is in the JSON.)
  2. contract(subquestions) → rewrite into one standalone question, folding zero-dependency subquestions in as stated facts.
  3. solve(q) → answer the final atomic question.

Build order:

  1. Get decompose returning clean JSON (use the model’s JSON/structured-output mode). Validate the dependency indices.
  2. Write partition: in-degree 0 → independent, else dependent. Trivial Python.
  3. Get contract producing a shorter question that a human agrees is equivalent. This is the hard part.
  4. Loop 2–3 times, then solve.

The 1–2 genuinely hard parts:

  • Answer-equivalence in contraction. It’s easy to write a contract prompt that drops information and silently changes the problem. You need few-shot examples and ideally a verification step that confirms the contracted question is equivalent. Get this wrong and accuracy tanks.
  • Bad decompositions with no recovery. The paper’s own limitation. Without a reflection loop, one wrong dependency edge corrupts everything downstream.

Reach for: any model with reliable JSON mode (gpt-4o-mini, gpt-4.1-mini, Claude Haiku). Use a structured-output library (Instructor / Pydantic, or the provider’s native JSON schema) so the DAG parses deterministically. No graph DB, no fine-tuning, no RL — this is pure prompt orchestration.

How to Improve It

The limitations are the roadmap. Each of these is testable:

  1. Add the reflection/repair the authors admit is missing. After decomposition, run a critic that checks the DAG for missing parallelism or spurious dependencies before contracting. Re-decompose on failure (Reflexion-style verbal feedback). The ablation says bad structure is worse than none — so a cheap “is this DAG sane?” gate could be the single highest-ROI add.

  2. Verify answer-equivalence explicitly at contraction. Insert a lightweight check: solve Q_i and Q_{i+1} independently on easy cases and confirm they match; roll back the contraction if they diverge. Turns the silent failure mode into a caught one.

  3. Parallelize independent subquestions. Independent (in-degree 0) subquestions have no dependencies — they can be solved concurrently before contraction instead of folded in as assumed-known. This attacks the latency problem (many sequential calls) without changing the math.

  4. Learn the termination policy instead of the consistency heuristic. The current early-stop is an LLM agreement check. A small learned classifier (or even a confidence threshold) on “is Q_i directly solvable now?” would cut wasted iterations and is easy to train from logged runs.

  5. Combine with a verifier/process reward. AOT discards derivations to stay Markov — but that means errors in folded-in “facts” are unrecoverable. A verifier scoring each contracted question’s facts before they become permanent would harden the chain against compounding errors (the paper’s stated failure mode).

Glossary

  • Test-time scaling — improving model output by spending more compute at inference (longer reasoning, more samples), as opposed to training a bigger model.
  • Chain-of-Thought (CoT) — prompting the model to produce intermediate reasoning steps before the answer.
  • Tree / Graph / Forest of Thoughts — frameworks that explore multiple reasoning paths in a tree, graph, or ensemble-of-trees structure, keeping the structure around as it grows.
  • Markov process / Markov property — a sequence of states where each state depends only on the immediately preceding state, not the full history (“memoryless”).
  • Atomic question / atomic state — a self-contained question equivalent to the original problem but with solved parts baked in as known facts; depends on nothing prior.
  • DAG (Directed Acyclic Graph) — nodes connected by one-way edges with no cycles; here, subquestions with “needs-this-first” dependency edges.
  • Decomposition (in AOT) — one LLM call that breaks the current question into subquestions plus their dependency edges (builds the DAG).
  • Contraction (in AOT) — one LLM call that collapses the DAG into a single simpler equivalent question, folding independent subquestions in as conditions.
  • Independent / dependent subquestion — independent = no incoming edges (solvable now); dependent = has prerequisites still unsolved.
  • Answer equivalence — the property that each rewritten question Q_i has the same answer as the original Q0; what makes AOT safe to plug into other methods.
  • Plug-in enhancement — using AOT as a preprocessing step that simplifies a question before handing it to another reasoning method (e.g., FoT).
  • Ablation study — systematically removing a component to measure its contribution.
  • F1 / Hit rate — F1 is the harmonic mean of precision and recall (used for QA answer overlap); Hit = fraction of answers with F1 > 0 (any overlap at all).