# Long Context Rl Reasoning > Train LLMs for long-context reasoning using KeyChain synthesis: convert short multi-hop QA into long-context tasks by embedding UUID chains in distractor documents, enabling 16K→128K generalization. - Skill: `adu2021/long-context-rl-reasoning` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/long-context-rl-reasoning` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/long-context-rl-reasoning/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: AI & ML - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/long-context-rl-reasoning --- # Technique: KeyChain Synthesis for Long-Context RL Training Standard RL training for reasoning uses short-context tasks (e.g., 2-4 hops). Scaling to long contexts requires massive data, but LoongRL introduces **KeyChain synthesis**: automatically convert short multi-hop tasks into challenging long-context tasks by embedding a chain of UUIDs that must be traced through hundreds of irrelevant documents. The method induces models to learn planning, retrieval, reasoning, and verification patterns without requiring expensive long-context annotation. Models trained at 16K context length effectively handle 128K-length tasks, and smaller models (7B) achieve performance comparable to frontier models (o3-mini, DeepSeek-R1). ## Core Concept LoongRL operates on three principles: - **KeyChain Synthesis**: Generate long-context tasks by embedding UUID chains in document collections - **Emergent Reasoning Patterns**: Training induces planning → retrieval → reasoning → verification steps - **Efficient Data Generation**: No manual long-context annotations needed; synthetic tasks are free to generate - **Generalization**: Models trained on 16K context handle 8× longer contexts without fine-tuning The insight is that the structure of the task (finding a chain through noise) naturally encourages the reasoning patterns needed for long-context understanding. ## Architecture Overview - **Task Synthesizer**: Convert short QA tasks into long-context KeyChain variants - **Document Generator**: Create distractor documents with planted UUID chain - **RL Trainer**: Standard RL loop with synthetic long-context tasks - **Verifier**: Check if model correctly traced chain and answered question - **Curriculum Scheduler**: Gradually increase context length during training ## Implementation Steps The core innovation is KeyChain synthesis. This example shows how to generate tasks and train on them. ```python from typing import List, Dict, Tuple import random import uuid class KeyChainTask: """Represents a KeyChain long-context reasoning task.""" def __init__( self, original_question: str, answer: str, num_documents: int = 100, context_length: int = 16000 ): self.question = original_question self.answer = answer self.num_docs = num_documents self.context_length = context_length # Generate the UUID chain (secret path) self.chain = [str(uuid.uuid4())[:8] for _ in range(5)] # 5-step chain self.documents = [] self.task_prompt = "" def generate_documents_with_chain(self, tokenizer) -> str: """ Create document collection with hidden UUID chain. Chain must be traced step-by-step to find the answer. """ documents = [] # Step 1: Create documents containing sequential UUIDs for step_idx, uuid_id in enumerate(self.chain): if step_idx == 0: # First document contains question and first UUID doc = f""" Document: Research Paper Abstract Title: Important Study on Machine Learning Question embedded here: {self.question} Reference ID: {uuid_id} Next step: Search for documents referencing {self.chain[step_idx + 1] if step_idx + 1 < len(self.chain) else "final answer"} Content: Lorem ipsum dolor sit amet... """ elif step_idx == len(self.chain) - 1: # Last document contains answer doc = f""" Document: Conclusion Summary Previous Reference: {self.chain[step_idx - 1]} Current ID: {uuid_id} ANSWER TO THE QUESTION: {self.answer} This document provides the final answer after following the chain. """ else: # Intermediate documents link chain doc = f""" Document: Supporting Material Previous Step: {self.chain[step_idx - 1]} Current Reference: {uuid_id} Next Reference: {self.chain[step_idx + 1]} This document connects the reasoning chain. """ documents.append(doc) # Add distractors (irrelevant documents) num_distractors = self.num_docs - len(self.chain) for _ in range(num_distractors): distractor = f""" Document: Unrelated Content ID: {uuid.uuid4()} This is a completely unrelated document that should be skipped. Lorem ipsum dolor sit amet, consectetur adipiscing elit... """ documents.append(distractor) # Shuffle documents random.shuffle(documents) # Combine into context context = "\n\n---\n\n".join(documents) # Truncate/pad to desired context length tokens = tokenizer.encode(context) if len(tokens) > self.context_length: tokens = tokens[:self.context_length] elif len(tokens) < self.context_length: # Pad with irrelevant content padding = tokenizer.encode("Extra information: " * 100) tokens = tokens + padding[:self.context_length - len(tokens)] self.documents = documents self.task_prompt = f""" You have {len(documents)} documents. Your task is to: 1. Find the UUID chain: {self.chain[0]} -> {self.chain[1]} -> ... -> {self.chain[-1]} 2. Trace through documents using the references 3. Answer the question: {self.question} Documents: {context[:self.context_length]} Answer: """ return self.task_prompt class LongContextRLTrainer: """ RL trainer for long-context reasoning using KeyChain tasks. """ def __init__(self, model, tokenizer): self.model = model self.tokenizer = tokenizer def create_keychain_dataset( self, short_qa_pairs: List[Tuple[str, str]], num_context_lengths: int = 4 ) -> List[KeyChainTask]: """ Convert short QA pairs into long-context KeyChain tasks. Generate variants at different context lengths. """ tasks = [] context_lengths = [4000, 8000, 16000, 32000] # Progressive difficulty for question, answer in short_qa_pairs: for context_length in context_lengths[:num_context_lengths]: task = KeyChainTask( original_question=question, answer=answer, num_documents=int(context_length / 150), # ~150 tokens per doc context_length=context_length ) tasks.append(task) return tasks def train_step(self, task: KeyChainTask) -> Dict: """ Single training step: model attempts KeyChain task. """ prompt = task.generate_documents_with_chain(self.tokenizer) # Generate response response = self.model.generate( prompt, max_tokens=200, temperature=0.7 ) # Verify: did model correctly trace chain AND answer question? chain_correct = all(uuid in response for uuid in task.chain) answer_correct = task.answer.lower() in response.lower() # Reward: both chain tracing and correct answer necessary reward = float(chain_correct and answer_correct) # RL loss (simplified DPO-style) loss = -reward # Reward maximization return { "reward": reward, "chain_correct": chain_correct, "answer_correct": answer_correct, "response_length": len(self.tokenizer.encode(response)) } def train_long_context_reasoning( model, tokenizer, short_qa_pairs: List[Tuple[str, str]], num_epochs: int = 3, context_length_start: int = 4000, context_length_end: int = 16000 ): """ Train model on progressively longer contexts using KeyChain synthesis. """ trainer = LongContextRLTrainer(model, tokenizer) # Generate tasks at multiple context lengths tasks = trainer.create_keychain_dataset(short_qa_pairs) for epoch in range(num_epochs): # Curriculum: increase context length over epochs current_length = int( context_length_start + (context_length_end - context_length_start) * epoch / num_epochs ) epoch_reward = 0.0 num_batches = 0 for task in tasks: # Adjust task to current curriculum length if task.context_length <= current_length: metrics = trainer.train_step(task) epoch_reward += metrics["reward"] num_batches += 1 avg_reward = epoch_reward / max(num_batches, 1) print(f"Epoch {epoch + 1}: Avg Reward={avg_reward:.4f}, " f"Context Length={current_length}") return model ``` The key insight is that UUID chain tracing induces the exact reasoning patterns needed for long-context understanding: planning (find first UUID), retrieval (locate next UUID), reasoning (interpret content), verification (confirm chain). ## Practical Guidance | Training Context | Test Context | Generalization | |-----------------|-------------|-----------------| | 4K | 4K | Baseline | | 8K | 16K | 2× extrapolation | | 16K | 64K | 4× extrapolation | | 16K | 128K | 8× extrapolation | **When to Use:** - Need long-context reasoning without expensive annotation - Models training on 16K, deploy on 128K contexts - Synthetic data generation is acceptable (no domain-specific docs required) - RL infrastructure available for training **When NOT to Use:** - Domain-specific long-context understanding (medical records, legal contracts) - Short-context tasks (KeyChain synthesis adds unnecessary overhead) - Batch processing (KeyChain works best with RL trajectory sampling) **Common Pitfalls:** - UUID chain too short (2-3 hops) → doesn't induce planning behavior - Too many distractor documents → overwhelming noise, model gives up - Not verifying BOTH chain tracing and answer correctness → incomplete learning - Fixed context length throughout training → poor generalization (use curriculum) - Distractor documents too similar to chain documents → model memorizes ## Reference [LoongRL: RL for Advanced Reasoning over Long Contexts](https://arxiv.org/abs/2510.19363)