Most of our experience with Large Language Models (LLMs) involves a frozen brain. You send a prompt, and the model — with weights fixed in time — predicts the next token. If the model fails at a complex coding task or a mathematical proof, we simply prompt it again or try “Best-of-N” sampling, hoping it “gets lucky.”

The release of TTT-Discover (from the paper “Learning to Discover at Test Time” https://arxiv.org/pdf/2601.16175) introduces a radical alternative: What if the model could learn from its own peak in real-time to fine-tune the model for a single problem?

The Problem: The “Frozen Brain” Plateau

Standard inference techniques like Best-of-N are “democratic.” They treat every attempt as an independent event. If you ask an LLM to optimize a GPU kernel, it might give you 100 versions. 99 might be slow or broken, and one might be a breakthrough.

In a standard system, the model never realizes it found a miracle. It doesn’t learn from that success to make its next 100 attempts even better. It stays stagnant.

The Solution: TTT-Discover & The J-Beta Objective

TTT-Discover turns the LLM into a Test-Time Optimizer. Instead of just sampling, the model enters a loop of acting, being evaluated by an external verifier (like a compiler), and then — crucially — updating its own weights before the next attempt.

The heart of this discovery is the J_beta Entropic Utility objective.

Standard RL vs. J-Beta

In standard Reinforcement Learning (REINFORCE), we maximize the average reward. This is great for making a chatbot “general purpose.”

But for discovery, the average is noise. You don’t want an “average” GPU kernel; you want the one that breaks the world record.

By exponentially weighting the rewards, the model focuses its entire gradient update on the outliers — the rare, high-performing solutions. It ignores the 99 failures and “specializes” its weights toward the logic of the one success.

The Implementation: End-to-End Discovery with LoRA

We use LoRA (Low-Rank Adaptation) to update only a small set of “specialist” weights during the test session. Below is a functional Python skeleton of the TTT-Discover loop:

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
import numpy as np

# 1. THE EVALUATOR (The "Environment")
def evaluate_kernel_performance(code_string):
    """
    Placeholder: In a real blog, this would be a subprocess 
    calling 'triton' or 'g++' to benchmark the code.
    """
    # Simulate a reward based on code length or a random discovery
    reward = len(code_string) / 1000.0 + np.random.uniform(0, 0.5)
    return reward

# 2. LOG-PROBABILITY EXTRACTOR
def get_sequence_log_probs(model, tokenizer, prompt, generations):
    """Calculates log P(generation | prompt) for the policy gradient."""
    # Concatenate prompt + generation
    texts = [prompt + gen for gen in generations]
    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True).to(model.device)
    
    # Forward pass
    outputs = model(**inputs)
    logits = outputs.logits # [B, Seq, Vocab]
    
    # Identify where the generation starts
    prompt_len = tokenizer(prompt, return_tensors="pt")['input_ids'].shape[1]
    
    # Shift to align logits with target tokens
    shift_logits = logits[:, prompt_len-1:-1, :].contiguous()
    shift_labels = inputs.input_ids[:, prompt_len:].contiguous()
    
    # Gather log_probs of the actual tokens sampled
    log_probs = F.log_softmax(shift_logits, dim=-1)
    token_log_probs = torch.gather(log_probs, 2, shift_labels.unsqueeze(-1)).squeeze(-1)
    
    # Mask out padding tokens
    mask = (shift_labels != tokenizer.pad_token_id).float()
    return (token_log_probs * mask).sum(dim=-1)

# 3. THE TTT-DISCOVER ENGINE
class DiscoveryAgent:
    def __init__(self, model_id="meta-llama/Llama-3-8B"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        
        # Initialize with LoRA
        base_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
        peft_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"])
        self.model = get_peft_model(base_model, peft_config)
        
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)
        self.buffer = [] # To store {'code': str, 'reward': float, 'visits': int}

    def run_discovery(self, problem_prompt, steps=50, batch_size=4):
        print(f"--- Starting TTT-Discover for Problem ---\n")
        
        for i in range(steps):
            # A. Select Seed (PUCT Logic)
            seed = self.select_seed()
            current_prompt = f"{problem_prompt}\nBest current solution:\n{seed['code']}" if seed else problem_prompt
            
            # B. Act (Sample Rollouts)
            inputs = self.tokenizer(current_prompt, return_tensors="pt").to(self.model.device)
            output_tokens = self.model.generate(
                **inputs, max_new_tokens=128, num_return_sequences=batch_size, do_sample=True, temperature=0.8
            )
            
            # Decode only the NEW tokens
            generations = [self.tokenizer.decode(o[inputs.input_ids.shape[1]:], skip_special_tokens=True) for o in output_tokens]
            
            # C. Evaluate (External Reward)
            rewards = torch.tensor([evaluate_kernel_performance(g) for g in generations]).to(self.model.device)
            
            # Save to buffer
            for g, r in zip(generations, rewards):
                self.buffer.append({'code': g, 'reward': r.item(), 'visits': 0})
            
            # D. Test-Time Training (J-Beta Update)
            # Calculate Adaptive Beta based on reward spread
            beta = 1.0 / (rewards.std() + 1e-6)
            weights = F.softmax(beta * rewards, dim=0)
            
            # Get Log-Probs for gradients
            log_probs = get_sequence_log_probs(self.model, self.tokenizer, current_prompt, generations)
            
            # Loss = - weighted log_prob (maximizing reward)
            loss = -(weights.detach() * log_probs).sum()
            
            loss.backward()
            self.optimizer.step()
            self.optimizer.zero_grad()
            
            print(f"Step {i} | Best Reward: {max(r['reward'] for r in self.buffer):.4f} | Beta: {beta:.2f}")

    def select_seed(self):
        if not self.buffer: return None
        # Simplified PUCT: Pick the highest reward solution from buffer
        return max(self.buffer, key=lambda x: x['reward'])

# 4. MAIN EXECUTION
if __name__ == "__main__":
    agent = DiscoveryAgent()
    problem = "Write a high-performance Triton kernel for 2D Depthwise Convolution."
    agent.run_discovery(problem, steps=10)

By allowing models to learn at test time, we enable them moving toward Agentic Autonomy

  • Solve OOD Problems: Solve tasks that weren’t in their training data by experimenting.
  • Surpass Human Baselines: TTT-Discover has already set new state-of-the-art records in GPU kernel optimization and mathematical problems (like Erdős’ minimum overlap).