Cognition-aware Supervised Fine-tuning (CoSFT)

• The paper introduces CoSFT, a supervised learning approach that instills structured, level-specific cognitive patterns in LLMs for step-wise decision-making.
• It enforces a structured triple output—cognitive level, reasoning text, and action—using randomized level sampling and maximum likelihood estimation.
• Empirical results highlight that while CoSFT reliably formats outputs and guides multi-level reasoning, it increases token usage and needs downstream policy optimization for adaptive efficiency.

Cognition-aware Supervised Fine-tuning (CoSFT) is a supervised learning methodology designed to instill stable, level-specific cognitive patterns in LLM agents operating in step-wise, multi-turn decision-making environments. Originating within the CogRouter framework, CoSFT is grounded in the ACT-R theory of human cognitive architecture and implements a structured multi-level reasoning scheme that precedes any policy optimization or reinforcement learning. The primary purpose of CoSFT is to enable models to produce structured outputs that reflect varying cognitive depths, providing both the syntax and initial balance necessary for downstream adaptive reasoning strategies (Yang et al., 13 Feb 2026).

1. Theoretical Foundation: ACT-R Hierarchy and Design Motivation

CoSFT is motivated by the observation that long-horizon, sequential tasks (such as ALFWorld and ScienceWorld) exhibit significant “step-wise heterogeneity” in cognitive requirements; some steps necessitate only reflexive actions while others demand detailed planning or integrative reasoning. Drawing on the ACT-R cognitive architecture (Anderson 1982, 1993, 1998; Taatgen 2013), CoSFT structures reasoning into four discrete levels (Sec 3.2):

• Level 1 (Instinctive Response): Immediate, rule-based actions with no explicit deliberation.
• Level 2 (Situational Awareness): Assessment of current state and action space before response.
• Level 3 (Experience Integration): Retrieval and integration of prior outcomes into current decisions.
• Level 4 (Strategic Planning): Prospective simulation of downstream effects for candidate actions.

This hierarchy underpins the multi-level reasoning format of CoSFT. Empirical analysis revealed two critical requirements: without explicit, level-specific exemplars, models either disregard structured formats or collapse to a single favored reasoning level, and biases in pretrained LLMs impede the recovery of balanced, stable cognitive patterns even with subsequent RL. CoSFT thus precedes policy optimization, ensuring initial exposure and proficiency at all reasoning levels (Yang et al., 13 Feb 2026).

2. Structured Supervised Learning Pipeline

The CoSFT stage formalizes output as a structured triple $y_t = [l_t, th_t, a_t]$ at each time step $t$, where $l_t$ is the cognitive level (enclosed in <level> tags), $th_t$ the corresponding reasoning text (in >), and $a_t$ the final action (in <action>). This tripartite format is enforced by system prompts and sample templates (App A.3).

Data Collection and Annotation

The training dataset $\mathcal{D}_{cog}$ is generated as follows (Sec 3.3, App B.1):

1. Sample 500 training environments per benchmark.
2. Deploy an expert agent (GPT-4o) to generate successful trajectories $$\tau^\ast = \{x, (o_0,a_0^\ast),\dots,(o_T,a_T^\ast)\}$$.
3. At each time step $t$, randomly select a level $l_t \in \{1,2,3,4\}$ (ensuring $\approx$25% coverage of each).
4. Prompt the expert to produce the reasoning $t$0 for $t$1 using level-specific templates.
5. Aggregate all quadruples $t$2 into $t$3.

Loss Function and Optimization

Fine-tuning is performed using standard maximum likelihood estimation (MLE, cross-entropy) over the structured output sequences (Sec 3.3):

$t$4

where the three output segments are concatenated and predicted auto-regressively.

Training Configuration

Key hyperparameters and choices (App B.2, Table 6):

• Optimizer: AdamW
• Learning rate: $t$5
• Batch size: 32
• Epochs: 3
• Max sequence length: 8,192 tokens
• Base models: Qwen2.5-7B-Instruct, Llama3.1-8B-Instruct
• Data size: $t$625,000 step-level examples (500 environments $t$7 $t$850 steps)

Randomized, rather than expert-driven, level selection is crucial for balanced initialization (“cold-start”) and to avoid early convergence to biased strategies.

3. Empirical Outcomes and Comparative Analyses

Table 1 (Sec 4.2) presents offline evaluation of SFT-only (single-level, ReAct-style), CoSFT (random-level), and CoSFT_exp (expert-level sampled) variants on Qwen2.5-7B:

Method	Success Rate (%)	Avg Tokens per Step
SFT	62.5	494.1
CoSFT	57.0	3912.7
CoSFT_exp	78.5	3151.0

SFT, utilizing almost exclusively shallow behaviors, achieves moderate success rates with minimal token usage but lacks any mechanism to enforce multi-level structure or consistency. CoSFT, through random-level supervision, enforces valid structured output but increases token counts due to unnecessary overthinking on many steps, negatively impacting performance metrics. CoSFT_exp, by mimicking the expert's level distribution (with a bias toward Levels 2 and 3), boosts success rates but retains inefficiencies—overusing moderate reasoning even on trivial steps, and inheriting teacher allocation biases.

Usage distributions visualized in Figure 1 (Sec 4.2) confirm that neither variant achieves the context-sensitive, adaptive cognitive depth required for optimal performance; uniform (CoSFT) or expert-skewed (CoSFT_exp) level allocations remain inflexible and agnostic to real-time task demands.

Qualitative examples (App B.1) demonstrate that CoSFT reliably induces correct three-part format outputs and level-specific linguistic patterns, but models do not yet adaptively select cognitive depth by context.

4. Design Choices and Ablation Studies

Random vs. Expert Level Sampling: Random sampling for level assignment (CoSFT) produces a balanced distribution across all reasoning levels, essential for enabling policy optimization to later “discover” when and how much to think via reinforcement learning (Sec 6.2, Table 4). Expert-driven sampling (CoSFT_exp) produces level collapse—particularly toward Level 3—reducing diversity and impeding effective exploration during downstream RL.

Format Consistency: To ensure syntactic invariance, all outputs must follow strict <level>, <think>, <action> tagging. During CoSFT, misformatted examples are filtered out, resulting in near-zero format errors. In subsequent CoPO, a formatting reward penalty $t$9 is used if deviations occur.

Role of CoSFT SFT Stage: Ablation experiments wherein the “cold start” (CoSFT) is skipped result in catastrophic format collapse and irrecoverable biases toward deep thinking during RL—models fail to learn structured multi-level outputs, and downstream adaptive reasoning does not emerge.

5. Practical Implications and Limitations

CoSFT is foundational for enabling hierarchical, multi-level reasoning in autonomous LLM agents (Sec 4, Sec 6). It encodes the “syntax” of variable cognitive depth and supplies stable, balanced exemplars for each reasoning mode. However, CoSFT alone does not endow the agent with the ability to choose context-appropriate levels; it “over-thinks” on many steps due to enforced uniformity. As a result, CoSFT, while indispensable, must be paired with a downstream optimization stage (CoPO) to yield context-adaptive, efficient cognitive profiling. The observed inefficiency and overuse of deep reasoning in the standalone CoSFT regime articulate the boundary between syntactic structure acquisition and adaptive cognitive control.

6. Integration into Adaptively Reasoning Agents

CoSFT provides supervised anchors for policy optimization algorithms that subsequently learn context-sensitive allocation of cognitive depth (Sec 3.4). By establishing both the format and level-specific linguistic bases for reasoning steps, CoSFT enables reinforcement learners (e.g., CoPO), to perform per-step credit assignment, confidence calibration, and efficiency tuning. Experimental results on the ALFWorld and ScienceWorld benchmarks underscore this synergy: only the combination of CoSFT and adaptive policy optimization produces agents with both superior stepwise reasoning accuracy and dramatically reduced computation (62% fewer tokens) relative to leading baselines (Yang et al., 13 Feb 2026). The workflow demonstrates the centrality of CoSFT as the precursor stage in hierarchical, cognition-aware agent training pipelines.