Disciplined CoT: Tag-Driven Reasoning Distillation

• Disciplined CoT is a training framework that interleaves explicit control tags into chain-of-thought trajectories to curb overthinking and reasoning drift in small language models.
• It applies tags like <TEMP_LOW>, <TEMP_MID>, and <TEMP_HIGH> as structural scaffolds, guiding multi-step reasoning and dynamically modulating sampling temperature.
• Empirical results demonstrate up to 10% accuracy gains and significant token reductions on benchmarks such as MMLU-Pro and GPQA-diamond, underscoring its efficiency.

Disciplined Chain-of-Thought (D-CoT) is a training framework designed to mitigate "overthinking" and excessive token consumption in Small LLMs (SLMs) during Chain-of-Thought (CoT) distillation from large teacher models. D-CoT enforces a structured multi-step reasoning process by interleaving explicit control tags, such as <TEMP_LOW>, <TEMP_MID>, and <TEMP_HIGH>, as auxiliary scaffolding within the reasoning trajectory. This tag-based discipline optimizes the CoT trajectory, suppresses reasoning drift, and leads to both accuracy improvements and token reduction, as demonstrated on benchmarks such as MMLU-Pro and GPQA-diamond (Ubukata, 25 Feb 2026).

1. Formalism and Training Objectives

Let 𝒟 = {(xᵢ, τᵢ)}_{i=1}^N denote the set of N training examples where each xᵢ is a prompt and τᵢ is a Chain-of-Thought trajectory—a sequence τ = (c₁, c₂, …, c_L) of thought segments c_j, each representing a token subsequence. In standard CoT distillation, the objective is to maximize

$$\mathcal{L}_{\mathrm{SFT}}(\theta) = - \sum_{(x,\tau) \in \mathcal{D}} \log p_\theta(\tau \mid x).$$

D-CoT augments each trajectory by interleaving mode tags t ∈ 𝒯, producing an augmented trajectory

$$\bar{\tau} = (t_0, c_1, t_1, c_2, \ldots, t_{L-1}, c_L, t_L)$$

with control tag set 𝒯 = {<TEMP_LOW>, <TEMP_MID>, <TEMP_HIGH>} applied at segment boundaries. The model is trained to emit both content tokens and the corresponding mode tags.

D-CoT employs Odds Ratio Preference Optimization (ORPO) to directly optimize the log-probability margin between a preferred (“chosen”) tag-augmented response 𝒴⁺ and its “rejected” counterpart 𝒴⁻. For each prompt x:

$$\ell_{\mathrm{ORPO}}(\theta; x, \mathcal{Y}^+, \mathcal{Y}^-) = - [\beta \log p_\theta(\mathcal{Y}^{+}|x) - \log p_\theta(\mathcal{Y}^{-}|x)]$$

where β > 0 controls the strength of the preference. The combined dataset objective with optional weight decay R(θ) is

$$\mathcal{L}(\theta) = \sum_{i=1}^N \ell_{\mathrm{ORPO}}(\theta; x_i, \mathcal{Y}^+_i, \mathcal{Y}^-_i) + \lambda R(\theta).$$

No external reference model is required; R(θ) is typically an 𝓁₂-norm on the LoRA adapters (Ubukata, 25 Feb 2026).

2. Control Tag Semantics and Decoding Regimes

D-CoT operationalizes structured reasoning through three discrete, temperature-linked control tags:

Tag	Function	Sampling Temperature
<TEMP_LOW>	Fact-checking, constraint listing	0.3
<TEMP_MID>	Algorithmic processing, calculation	0.6
<TEMP_HIGH>	Creative exploration, multi-viewpoint	0.8

During training, segment tags precede their corresponding reasoning steps in all preferred responses. At inference, tags can either be supplied explicitly to direct reasoning and sampling temperature, or omitted if the model has internalized the discipline. Dynamic decoding based on detected tags adaptively sets the sampling temperature per the control tag—see Algorithm 1 in (Ubukata, 25 Feb 2026).

3. End-to-End Training Methodology

Data Sources and Teacher Outputs:

Training utilizes seven non-overlapping domains, such as Supply Chain Logistics and Cybersecurity Incident Response, to prevent benchmark data contamination. Teacher outputs (from Qwen3-235B) include:

• user_prompt
• thought_chosen + chosen_response (explicit tags)
• thought_rejected + rejected_response
• reasoning summary (for QA only)

Rejection categories (31 in total) are used to enforce fine-grained model discrimination, including both tag-usage and domain-related errors. Post-filtering for high similarity (cosine > 0.55 or 13-gram overlap), 5,079 samples remain.

Optimization Details:

• Base model: Qwen3-8B with LoRA adapters.
• Optimizer: Lion (8-bit).
• Learning rate: 4 × 10⁻⁶.
• β_ORPO = 0.1.
• Batch size = 1, gradient accumulation = 8, bfloat16 precision.
• Warmup ratio = 0.1, cosine decay scheduler.
• Epochs = 2.

Each mini-batch involves computing the log-probabilities for both chosen and rejected responses, calculating the per-sample ORPO loss, and updating θ via the Lion optimizer.

4. Theoretical Motivation

SLMs experience "reasoning-drift" or "overthinking" under long CoT trajectories due to per-step errors ε compounding as

$P(\text{correct}) \approx (1-\epsilon)^L,$

with performance peaking and then declining as L increases (inverted U-curve). D-CoT aims to bound this drift by structurally segmenting reasoning with control tags, thereby localizing errors and suppressing undirected loops. This approach is hypothesized to reduce effective trajectory length per error accumulation event and to improve local segment accuracy, increasing the overall product ∏(1–ε_j) across segments. Although the paper provides no formal error-bound theorem, these structural insights underpin the framework's design (Ubukata, 25 Feb 2026).

5. Empirical Evaluation and Results

Benchmarks:

• MMLU-Pro (12,000 multi-step, 10-way, 0-shot)
• GPQA-diamond (198 graduate-level, 4-way, 5-seed, 0-shot)

Key Results (Qwen3-8B):

Benchmark	Configuration	Accuracy	Avg Tokens	Null-rate
MMLU-Pro	Base (T=0.6, Base Prompt)	55.66%	1742	—
MMLU-Pro	D-CoT (Locked, Base)	64.73%	1496	—
GPQA-diamond	Base (Locked, Base Prompt)	43.03%	5875	30.91%
GPQA-diamond	D-CoT (Dyn, Custom Prompt)	52.93%	2153	3.54%

D-CoT achieves a +9.07 point accuracy gain and 14.2% token reduction on MMLU-Pro, and a +9.90 point accuracy gain and 63.4% token reduction on GPQA-diamond relative to baseline models. In both cases, D-CoT points dominate baseline configurations with respect to the accuracy–token Pareto frontier (Ubukata, 25 Feb 2026).

6. Ablation Studies and Operational Insights

Dynamic temperature decoders provide only marginal (+0.5–1.5 points) further improvement over locked-temperature strategies when tags are present, indicating that the disciplinary effect is primarily a function of training structure, rather than runtime entropy control. Notably, omitting control tags at inference—after training with D-CoT tags—yields the highest MMLU-Pro accuracy (64.73%), demonstrating that the model internalizes disciplined reasoning and no longer requires explicit token-level scaffolding when deployed.

7. Limitations and Prospects

Current evaluations are limited to Qwen3-8B and the specified benchmarks, with no broad analysis of generalization bounds or a detailed study of the impact of control tag set size and composition. Potential directions for further research include extending D-CoT to additional SLM architectures and larger parameter scales, analyzing attention maps for evidence of discipline internalization, and experimenting with expanded or more specialized tag sets (e.g., for safety or bias-checking) and finer-grained mode control.

D-CoT represents a methodologically lightweight but structurally impactful enhancement to conventional CoT distillation. It leverages tag-based scaffolding and preference optimization to suppress overthinking, minimize token usage, and promote the internalization of disciplined reasoning trajectories, ultimately enabling efficient and robust SLM inference even in the absence of explicit control tokens at deployment (Ubukata, 25 Feb 2026).