Test-Time Compute Deep Dive: Engineering Practices for Making Models Think Longer

# Traditional paradigm: better performance = bigger model + more training data
performance = f(model_size, training_compute)

# TTC paradigm: better performance = more compute at inference
performance = f(model_size, training_compute, inference_compute)

import openai

def solve_with_cot(problem: str, client: openai.OpenAI) -> str:
    """Solve complex problems using Chain-of-Thought prompting"""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a rigorous reasoning expert. When solving problems:\n"
                    "1. Explicitly list known conditions\n"
                    "2. Derive step by step, stating the basis for each step\n"
                    "3. Verify the final answer for reasonableness"
                )
            },
            {
                "role": "user",
                "content": f"Please solve the following problem step by step:\n\n{problem}"
            }
        ],
        temperature=0.0
    )
    return response.choices[0].message.content

import asyncio
from collections import Counter
from typing import List

async def self_consistency_solve(
    problem: str,
    client: openai.AsyncOpenAI,
    num_samples: int = 5,
    temperature: float = 0.7
) -> dict:
    """Self-Consistency: multi-path sampling + majority voting"""
    
    async def sample_one() -> str:
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "Reason step by step. Wrap final answer in \\boxed{}."},
                {"role": "user", "content": problem}
            ],
            temperature=temperature
        )
        return response.choices[0].message.content
    
    # Sample N independent reasoning paths in parallel
    paths = await asyncio.gather(*[sample_one() for _ in range(num_samples)])
    
    # Extract final answers and vote
    answers = [extract_boxed_answer(path) for path in paths]
    vote_counts = Counter(answers)
    best_answer = vote_counts.most_common(1)[0][0]
    confidence = vote_counts[best_answer] / num_samples
    
    return {
        "answer": best_answer,
        "confidence": confidence,
        "num_paths": num_samples,
        "vote_distribution": dict(vote_counts)
    }


def extract_boxed_answer(text: str) -> str:
    """Extract answer from LaTeX \\boxed{} format"""
    import re
    match = re.search(r'\\boxed\{(.+?)\}', text)
    return match.group(1).strip() if match else text.strip().split('\n')[-1]

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ThoughtNode:
    """Reasoning tree node"""
    thought: str
    score: float = 0.0
    children: list = field(default_factory=list)
    parent: Optional['ThoughtNode'] = None
    depth: int = 0

class TreeOfThought:
    """Tree-of-Thought reasoning search engine"""
    
    def __init__(self, client: openai.OpenAI, max_depth: int = 3, branch_factor: int = 3):
        self.client = client
        self.max_depth = max_depth
        self.branch_factor = branch_factor
    
    def solve(self, problem: str) -> ThoughtNode:
        root = ThoughtNode(thought=f"Problem: {problem}")
        self._expand(root, problem)
        return self._best_leaf(root)
    
    def _expand(self, node: ThoughtNode, problem: str):
        """Recursively expand the reasoning tree"""
        if node.depth >= self.max_depth:
            return
        
        candidates = self._generate_thoughts(problem, node)
        
        for thought_text in candidates:
            child = ThoughtNode(
                thought=thought_text,
                parent=node,
                depth=node.depth + 1
            )
            child.score = self._evaluate_thought(problem, child)
            node.children.append(child)
        
        # Only expand top-scoring branches (beam search)
        node.children.sort(key=lambda x: x.score, reverse=True)
        for child in node.children[:2]:  # beam width = 2
            self._expand(child, problem)
    
    def _generate_thoughts(self, problem: str, node: ThoughtNode) -> list:
        """Generate candidate next-step thoughts for the current node"""
        path = self._get_path(node)
        prompt = (
            f"Problem: {problem}\n\n"
            f"Reasoning so far:\n{path}\n\n"
            f"Propose {self.branch_factor} different next reasoning directions. "
            f"Label each as [Thought N]."
        )
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.8
        )
        return self._parse_thoughts(response.choices[0].message.content)
    
    def _evaluate_thought(self, problem: str, node: ThoughtNode) -> float:
        """Use LLM as evaluator to score (0-1)"""
        path = self._get_path(node)
        prompt = (
            f"Rate the quality of this reasoning path (0-1):\n"
            f"Problem: {problem}\nReasoning: {path}\n"
            f"Criteria: logical coherence, correct direction, no obvious errors.\n"
            f"Output only a number between 0 and 1."
        )
        response = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0
        )
        try:
            return float(response.choices[0].message.content.strip())
        except ValueError:
            return 0.5
    
    def _get_path(self, node: ThoughtNode) -> str:
        """Backtrack to get full path from root to current node"""
        path = []
        current = node
        while current.parent:
            path.append(current.thought)
            current = current.parent
        return "\n→ ".join(reversed(path)) if path else "(start)"
    
    def _best_leaf(self, root: ThoughtNode) -> ThoughtNode:
        """Find the highest-scoring leaf node"""
        best = root
        stack = [root]
        while stack:
            node = stack.pop()
            if not node.children and node.score > best.score:
                best = node
            stack.extend(node.children)
        return best
    
    def _parse_thoughts(self, text: str) -> list:
        import re
        thoughts = re.findall(r'\[Thought\s*\d+\]\s*(.+?)(?=\[Thought|$)', text, re.DOTALL)
        return [t.strip() for t in thoughts] if thoughts else [text]

# DeepSeek R1's training philosophy (pseudocode)
class DeepSeekR1Training:
    """
    R1's core innovations:
    1. No expensive Reward Model (unlike RLHF)
    2. Uses GRPO (Group Relative Policy Optimization)
    3. Reasoning process fully visible (<think>...</think> tags)
    """
    
    def grpo_step(self, problem, model):
        # Sample a group of candidate answers
        group = [model.generate(problem) for _ in range(16)]
        
        # Score using rule-based verifiers (not a Reward Model)
        scores = [self.rule_verifier(problem, answer) for answer in group]
        
        # Use within-group relative ranking as reward signal
        baseline = sum(scores) / len(scores)
        advantages = [s - baseline for s in scores]
        
        # Policy gradient update
        model.update(group, advantages)
    
    def rule_verifier(self, problem, answer):
        """Rule-based verifier (math: compare to ground truth; code: run tests)"""
        if problem.type == "math":
            return 1.0 if answer.final == problem.ground_truth else 0.0
        elif problem.type == "code":
            return run_tests(answer.code, problem.test_cases)

import OpenAI from 'openai';

interface RefinementConfig {
  maxIterations: number;
  qualityThreshold: number;
  model: string;
  critiqueModel: string;
}

interface RefinementResult {
  finalOutput: string;
  iterations: number;
  scores: number[];
  totalTokens: number;
}

async function iterativeRefinement(
  task: string,
  config: RefinementConfig,
  client: OpenAI
): Promise<RefinementResult> {
  const { maxIterations, qualityThreshold, model, critiqueModel } = config;
  let currentOutput = '';
  const scores: number[] = [];
  let totalTokens = 0;

  // Initial generation
  const initial = await client.chat.completions.create({
    model,
    messages: [
      { role: 'system', content: 'You are an expert problem solver.' },
      { role: 'user', content: task }
    ],
    temperature: 0.7
  });
  currentOutput = initial.choices[0].message.content || '';
  totalTokens += initial.usage?.total_tokens || 0;

  for (let i = 0; i < maxIterations; i++) {
    // Evaluate current output quality
    const critique = await client.chat.completions.create({
      model: critiqueModel,
      messages: [
        {
          role: 'system',
          content: `Evaluate the output quality. Return JSON:
            {"score": 0.0-1.0, "issues": ["issue1", ...], "suggestions": ["suggestion1", ...]}`
        },
        {
          role: 'user',
          content: `Task: ${task}\n\nOutput: ${currentOutput}`
        }
      ],
      temperature: 0.0,
      response_format: { type: 'json_object' }
    });
    totalTokens += critique.usage?.total_tokens || 0;

    const evaluation = JSON.parse(critique.choices[0].message.content || '{}');
    scores.push(evaluation.score || 0);

    // Exit early if quality threshold is met
    if (evaluation.score >= qualityThreshold) {
      return { finalOutput: currentOutput, iterations: i + 1, scores, totalTokens };
    }

    // Refine based on feedback
    const refinement = await client.chat.completions.create({
      model,
      messages: [
        {
          role: 'system',
          content: 'Improve your output based on the following feedback. Keep what works, fix the issues.'
        },
        { role: 'user', content: `Original task: ${task}` },
        { role: 'assistant', content: currentOutput },
        {
          role: 'user',
          content: `Feedback:\nIssues: ${evaluation.issues?.join(', ')}\nSuggestions: ${evaluation.suggestions?.join(', ')}\n\nPlease output the improved complete version.`
        }
      ],
      temperature: 0.5
    });
    totalTokens += refinement.usage?.total_tokens || 0;
    currentOutput = refinement.choices[0].message.content || currentOutput;
  }

  return { finalOutput: currentOutput, iterations: maxIterations, scores, totalTokens };
}

// Usage example
const result = await iterativeRefinement(
  'Implement a thread-safe LRU Cache in Python with TTL expiration',
  {
    maxIterations: 3,
    qualityThreshold: 0.85,
    model: 'gpt-4o',
    critiqueModel: 'gpt-4o-mini'
  },
  new OpenAI()
);

console.log(`Iterations: ${result.iterations}, Final score: ${result.scores.at(-1)}`);

import math
import random
from dataclasses import dataclass, field
from typing import List, Optional

@dataclass
class MCTSNode:
    """MCTS reasoning node"""
    state: str
    parent: Optional['MCTSNode'] = None
    children: List['MCTSNode'] = field(default_factory=list)
    visits: int = 0
    value: float = 0.0
    reasoning_step: str = ""
    
    @property
    def ucb1(self) -> float:
        if self.visits == 0:
            return float('inf')
        exploitation = self.value / self.visits
        exploration = math.sqrt(2 * math.log(self.parent.visits) / self.visits)
        return exploitation + exploration


class ReasoningMCTS:
    """MCTS-based reasoning search engine"""
    
    def __init__(
        self,
        client: openai.OpenAI,
        num_simulations: int = 50,
        max_depth: int = 5,
        expansion_width: int = 3
    ):
        self.client = client
        self.num_simulations = num_simulations
        self.max_depth = max_depth
        self.expansion_width = expansion_width
    
    def search(self, problem: str) -> str:
        """Run MCTS search, return optimal reasoning path"""
        root = MCTSNode(state=problem)
        
        for _ in range(self.num_simulations):
            # 1. Selection: follow UCB1 policy
            leaf = self._select(root)
            
            # 2. Expansion: generate new reasoning steps
            if leaf.visits > 0 and leaf.children == []:
                self._expand(leaf, problem)
                if leaf.children:
                    leaf = random.choice(leaf.children)
            
            # 3. Simulation: quick evaluation
            value = self._simulate(leaf, problem)
            
            # 4. Backpropagation: update ancestors
            self._backpropagate(leaf, value)
        
        return self._extract_best_path(root)
    
    def _select(self, node: MCTSNode) -> MCTSNode:
        """UCB1 selection policy"""
        current = node
        while current.children:
            current = max(current.children, key=lambda c: c.ucb1)
        return current
    
    def _expand(self, node: MCTSNode, problem: str):
        """Expand node: generate multiple candidate next steps"""
        path = self._reconstruct_path(node)
        
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[{
                "role": "user",
                "content": (
                    f"Problem: {problem}\n"
                    f"Current reasoning: {path}\n\n"
                    f"Generate {self.expansion_width} different next reasoning directions. "
                    f"Separate each with ---."
                )
            }],
            temperature=0.9
        )
        
        steps = response.choices[0].message.content.split('---')
        for step in steps[:self.expansion_width]:
            child = MCTSNode(
                state=node.state + "\n" + step.strip(),
                parent=node,
                reasoning_step=step.strip()
            )
            node.children.append(child)
    
    def _simulate(self, node: MCTSNode, problem: str) -> float:
        """Use LLM for quick terminal value estimation"""
        path = self._reconstruct_path(node)
        
        response = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": (
                    f"Rate the likelihood this reasoning leads to a correct answer (0-1):\n"
                    f"Problem: {problem}\n"
                    f"Current reasoning: {path}\n"
                    f"Output only a number."
                )
            }],
            temperature=0.0
        )
        
        try:
            return float(response.choices[0].message.content.strip())
        except ValueError:
            return 0.5
    
    def _backpropagate(self, node: MCTSNode, value: float):
        """Backpropagate value updates"""
        current = node
        while current:
            current.visits += 1
            current.value += value
            current = current.parent
    
    def _reconstruct_path(self, node: MCTSNode) -> str:
        """Reconstruct reasoning path from root to current node"""
        steps = []
        current = node
        while current.parent:
            steps.append(current.reasoning_step)
            current = current.parent
        return " → ".join(reversed(steps)) if steps else "(start)"
    
    def _extract_best_path(self, root: MCTSNode) -> str:
        """Extract the most-visited path (most reliable)"""
        path_steps = []
        current = root
        while current.children:
            current = max(current.children, key=lambda c: c.visits)
            path_steps.append(current.reasoning_step)
        return "\n".join(path_steps)

import OpenAI from 'openai';

// OpenAI o1 series: control reasoning budget via max_completion_tokens
async function o1ReasoningWithBudget(
  problem: string,
  thinkingBudget: 'low' | 'medium' | 'high',
  client: OpenAI
) {
  const budgetMap = {
    low: 4096,     // Quick answer, minimal reasoning
    medium: 16384, // Balanced mode
    high: 65536    // Deep reasoning, cost no object
  };

  const response = await client.chat.completions.create({
    model: 'o1',
    messages: [{ role: 'user', content: problem }],
    max_completion_tokens: budgetMap[thinkingBudget]
  });

  return {
    answer: response.choices[0].message.content,
    reasoningTokens: response.usage?.completion_tokens_details?.reasoning_tokens,
    outputTokens: response.usage?.completion_tokens_details?.accepted_prediction_tokens,
    totalCost: calculateCost(response.usage)
  };
}

// DeepSeek R1: visible reasoning process
async function deepseekR1Reasoning(problem: string) {
  const client = new OpenAI({
    baseURL: 'https://api.deepseek.com/v1',
    apiKey: process.env.DEEPSEEK_API_KEY
  });

  const response = await client.chat.completions.create({
    model: 'deepseek-reasoner',
    messages: [{ role: 'user', content: problem }]
  });

  // DeepSeek R1 returns reasoning_content (thinking) and content (final answer)
  const message = response.choices[0].message as any;
  return {
    thinking: message.reasoning_content,  // Full thinking process
    answer: message.content               // Final answer
  };
}

function calculateCost(usage: any): number {
  const reasoningCost = (usage?.completion_tokens_details?.reasoning_tokens || 0) * 0.015 / 1000;
  const outputCost = (usage?.completion_tokens || 0) * 0.06 / 1000;
  const inputCost = (usage?.prompt_tokens || 0) * 0.015 / 1000;
  return reasoningCost + outputCost + inputCost;
}

import subprocess
import tempfile
from typing import Tuple

class CodeGenerationWithVerification:
    """Code generation with verification loops (canonical TTC for code)"""
    
    def __init__(self, client: openai.OpenAI, max_attempts: int = 5):
        self.client = client
        self.max_attempts = max_attempts
    
    def generate_and_verify(
        self, 
        task: str, 
        test_code: str
    ) -> Tuple[str, int]:
        """Generate code and verify with tests, iterate on failure"""
        
        code = ""
        for attempt in range(self.max_attempts):
            if attempt == 0:
                code = self._generate_initial(task)
            else:
                code = self._fix_code(task, code, error_output)
            
            success, error_output = self._run_tests(code, test_code)
            
            if success:
                return code, attempt + 1
        
        return code, self.max_attempts
    
    def _generate_initial(self, task: str) -> str:
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "Write Python code. Output only code, no explanations."},
                {"role": "user", "content": task}
            ],
            temperature=0.3
        )
        return self._extract_code(response.choices[0].message.content)
    
    def _fix_code(self, task: str, code: str, error: str) -> str:
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "Fix the code error. Output only the corrected complete code."},
                {"role": "user", "content": f"Task: {task}\n\nCode:\n{code}\n\nError:\n{error}"}
            ],
            temperature=0.2
        )
        return self._extract_code(response.choices[0].message.content)
    
    def _run_tests(self, code: str, test_code: str) -> Tuple[bool, str]:
        """Run code and tests in a sandbox"""
        full_code = f"{code}\n\n{test_code}"
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
            f.write(full_code)
            f.flush()
            result = subprocess.run(
                ['python', f.name],
                capture_output=True, text=True, timeout=10
            )
        
        if result.returncode == 0:
            return True, ""
        return False, result.stderr
    
    def _extract_code(self, text: str) -> str:
        if '```python' in text:
            return text.split('```python')[1].split('```')[0].strip()
        return text.strip()

from enum import Enum

class DifficultyLevel(Enum):
    TRIVIAL = "trivial"
    EASY = "easy"
    MEDIUM = "medium"
    HARD = "hard"
    EXPERT = "expert"

class AdaptiveComputeAllocator:
    """Adaptive inference compute allocator"""
    
    STRATEGIES = {
        DifficultyLevel.TRIVIAL: {
            "method": "direct",
            "samples": 1,
            "max_tokens": 256,
            "model": "gpt-4o-mini"
        },
        DifficultyLevel.EASY: {
            "method": "cot",
            "samples": 1,
            "max_tokens": 1024,
            "model": "gpt-4o-mini"
        },
        DifficultyLevel.MEDIUM: {
            "method": "self_consistency",
            "samples": 3,
            "max_tokens": 2048,
            "model": "gpt-4o"
        },
        DifficultyLevel.HARD: {
            "method": "self_consistency",
            "samples": 7,
            "max_tokens": 4096,
            "model": "gpt-4o"
        },
        DifficultyLevel.EXPERT: {
            "method": "mcts",
            "simulations": 30,
            "max_tokens": 8192,
            "model": "o1"
        }
    }
    
    def __init__(self, client: openai.OpenAI):
        self.client = client
    
    async def solve(self, problem: str) -> dict:
        difficulty = await self._classify_difficulty(problem)
        strategy = self.STRATEGIES[difficulty]
        
        if strategy["method"] == "direct":
            return await self._direct_solve(problem, strategy)
        elif strategy["method"] == "cot":
            return await self._cot_solve(problem, strategy)
        elif strategy["method"] == "self_consistency":
            return await self._sc_solve(problem, strategy)
        elif strategy["method"] == "mcts":
            return await self._mcts_solve(problem, strategy)
    
    async def _classify_difficulty(self, problem: str) -> DifficultyLevel:
        """Use a small model to quickly classify problem difficulty"""
        response = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": (
                    f"Classify this problem as trivial/easy/medium/hard/expert:\n"
                    f"{problem}\nOutput only one word."
                )
            }],
            temperature=0.0,
            max_tokens=10
        )
        level_str = response.choices[0].message.content.strip().lower()
        return DifficultyLevel(level_str) if level_str in [e.value for e in DifficultyLevel] else DifficultyLevel.MEDIUM

interface CascadeConfig {
  stages: Array<{
    model: string;
    maxTokens: number;
    confidenceThreshold: number;
  }>;
}

async function cascadeReasoning(
  problem: string,
  config: CascadeConfig,
  client: OpenAI
): Promise<{ answer: string; stage: number; totalCost: number }> {
  let totalCost = 0;

  for (let i = 0; i < config.stages.length; i++) {
    const stage = config.stages[i];
    
    const response = await client.chat.completions.create({
      model: stage.model,
      messages: [
        { role: 'system', content: 'Solve the problem and assess your confidence (0-1). Return JSON: {"answer": "...", "confidence": 0.X}' },
        { role: 'user', content: problem }
      ],
      max_tokens: stage.maxTokens,
      response_format: { type: 'json_object' }
    });

    totalCost += estimateCost(response.usage, stage.model);
    const result = JSON.parse(response.choices[0].message.content || '{}');

    if (result.confidence >= stage.confidenceThreshold) {
      return { answer: result.answer, stage: i + 1, totalCost };
    }
  }

  return { answer: 'Fallback to last stage', stage: config.stages.length, totalCost };
}

// Simple problems resolve at stage 1; complex ones escalate
const cascade = await cascadeReasoning(problem, {
  stages: [
    { model: 'gpt-4o-mini', maxTokens: 512, confidenceThreshold: 0.9 },
    { model: 'gpt-4o', maxTokens: 2048, confidenceThreshold: 0.8 },
    { model: 'o1', maxTokens: 16384, confidenceThreshold: 0.0 }
  ]
}, client);

async def early_stopping_consistency(
    problem: str,
    client: openai.AsyncOpenAI,
    max_samples: int = 10,
    consensus_threshold: int = 3
) -> dict:
    """Self-Consistency with early stopping: stop after N consecutive same answers"""
    
    answers = []
    consecutive_same = 0
    last_answer = None
    
    for i in range(max_samples):
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "Reason step by step. Wrap final answer in \\boxed{}."},
                {"role": "user", "content": problem}
            ],
            temperature=0.7
        )
        
        answer = extract_boxed_answer(response.choices[0].message.content)
        answers.append(answer)
        
        # Early stop: consensus_threshold consecutive identical answers
        if answer == last_answer:
            consecutive_same += 1
            if consecutive_same >= consensus_threshold:
                break
        else:
            consecutive_same = 1
            last_answer = answer
    
    vote_counts = Counter(answers)
    best = vote_counts.most_common(1)[0]
    
    return {
        "answer": best[0],
        "confidence": best[1] / len(answers),
        "samples_used": len(answers),
        "samples_saved": max_samples - len(answers),
        "early_stopped": len(answers) < max_samples
    }