Controlling LLM Reasoning Costs via Self-Reflection

# Illustrative sketch (not from the paper)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load a frozen backbone (e.g., deepseek-v4-flash) – no gradient updates
model_name = "deepseek-v4-flash"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
model.eval()
 tokenizer = AutoTokenizer.from_pretrained(model_name)

D = 5  # maximum reflection depth

def reflective_loop(question, max_depth=D):
    answer = None
    for step in range(1, max_depth + 1):
        # 1️⃣ generate answer (or revision) conditioned on previous critique
        prompt = f"Question: {question}\n"
        if answer is not None:
            prompt += f"Previous answer: {answer}\n"
        prompt += "Provide a concise answer."
        inputs = tokenizer(prompt, return_tensors="pt")
        gen = model.generate(**inputs, max_new_tokens=64)
        answer = tokenizer.decode(gen[0], skip_special_tokens=True).split("Answer:")[-1].strip()

        # 2️⃣ self‑critique using meta‑reward prompts (correctness, efficiency, depth, tool diversity)
        critique_prompt = (
            f"Answer: {answer}\n"
            "Critique the answer for correctness, efficiency, reflection depth, and tool‑call diversity. "
            "If the answer is satisfactory, reply with the sentinel 'CONFIRMED'."
        )
        crit_inputs = tokenizer(critique_prompt, return_tensors="pt")
        crit_gen = model.generate(**crit_inputs, max_new_tokens=32)
        critique = tokenizer.decode(crit_gen[0], skip_special_tokens=True).strip()

        # 3️⃣ early stopping if sentinel appears
        if "CONFIRMED" in critique:
            break
        # otherwise the loop continues and the next iteration will revise the answer
    return answer, critique, step

# Example usage
q = "If a train travels 60 km/h for 2 hours, how far does it go?"
final_answer, final_critique, used_steps = reflective_loop(q)
print(f"Answer after {used_steps} step(s): {final_answer}\nCritique: {final_critique}")

// Illustrative sketch (not from the paper)
const { AutoModelForCausalLM, AutoTokenizer } = require('@huggingface/transformers');
const torch = require('torch-js'); // placeholder for torch bindings

// Load a frozen backbone (no training)
const modelName = 'deepseek-v4-flash';
let model, tokenizer;
(async () => {
  model = await AutoModelForCausalLM.from_pretrained(modelName, { dtype: 'float16' });
  tokenizer = await AutoTokenizer.from_pretrained(modelName);
  model.eval();
})();

const D = 5; // maximum reflection depth

async function reflectiveLoop(question, maxDepth = D) {
  let answer = null;
  let critique = '';
  let step = 0;
  for (step = 1; step <= maxDepth; step++) {
    // 1️⃣ generate answer (or revision)
    let prompt = `Question: ${question}\n`;
    if (answer) prompt += `Previous answer: ${answer}\n`;
    prompt += 'Provide a concise answer.';
    const inputs = tokenizer(prompt, { return_tensors: 'pt' });
    const gen = await model.generate({ ...inputs, max_new_tokens: 64 });
    answer = tokenizer.decode(gen[0], { skip_special_tokens: true })
      .split('Answer:')
      .pop()
      .trim();

    // 2️⃣ self‑critique with meta‑reward prompts
    const critiquePrompt = `Answer: ${answer}\nCritique the answer for correctness, efficiency, reflection depth, and tool‑call diversity. If the answer is satisfactory, reply with the sentinel 'CONFIRMED'.`;
    const critInputs = tokenizer(critiquePrompt, { return_tensors: 'pt' });
    const critGen = await model.generate({ ...critInputs, max_new_tokens: 32 });
    critique = tokenizer.decode(critGen[0], { skip_special_tokens: true }).trim();

    // 3️⃣ early stopping if sentinel appears
    if (critique.includes('CONFIRMED')) break;
    // otherwise loop continues to revise answer
  }
  return { answer, critique, steps: step };
}

// Example usage
(async () => {
  const q = 'If a train travels 60 km/h for 2 hours, how far does it go?';
  const result = await reflectiveLoop(q);
  console.log(`Answer after ${result.steps} step(s): ${result.answer}`);
  console.log(`Critique: ${result.critique}`);
})();