🤖 AI

Diverse Thinking Schemata Elicit Better Reasoning in Large Language Models

The paper proposes DiScO, a framework that enhances large language model reasoning by explicitly defining and optimizing for diverse thinking schemata—comprising reasoning transitions and answer candidates—through a three-stage process of schemata awareness, reinforcement learning, and diverse inference, resulting in superior performance on mathematical benchmarks and improved error recovery compared to standard methods.

Original authors: Xinyue Liang, Yizhe Yang, Yu Bai, Bin Xu, Jiawei Li, Yang Gao

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xinyue Liang, Yizhe Yang, Yu Bai, Bin Xu, Jiawei Li, Yang Gao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Don't Just Think Harder, Think Differently

Imagine you are trying to solve a very tricky maze. Most AI models (Large Reasoning Models) are like a single runner who runs down one path, hits a wall, turns around, and tries to run down that same path again, just a little faster. They get stuck in loops or give up because they are too focused on their first idea.

This paper argues that to solve complex problems, AI shouldn't just run faster; it needs to be a better explorer. The authors introduce a concept called "Thinking Schemata." Think of this as the AI's internal "map-making" style.

They found that the AI performs best when its map-making style is diverse. Specifically, they looked at two things:

Reasoning Transitions: How often the AI changes its mind or switches strategies (e.g., "Wait, that math didn't work, let's try drawing a picture instead").
Answer Candidates: How many different possible answers the AI considers along the way before picking the final one.

The paper claims: The more the AI changes its mind and tries different angles, the more likely it is to find the correct answer.

The Problem: The "Stubborn Runner"

Current AI models often get stuck in a rut. If they start solving a math problem a certain way, they tend to keep doing it, even if they realize halfway through that they made a mistake. They might say, "I think the answer is 12," then realize 12 is wrong, but instead of trying a totally new approach, they just tweak the 12 slightly and say, "Okay, maybe 12.1?"

They lack the flexibility to say, "Okay, I was wrong. Let's throw this whole idea in the trash and start fresh from a different angle."

The Solution: DiScO (Diverse Schemata Policy Optimization)

The authors created a new training method called DiScO. You can think of DiScO as a coach that teaches the AI how to be a flexible explorer. It works in three steps:

1. Teaching the AI to "Label Its Thoughts" (Schemata-Aware SFT)

First, the AI is taught to wear a "thinking cap" that helps it see its own thought process. It learns to tag its own reasoning with special labels:

\AnswerCandidate: "I think the answer might be X."
\ReasoningTransition: "Wait, I'm changing my strategy now."

It's like teaching a student to write "I'm stuck" or "Let's try a different method" in the margins of their homework. This makes the AI aware of how it is thinking, not just what it is thinking.

2. Rewarding the "Change of Mind" (Diversity-Oriented RL)

Next, the AI plays a game where it gets points not just for getting the right answer, but for being diverse.

If the AI tries three different ways to solve the problem, it gets a bonus.
If it switches from "doing math" to "drawing a diagram," it gets a bonus.
If it keeps repeating the same sentence over and over, it gets no points.

This is like a coach telling a runner: "If you just run the same lap, you get a C. But if you try running through the woods, then the river, then the hills, you get an A+." The AI learns that exploring different paths is valuable.

3. The "Fresh Start" Trick (Inference-Time Diversity)

Finally, when the AI is actually taking a test (inference time), the system has a safety net. If the AI starts rambling or getting stuck in a loop (repeating itself), the system cuts off the beginning of its thought process and says, "Okay, forget the first 20% of what you just said. Start over with what you have left."

This forces the AI to "reset" and look at the problem with fresh eyes, preventing it from getting trapped in a mental dead end.

What Happened? (The Results)

The researchers tested this on hard math problems (like those found in high-level competitions).

Better Scores: The AI trained with DiScO got significantly higher scores than standard AI models, especially on the hardest problems.
Better Recovery: When the AI made a mistake early on, the DiScO-trained model was much better at realizing, "Oh no, I went the wrong way," and switching to a new path to fix it. Standard models often just kept walking down the wrong path.
More Options: The DiScO models didn't just guess one answer; they explored many different "Answer Candidates" before settling on the right one.

The Bottom Line

The paper concludes that diversity is key. Just like a human solving a hard puzzle benefits from trying different strategies, changing perspectives, and not sticking to the first idea, AI models perform much better when they are encouraged to think in varied, flexible ways.

The authors suggest that the future of AI reasoning isn't just about making models bigger or faster, but about making them more diverse thinkers.

Technical Summary: Diverse Thinking Schemata Elicit Better Reasoning in Large Language Models

Problem Statement

Large Reasoning Models (LRMs) have demonstrated significant capabilities in solving complex mathematical problems by generating extended chains of thought. However, current approaches often rely on overly predefined reasoning structures, such as strictly linear chains or explicitly organized trees, which may limit cognitive flexibility. The authors identify two critical yet underexplored aspects of the reasoning process: Reasoning Transitions (the distinct shifts between intermediate reasoning steps, e.g., "Let me try another approach") and Answer Candidates (the variety of solution paths or intermediate conclusions proposed by the model). The paper posits that existing models often converge on narrow reasoning patterns, failing to fully explore the solution space, which limits their ability to recover from errors or find optimal solutions in complex tasks.

Methodology: DiScO Framework

To address these limitations, the authors propose DiScO (Diverse Schemata Policy Optimization), a framework designed to enhance the diversity of "thinking schemata" (the latent reasoning structure governing transitions and solution paths) during both training and inference. The framework consists of three core components:

1. Schemata-Aware Supervised Fine-Tuning (SFT)

The model is first equipped with "schemata awareness" by training it to explicitly annotate its own reasoning chains. Two specialized tokens are introduced:

\ReasoningTransition: Marks shifts in reasoning perspective or strategy.
\AnswerCandidate: Marks potential final answers or intermediate conclusions.
This step renders the latent thinking structure observable and amenable to analysis without requiring external human annotation for every training step (annotations are generated by an LLM annotator, specifically Qwen3-Max).

2. Diversity-Oriented Reinforcement Learning (RL)

Following SFT, the model undergoes reinforcement learning using a modified Group Relative Policy Optimization (GRPO) algorithm. The core innovation is a composite reward function ( $R$ ) that balances accuracy with diversity:
$R = W_{div}R_{div} + W_{acc}R_{acc} + W_{form}R_{form}$

Diversity Reward ( $R_{div}$ ): Encourages the generation of distinct reasoning transitions and answer candidates. It measures the count of transitions ( $N_{thought}$ ), the number of answer candidates ( $N_{ans}$ ), the uniqueness of solutions ( $N_{uniq}$ ), and the correctness of intermediate candidates ( $N_{true}$ ). Truncation caps are applied to prevent reward hacking via marker repetition.
Accuracy Reward ( $R_{acc}$ ): Standard binary reward for a correct final answer.
Format Reward ( $R_{form}$ ): Ensures the model adheres to the required annotation tags and final answer formatting.

3. Inference-Time Diversity Enhancement (ITDE)

To further promote diversity during inference, two strategies are employed when the model reaches the maximum generation length without a final answer:

Initial Truncation: Discards the earliest 20% of tokens to simulate a "cognitive reset," allowing the model to reframe the problem from a fresh perspective.
Truncation with Repetition Elimination: Detects and removes repetitive reasoning patterns (e.g., the same 15-word phrase) to break local loops and encourage exploration of new paths.

Key Contributions

Concept Definition: The paper introduces "Thinking Schemata," defining them through observable Reasoning Transitions and Answer Candidates. It empirically demonstrates a positive correlation between the diversity of these schemata and model performance across various LRMs.
DiScO Framework: The authors propose a novel framework that integrates schemata-aware SFT, diversity-oriented RL, and inference-time strategies to systematically promote diverse reasoning.
Empirical Validation: Extensive experiments across five mathematical reasoning benchmarks and three out-of-domain (OOD) benchmarks demonstrate that DiScO consistently outperforms standard GRPO and other diversity-oriented baselines.

Experimental Results

The authors evaluated DiScO on DeepSeek-R1-Distill-Qwen models at 7B and 32B scales, comparing them against frontier proprietary models (e.g., GPT-5, Claude 3.5) and open-source baselines.

Performance Gains: DiScO consistently achieved state-of-the-art results among open-source models. For instance, DiScO-7B achieved 95.6% on MATH-500 and 70.0% on AIME 2024, significantly outperforming other 7B models. DiScO-32B reached 83.5% average accuracy across benchmarks, surpassing the base DeepSeek-R1-Distill-Qwen-32B.
Hard Task Improvement: Gains were most pronounced on challenging competition-level benchmarks (AIME, AMC), where substantial room for improvement exists. DiScO-32B showed a 6.7% absolute improvement over the base model on AIME 2024 and 2025.
Generalization: Despite being trained exclusively on mathematical data, DiScO generalized effectively to out-of-domain tasks (ARC-c, GPQA-diamond, MMLU-Pro), suggesting that schemata-level diversity promotes transferable reasoning behaviors.
Error Recovery: Human-annotated analysis revealed that DiScO significantly improves the model's ability to recover from erroneous initial attempts (measured by the FWFC metric: First Wrong, Final Correct). DiScO models were less likely to discard correct intermediate answers and more likely to self-correct compared to baselines.
Pass@k Performance: DiScO showed improved performance in multi-attempt scenarios (Pass@k), indicating a genuine expansion of the reasoning boundary rather than just reshaping the single-shot distribution.

Significance and Claims

The paper claims that the diversity of thinking schemata plays a crucial role in successful problem solving, acting as a complementary axis to standard model scaling (parameters and data). The authors argue that:

Diversity is a Scalable Resource: Promoting diverse thinking schemata allows models to explore alternative inferential routes and recover from errors more effectively, leading to better performance on complex tasks without necessarily increasing model size.
Beyond Token-Level Diversity: Unlike previous methods that focus on token-level entropy or trajectory-level sampling across independent rollouts, DiScO targets diversity within a single reasoning trajectory, fostering semantic shifts and perspective changes that mirror human-like problem-solving.
Practical Impact: The framework offers a practical path for improving open-source reasoning models, narrowing the gap with frontier proprietary systems, and enhancing the interpretability of reasoning processes through explicit schemata annotation.

The authors acknowledge limitations, including the non-trivial cost of LLM-based annotation, the modest increase in inference latency (approx. 10% more tokens), and the current focus on mathematical reasoning, though they note consistent gains in out-of-domain tasks. They suggest that future work could explore cheaper annotation methods (e.g., probing hidden states) and the application of these principles to other reasoning regimes.