long-cot-training

name	long-cot-training
title	Through the Valley: Path to Effective Long CoT Training for Small LLMs
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2506.07712
keywords	["chain-of-thought","small-models","training-instability","error-accumulation","scaling"]
description	Navigate Long CoT Degradation phenomenon when training small models on extended reasoning, understanding recovery dynamics and implementing strategies to maintain performance.

Through the Valley: Path to Effective Long CoT Training for Small LLMs

Core Concept

Small language models (≤3B parameters) suffer from Long CoT Degradation when trained on extended chain-of-thought data with insufficient examples. Performance deteriorates sharply early in training, then gradually recovers—but smaller models often fail to return to baseline even with 220k examples. The underlying cause is error accumulation: longer outputs increase the probability of compounding mistakes throughout the reasoning chain. Understanding this "valley" and its recovery dynamics enables strategies to train effective small-model reasoners without falling into the degradation trap.

Architecture Overview

• Error Accumulation Mechanism: Longer reasoning chains multiply per-token error probabilities
• Model-Size Dependent Recovery: Larger models recover faster and more completely
• Training Instability Detection: Monitor degradation valley depth and recovery curve
• Reflection Pattern Recognition: Surface-level reflection adoption without genuine reasoning
• Data Scaling Requirements: Determine minimum examples needed for stable recovery
• RL Integration: Downstream RL training sensitive to pre-training CoT trajectory

Implementation

Step 1: Analyze Error Accumulation Dynamics

Implement diagnostic tools to understand degradation:

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
from collections import defaultdict

class CoTErrorAnalyzer:
    """Analyze how errors accumulate in long chain-of-thought reasoning"""

    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def compute_token_error_rate(self, generated_cot, target_cot):
        """
        Compute per-token error rate in CoT sequence.

        Error rate = (incorrect_tokens / total_tokens)
        """
        gen_tokens = self.tokenizer.encode(generated_cot)
        tgt_tokens = self.tokenizer.encode(target_cot)

        # Align sequences
        min_len = min(len(gen_tokens), len(tgt_tokens))
        matches = sum(1 for i in range(min_len) if gen_tokens[i] == tgt_tokens[i])

        token_error_rate = 1.0 - (matches / min_len)

        return token_error_rate

    def compute_sequence_accuracy(self, generated_seq, target_seq, error_rate_per_token=0.1):
        """
        Model sequence accuracy as product of per-token probabilities.

        Accuracy(seq_len) = (1 - error_rate_per_token) ^ seq_len

        This explains why longer sequences are exponentially more likely to fail.
        """
        seq_length = (.tokenizer.encode(target_seq))

        accuracy = ( - error_rate_per_token) ** seq_length

         accuracy

     ():
        
        baseline_acc = evaluate_on_dataset(model, eval_dataset)
        trajectory = {: [], : [], : []}

         step  (num_steps):
            
            batch = ((train_dataset))
            train_one_step(model, batch)

            
             step %  == :
                current_acc = evaluate_on_dataset(model, eval_dataset)
                avg_cot_length = compute_avg_cot_length(model, eval_dataset)

                trajectory[].append(step)
                trajectory[].append(current_acc)
                trajectory[].append(avg_cot_length)

                
                 current_acc < baseline_acc * :
                    ()

         trajectory

     ():
        
        
        a, b = ,   

        recovery_tokens = a * (model_size ** (-b))

        
        tokens_per_example = 
        num_examples = recovery_tokens / tokens_per_example

         num_examples

Step 2: Implement Reflection Detection

Identify surface-level vs. genuine reasoning:

class ReflectionDetector:
    """Detect whether model is genuinely reflecting or pattern-matching"""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.reflection_keywords = [
            "let me think", "wait", "actually", "i made a mistake",
            "let me reconsider", "hmm", "on second thought"
        ]

    def detect_reflection_patterns(self, cot_text):
        """
        Identify whether CoT contains reflection keywords.

        High keyword density ≠ genuine reasoning; often surface-level pattern adoption.
        """
        lower_text = cot_text.lower()

        keyword_count = sum(1 for kw in self.reflection_keywords if kw in lower_text)
        total_tokens = len(self.tokenizer.encode(cot_text))

        reflection_density = keyword_count / max(total_tokens, 1)

        return {
            'has_reflection_keywords': keyword_count > 0,
            'reflection_density': reflection_density,
            'is_likely_surface_level': reflection_density > 0.05
        }

    def llm_based_reflection_quality(self, cot_text, model, tokenizer):
        """
        Use an evaluator LLM to assess reasoning quality.

        Prompt: "Is this reasoning genuine or surface-level pattern matching?"
        """
        prompt = 

        
        input_ids = tokenizer.encode(prompt, return_tensors=)
        outputs = model.generate(input_ids, max_new_tokens=, return_dict_in_generate=,
                                output_scores=)

        
        generated_token = outputs.sequences[, -]
        score = (tokenizer.decode(generated_token))  tokenizer.decode(generated_token).isdigit()  

         score

Step 3: Implement Adaptive Training Strategy

Create curriculum to navigate degradation valley:

class AdaptiveCoTTrainer:
    """
    Adaptively train on long CoT to minimize degradation valley impact.

    Strategy: Gradually increase CoT length, monitor for degradation, adjust batch composition.
    """

    def __init__(self, model, tokenizer, device='cuda'):
        self.model = model.to(device)
        self.tokenizer = tokenizer
        self.device = device
        self.reflection_detector = ReflectionDetector(tokenizer)
        self.error_analyzer = CoTErrorAnalyzer(model, tokenizer)

    def train_with_curriculum(self, dataset, epochs=3, target_cot_length=256):
        """
        Curriculum learning: start with short CoTs, gradually increase.

        Avoids hitting degradation valley too suddenly.
        """
        # Sort dataset by CoT length
        dataset_by_length = self.bucket_dataset_by_cot_length(dataset)

        current_max_length = 64  # Start with short reasoning chains

        for epoch in range(epochs):
            print(f"Epoch {epoch+1}, max CoT length: {current_max_length}")

            # Select training examples up to current length limit
            train_bucket = dataset_by_length[current_max_length]

            # Train for an epoch
            self.train_epoch(train_bucket)

            # Evaluate and check for degradation
            eval_acc = self.evaluate_on_all_lengths(dataset_by_length)

            print()

            
             current_max_length < target_cot_length:
                current_max_length = (current_max_length * , target_cot_length)

     ():
        
        buckets = defaultdict()

         sample  dataset:
            cot_length = (.tokenizer.encode(sample[]))

            
            bucket_idx =  ** (cot_length.bit_length() - )
            bucket_idx = (bucket_idx, )  

            buckets[bucket_idx].append(sample)

         buckets

     ():
        
        optimizer = torch.optim.AdamW(.model.parameters(), lr=)

        baseline_loss = 

         step, sample  (training_samples):
            prompt = sample[]
            cot = sample[]
            answer = sample[]

            
            full_text = 
            input_ids = .tokenizer.encode(full_text, return_tensors=).to(.device)

            optimizer.zero_grad()

            outputs = .model(input_ids=input_ids, labels=input_ids)
            loss = outputs.loss

            
             baseline_loss  :
                baseline_loss = loss.item()

            
             loss.item() > baseline_loss * :
                ()

            loss.backward()
            torch.nn.utils.clip_grad_norm_(.model.parameters(), )
            optimizer.step()

     ():
        
        results = {}

        .model.()

         length, samples  buckets.items():
              samples:
                

            correct = 

             torch.no_grad():
                 sample  samples:
                    prompt = sample[]
                    target_answer = sample[]

                    generated = .model.generate(
                        .tokenizer.encode(prompt, return_tensors=).to(.device),
                        max_new_tokens=,
                        do_sample=
                    )

                    generated_text = .tokenizer.decode(generated[])

                    
                     target_answer  generated_text:
                        correct += 

            accuracy = correct / (samples)
            results[length] = accuracy

         results

Step 4: Downstream RL Integration

Prepare models for subsequent reinforcement learning:

class CoTPretrainingForRL:
    """
    Evaluate CoT pre-training as foundation for downstream RL.

    RL training sensitivity to pre-training trajectory is significant.
    """

    def __init__(self, pretrained_model, tokenizer):
        self.model = pretrained_model
        self.tokenizer = tokenizer

    def assess_rl_readiness(self, model_checkpoint, rl_benchmark):
        """
        Measure how well CoT pre-training prepared model for RL fine-tuning.

        Key metric: How quickly does RL training converge and what's final performance?
        """
        # Load checkpoint
        model = AutoModelForCausalLM.from_pretrained(model_checkpoint)

        # Run brief RL training
        rl_trainer = RLTrainer(model, self.tokenizer)
        convergence_curve = rl_trainer.train(rl_benchmark, num_steps=1000)

        # Metrics
        convergence_speed = self.estimate_convergence_speed(convergence_curve)
        final_performance = convergence_curve['rewards'][-1]
        stability = self.compute_training_stability(convergence_curve)

        readiness_score = (convergence_speed + final_performance + stability) / 3

        return {
            'convergence_speed': convergence_speed,
            'final_performance': final_performance,
            'stability': stability,
            'readiness_score': readiness_score
        }

    def estimate_convergence_speed(self, convergence_curve):
        """How quickly does performance improve during RL?"""
        rewards = np.array(convergence_curve['rewards'])

        # Fit exponential: reward(t) = asymptote * (1 - exp(-decay * t))
         scipy.optimize  curve_fit

         ():
             asymptote * ( - np.exp(-decay * t))

        :
            popt, _ = curve_fit(exponential, np.arange((rewards)), rewards,
                              p0=[rewards[-], ], maxfev=)
            decay_rate = popt[]
        :
            decay_rate = 

        
         (decay_rate / , )

     ():
        
        rewards = np.array(convergence_curve[])

        
        window = (, (rewards) // )
        rolling_std = pd.Series(rewards).rolling(window).std()

        
        avg_std = rolling_std.mean()
        stability =  / ( + avg_std)

         stability

Practical Guidance

• Model Size Matters: <1B models are at highest risk; 7B+ models recover reliably
• Degradation Prevention: Start with short CoTs (64-128 tokens), gradually increase
• Valley Depth: Smaller models experience 20-40% performance drop during valley
• Recovery Timeline: 50-220k examples needed depending on model size
• Reflection Quality: Monitor keyword density; >5% suggests surface-level patterns
• Batch Composition: Mix short and long CoTs to maintain stability
• RL Sensitivity: Models that degrade significantly may struggle in subsequent RL
• Best Practices: Curriculum learning, error monitoring, targeted data selection

Reference

• Error accumulation is exponential with sequence length: accuracy = (1 - error_rate)^length
• Recovery dynamics follow power law: recovery_tokens ∝ model_size^(-0.5)
• Surface-level reflection adoption is common; requires evaluator-based verification
• CoT pre-training trajectory significantly impacts downstream RL fine-tuning performance