💬 NLP

Think Fast and Slow: Step-Level Cognitive Depth Adaptation for LLM Agents

This paper introduces CogRouter, a framework that enhances LLM agent efficiency and performance on long-horizon tasks by dynamically adapting cognitive depth at each step through a two-stage training process grounded in ACT-R theory, achieving state-of-the-art results with significantly reduced token usage.

Original authors: Ruihan Yang, Fanghua Ye, Xiang We, Ruoqing Zhao, Kang Luo, Xinbo Xu, Bo Zhao, Ruotian Ma, Shanyi Wang, Zhaopeng Tu, Xiaolong Li, Deqing Yang, Linus

Published 2026-02-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruihan Yang, Fanghua Ye, Xiang We, Ruoqing Zhao, Kang Luo, Xinbo Xu, Bo Zhao, Ruotian Ma, Shanyi Wang, Zhaopeng Tu, Xiaolong Li, Deqing Yang, Linus

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

Imagine you are hiring a personal assistant to help you solve a series of complex problems, like organizing a house move, fixing a broken appliance, and planning a surprise party.

Currently, most AI assistants (Large Language Models) are stuck in one of two rigid modes:

The "Reflex" Assistant: They act instantly without thinking. Great for simple tasks like "turn on the light," but they often make mistakes on complex problems because they didn't plan ahead.
The "Over-thinker" Assistant: They spend hours writing a 50-page essay for every single step, even when the answer is obvious. They are brilliant at solving hard puzzles but incredibly slow and expensive to run, wasting energy on simple tasks like "pick up the cup."

The Problem: Real life isn't black and white. Sometimes you need a split-second reaction; sometimes you need a deep strategy. But current AI treats every moment the same way.

The Solution: COGROUTER
The paper introduces a new framework called COGROUTER (Cognitive Router). Think of it as teaching an AI assistant to be a smart driver who knows exactly when to cruise on autopilot and when to take full manual control.

The Four "Gears" of Thinking

Based on human psychology (specifically a theory called ACT-R), the researchers taught the AI to switch between four distinct "gears" of thinking:

Gear 1: The Reflex (Instinctive Response)
- Analogy: You touch a hot stove and pull your hand away instantly.
- AI Action: "The door is open, so I walk through it." No thinking required. Fast and cheap.
Gear 2: The Situational Check (Situational Awareness)
- Analogy: You are driving and see a red light. You check your speed and the traffic before stopping.
- AI Action: "I need to open the fridge, but first, let me check if it's locked." A quick, focused look at the current situation.
Gear 3: The Experience Review (Experience Integration)
- Analogy: You tried to fix a leaky faucet yesterday and failed. Today, you remember that specific trick didn't work, so you try a different tool.
- AI Action: "I tried looking in the kitchen for the thermometer, but it wasn't there. I should check the workshop instead." Learning from past mistakes.
Gear 4: The Grand Strategist (Strategic Planning)
- Analogy: You are planning a cross-country road trip. You map out the route, check gas stations, and plan for traffic jams before you even start the car.
- AI Action: "To solve this science experiment, I need to find a thermometer, then a substance, then a heat source. If I go to the wrong room first, I waste time. Let's map the best path." Deep, future-looking planning.

How It Learns: The "Confidence" Trick

The hardest part is teaching the AI when to switch gears. If you just tell it to "think harder," it will overthink everything. If you tell it to "be fast," it will be careless.

The researchers used a clever two-step training process:

The Classroom (Supervised Fine-Tuning): First, they showed the AI examples of experts solving problems, labeling each step with the correct "gear." This taught the AI what each gear looks like.
The Gym (Policy Optimization): This is the magic part. The AI plays the game, but the researchers use a special scoring system based on Confidence.
- The Rule: If the AI is using "Gear 1" (Reflex) and it predicts the next move with high confidence, it gets a reward. If it uses "Gear 4" (Strategist) for a simple task and is still confident, it's fine, but it wasted energy.
- The Result: The AI learns that if a task is easy, it should use the "Reflex" gear to save energy. If it gets stuck or the situation is complex, it naturally switches to "Strategist" gear to get the job done.

Why This Matters

The results are impressive. Using a relatively small AI model (Qwen2.5-7B), COGROUTER:

Solved more problems than massive, expensive models like GPT-4o or OpenAI's o3.
Used 62% fewer "tokens" (words/computational energy).
Was faster because it didn't waste time overthinking simple steps.

In a nutshell: COGROUTER stops AI from being a "one-trick pony." It teaches them to be economical thinkers, saving their brainpower for the hard moments and cruising on autopilot for the easy ones. It's the difference between a student who studies for 10 hours for a pop quiz and one who knows exactly how much effort to apply to get an A.

(reasoning process), and`.

Significance: This prevents the model from inheriting the expert's bias toward specific levels and establishes a "cold start" foundation for RL.

Stage 2: Cognition-aware Policy Optimization (COPO)

Goal: Optimize the selection of cognitive levels via Reinforcement Learning.
Core Innovation: Unlike standard Group Relative Policy Optimization (GRPO) which assigns a single advantage to an entire trajectory, COPO performs step-level credit assignment.
Mechanism:
1. Cognitive Group Expansion: For successful trajectories, the model generates "counterfactual" thinking processes for all four cognitive levels at each step, keeping the observation and action fixed.
2. Confidence-Aware Reweighting: The model calculates the prediction confidence (average log-probability) of the action under each of the four generated thinking processes.
3. Advantage Modulation: The trajectory-level advantage is redistributed to specific steps/levels based on this confidence. If a specific level (e.g., $L_1$ ) yields high confidence for a routine step, it receives a higher weight. If a complex step requires $L_4$ for high confidence, that level is reinforced.
4. Objective: Maximize the reward while encouraging the model to choose the minimum necessary cognitive depth that still guarantees high-confidence action prediction.

3. Key Contributions

Formalization of Cognitive Rigidity: Identified and formalized the inefficiency of fixed reasoning patterns in long-horizon agentic tasks, demonstrating that binary switching is insufficient for multi-level cognitive demands.
COGROUTER Framework: Introduced a novel framework based on ACT-R theory that defines four distinct cognitive levels and a two-stage training pipeline (COSFT + COPO) to enable dynamic adaptation.
COPO Algorithm: Developed a novel RL algorithm that solves the step-level credit assignment problem in sparse-reward environments by using confidence-aware advantage reweighting. This prevents the "mode collapse" to uniform deep thinking seen in standard RL methods.

4. Experimental Results

Experiments were conducted on ALFWorld (embodied household tasks) and ScienceWorld (scientific reasoning) using Qwen2.5-7B and Llama3.1-8B as base models.

Performance: COGROUTER achieved state-of-the-art (SOTA) results.
- Qwen2.5-7B: Achieved an 82.3% average success rate.
- Comparison: Outperformed GPT-4o (+40.3%), OpenAI-o3 (+18.3%), and the standard GRPO baseline (+14.0%).
Efficiency:
- COGROUTER reduced token consumption by 62% compared to GRPO and 57% compared to GiGPO.
- It significantly outperformed "thinking" models (like DeepSeek-R1) which consumed ~10x more tokens for lower success rates.
Behavioral Analysis:
- Adaptive Distribution: Unlike GRPO and GiGPO, which collapsed to >75% usage of deep thinking ( $L_4$ ) regardless of task complexity, COGROUTER learned to allocate resources dynamically (e.g., using $L_1$ for ~55% of steps in ALFWorld).
- Task Complexity Scaling: The model correctly increased the usage of $L_3$ and $L_4$ for complex tasks while relying on $L_1$ for routine steps.
- Convergence: COPO converged to high success rates significantly faster than other RL baselines.

5. Significance

Efficiency vs. Capability: COGROUTER demonstrates that high-level agent performance does not require uniform deep reasoning. By mimicking human "fast and slow" thinking, agents can achieve SOTA performance with drastically reduced computational costs.
Solving Credit Assignment: The paper provides a robust solution to the sparse reward problem in multi-step agents by introducing a confidence-based mechanism for granular credit assignment, preventing the collapse of adaptive behaviors in RL.
Scalability: The approach is model-agnostic and effective even on smaller open-source models (7B/8B parameters), making advanced agentic reasoning accessible without relying on massive proprietary models.

In summary, COGROUTER represents a paradigm shift from static, uniform reasoning to dynamic, context-aware cognitive allocation, enabling LLM agents to be both highly capable and computationally efficient.