adu2021 Retrofit Pretrained Models with Recurrent Depth for Efficient Reasoning
Standard pretrained language models allocate the same computational budget at training and inference. Retrofitted recurrence decouples these budgets—a model can be trained efficiently while allocating more compute to inference when needed. By gradually introducing recurrence depth during training via curriculum learning, existing models learn to perform deeper internal reasoning and refine solutions iteratively.
The key insight is that pretrained models can be efficiently converted to recurrent variants without retraining from scratch. On mathematical reasoning tasks, this approach outperforms continued training of original non-recurrent models at the same inference compute budget, demonstrating that depth-recurrence provides meaningful efficiency gains.
Core Concept
Retrofitted recurrence treats inference-time reasoning depth as a learnable skill. Rather than the fixed depth of a traditional transformer stack, recurrent models can iterate multiple times over the same (or evolving) internal representations, refining outputs with each iteration. A curriculum training strategy gradually increases recurrence depth from 1 to the target maximum, allowing the model to learn iterative refinement patterns naturally.
The approach converts the inference process from "compute once and output" to "compute, reflect, refine, iterate"—enabling expensive reasoning at test time without expensive training.
Architecture Overview
- Base Pretrained Model: Non-recurrent transformer (frozen or fine-tuned)
- Recurrence Controller: Manages the number of reasoning iterations at inference time
- Depth Curriculum: Progressively increases recurrence depth during training (1 → 2 → 4 → 8 iterations)
- Refinement Mechanism: At each iteration, the model refines previous outputs or internal states
- Computational Budget: Separates training cost (linear in depth) from inference cost (configurable depth)
- Loss Aggregation: Supervises refinement at each intermediate depth, not just final output
Implementation Steps
Step 1: Define Recurrent Refinement Layer
Create a module that applies iterative refinement on top of base model outputs.
import torch
import torch.nn as nn
class RecurrentRefinementLayer(nn.Module):
"""
Enables iterative refinement of LLM outputs through recurrence.
"""
def __init__(self, hidden_dim, num_refinement_layers=2):
"""
Args:
hidden_dim: Dimension of model hidden states
num_refinement_layers: Number of refinement stacks
"""
super().__init__()
self.hidden_dim = hidden_dim
# Stacks for refining outputs at each recurrence step
self.refinement_blocks = nn.ModuleList([
nn.TransformerDecoderLayer(
d_model=hidden_dim,
nhead=8,
dim_feedforward=hidden_dim * 4,
batch_first=True
)
for _ in range(num_refinement_layers)
])
def forward(self, base_output, memory, current_depth):
"""
Refine outputs through recurrent iterations.
Args:
base_output: Initial output from base model [batch, seq_len, hidden_dim]
memory: Encoder memory from base model [batch, seq_len, hidden_dim]
current_depth: Number of iterations to perform (curriculum-based)
Returns:
refined_output: Progressively refined output
intermediate_outputs: Outputs at each refinement step
"""
refined = base_output
intermediate_outputs = [base_output]
for step in range(current_depth):
# Apply refinement layer
block_idx = step % len(self.refinement_blocks)
refined = self.refinement_blocks[block_idx](
refined, memory
)
intermediate_outputs.append(refined)
return refined, intermediate_outputs
Step 2: Implement Depth Curriculum
Define a curriculum that gradually increases recurrence depth during training.
class DepthCurriculum:
"""
Manages progressive increase in recurrence depth.
"""
def __init__(self, max_depth=8, warmup_steps=5000, growth_schedule='linear'):
"""
Args:
max_depth: Maximum recurrence iterations (e.g., 8)
warmup_steps: Steps over which to ramp up depth
growth_schedule: 'linear', 'exponential', or 'step' (phase-based)
"""
self.max_depth = max_depth
self.warmup_steps = warmup_steps
self.schedule = growth_schedule
self.current_step = 0
def get_current_depth(self):
"""
Return recurrence depth for current training step.
Returns:
depth: Integer in [1, max_depth]
"""
if self.schedule == 'linear':
# Linear increase: 1 + (max_depth-1) * (step / warmup)
progress = min(self.current_step / self.warmup_steps, 1.0)
depth = int(1 + (self.max_depth - 1) * progress)
elif self.schedule == 'exponential':
# Exponential: depth increases slowly then rapidly
progress = min(self.current_step / self.warmup_steps, 1.0)
depth = int(1 + (self.max_depth - 1) * (progress ** 0.5))
elif self.schedule == 'step':
# Phase-based: 1 → 2 → 4 → 8
phases = [1000, 2000, 3000, self.warmup_steps]
depth = 1
for phase_step in phases:
if self.current_step >= phase_step:
depth = min(depth + 1, self.max_depth)
else:
depth = self.max_depth
self.current_step += 1
return depth
Step 3: Supervised Loss with Intermediate Supervision
Train on outputs at each refinement depth, encouraging progressive improvement.
def compute_depth_curriculum_loss(
intermediate_outputs, target_output,
base_loss_fn=nn.CrossEntropyLoss(), depth_weights=None):
"""
Compute loss supervising refinement at all depths.
Args:
intermediate_outputs: List of outputs at each depth
[(batch, seq_len, vocab_size), ...]
target_output: Ground truth output (batch, seq_len)
base_loss_fn: Loss function (CrossEntropyLoss, etc.)
depth_weights: Optional weights for each depth (encourage deep outputs)
Returns:
total_loss: Weighted sum of losses at all intermediate depths
"""
if depth_weights is None:
# Default: final output weighted most
num_depths = len(intermediate_outputs)
depth_weights = torch.linspace(0.5, 1.0, num_depths)
depth_weights = depth_weights / depth_weights.sum()
total_loss = 0.0
for depth, (output, weight) in enumerate(zip(intermediate_outputs, depth_weights)):
# Reshape for loss computation
batch_size, seq_len = target_output.shape
logits = output.view(-1, output.shape[-1])
targets = target_output.view(-1)
# Compute loss
depth_loss = base_loss_fn(logits, targets)
total_loss += weight * depth_loss
return total_loss
Step 4: Integrated Training Loop
Combine curriculum scheduling with recurrent refinement training.
def train_with_depth_curriculum(
model, refinement_layer, train_dataset,
max_depth=8, num_epochs=3):
"""
Training loop applying depth curriculum.
Args:
model: Base pretrained language model
refinement_layer: RecurrentRefinementLayer instance
train_dataset: Training data with inputs and targets
max_depth: Maximum recurrence depth
num_epochs: Number training epochs
"""
import torch.optim as optim
optimizer = optim.Adam(
list(model.parameters()) + list(refinement_layer.parameters()),
lr=1e-5
)
curriculum = DepthCurriculum(max_depth=max_depth, warmup_steps=5000)
for epoch in range(num_epochs):
for batch in train_dataset:
input_ids = batch['input_ids']
target_ids = batch['target_ids']
# Forward pass through base model
base_output = model.forward(input_ids, return_hidden_states=True)
hidden_states = base_output.hidden_states[-1] # Last layer
# Get current curriculum depth
current_depth = curriculum.get_current_depth()
# Apply recurrent refinement
refined_output, intermediate = refinement_layer(
hidden_states, hidden_states, current_depth
)
# Convert hidden states to logits
logits_list = [
model.lm_head(h) for h in intermediate
]
# Compute depth-curriculum loss
loss = compute_depth_curriculum_loss(
logits_list, target_ids,
depth_weights=torch.linspace(0.5, 1.0, len(logits_list))
)
# Backward pass
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if curriculum.current_step % 500 == 0:
avg_loss = loss.item()
print(f"Step {curriculum.current_step}: Loss {avg_loss:.4f}, "
f"Depth {current_depth}")
return model, refinement_layer
Step 5: Inference with Configurable Depth
At test time, use desired recurrence depth independent of training.
def inference_with_recurrence(model, refinement_layer, prompt, target_depth=4):
"""
Generate text using recurrent refinement at specified depth.
Args:
model: Trained base model
refinement_layer: Trained RecurrentRefinementLayer
prompt: Input prompt text
target_depth: Number of refinement iterations (configurable)
Returns:
final_output: Refined text prediction
refinement_path: Outputs at each iteration (for analysis)
"""
input_ids = model.tokenize(prompt)
refinement_path = []
with torch.no_grad():
# Initial generation from base model
base_output = model.forward(input_ids, return_hidden_states=True)
hidden_states = base_output.hidden_states[-1]
# Apply recurrent refinement at desired depth
refined, intermediate = refinement_layer(
hidden_states, hidden_states, target_depth
)
# Convert to logits and decode
for h in intermediate:
logits = model.lm_head(h)
tokens = torch.argmax(logits, dim=-1)
text = model.decode(tokens)
refinement_path.append(text)
return refinement_path[-1], refinement_path
Practical Guidance
When to Use Retrofitted Recurrence:
- Mathematical and logical reasoning tasks (benefit from iterative refinement)
- Scenarios where inference compute budget is flexible (test-time scales)
- Models where existing architecture is effective (minimize retraining)
When NOT to Use:
- Real-time systems with strict latency requirements (recurrence adds inference time)
- Tasks with diminishing returns to depth (simple tasks don't benefit)
- Models requiring immediate single-pass inference
Hyperparameters and Configuration:
- Max depth: 4-8 for most tasks (balance benefit vs. latency)
- Warmup steps: 5000+ for stable curriculum (longer for larger models)
- Curriculum schedule: 'linear' for steady progression; 'exponential' for aggressive later growth
- Intermediate supervision weight: Increase toward final depth (e.g., 0.5 → 1.0)
Pitfalls to Avoid:
- Too-aggressive curriculum - Depth increases too quickly; model can't learn refinement patterns
- Ignoring latency - Each refinement iteration costs compute; set realistic depth budgets for inference
- Loss imbalance - Heavy weight on intermediate outputs can prevent deep refinement learning
- Cold start - Start with depth=1 before increasing; jumping to depth=4 causes training instability
Reference: https://arxiv.org/abs/2511.07384