adu2021 LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning
Problem Context
Transformers with fixed depth have fixed computation. Extended reasoning requires more depth, but retraining is expensive. LoopFormer enables variable-depth inference by repeatedly applying the same transformer blocks, with trajectories conditioned on normalized time (0 to 1 over steps) and step size. Shorter trajectories remain informative while longer ones refine.
Core Concept
LoopFormer uses: (1) a shared stack of K transformer blocks applied repeatedly, (2) trajectory-based conditioning via time t ∈ [0,1] and step size Δt indicating position in loop, (3) shortcut-consistency training that aligns trajectories of different lengths, (4) AdaLN-style modulation of attention/FFN residual strengths.
Implementation
Step 1: Trajectory-based conditioning
import torch
import torch.nn as nn
import math
from typing import Tuple
class TrajectoryConditioner:
"""Encode time and step size for loop iteration."""
def __init__(self, dim: int = 768):
self.dim = dim
def encode_trajectory_position(
self,
iteration_idx: int,
total_iterations: int,
step_size: float = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Encode normalized time t and step size Δt.
Args:
iteration_idx: Current loop iteration (0-based)
total_iterations: Total number of iterations
step_size: Size of current step (fraction of sequence)
Returns:
(time_embedding, step_embedding): [dim] tensors
"""
# Normalized time: 0 to 1 over course of all iterations
t = iteration_idx / max(1, total_iterations - 1)
# Step size (default: uniform 1/total_iterations)
if step_size is None:
step_size = 1.0 / total_iterations
# Sine-cosine positional encoding
time_embedding = self._sinusoidal_encode(t)
step_embedding = self._sinusoidal_encode(step_size)
return time_embedding, step_embedding
def _sinusoidal_encode(self, value: float) -> torch.Tensor:
"""Sine-cosine encoding of scalar value."""
embedding = torch.zeros(self.dim)
for i in range(0, self.dim, 2):
omega = 1 / (10000 ** (i / self.dim))
if i < self.dim:
embedding[i] = math.sin(value * omega)
if i + 1 < self.dim:
embedding[i + 1] = math.cos(value * omega)
return embedding
def condition_via_adaLN(
self,
hidden_states: torch.Tensor, # [batch, seq_len, dim]
time_embedding: torch.Tensor,
step_embedding: torch.Tensor
) -> torch.Tensor:
"""
Apply trajectory conditioning via learned modulation (AdaLN-style).
Learns: scale and bias that depend on time and step.
"""
combined = time_embedding + step_embedding # [dim]
# Learned projection: embedding -> (scale, bias)
# Would be part of model parameters
scale = 1.0 + 0.1 * combined.mean()
bias = 0.05 * combined.mean()
return hidden_states * scale + bias
Step 2: Looped transformer with shared blocks
class LoopedTransformer(nn.Module):
"""Transformer with shared block stack applied repeatedly."""
def __init__(
self,
dim: int = 768,
num_blocks: int = 6,
num_heads: int = 12,
max_loops: int = 16
):
super().__init__()
self.dim = dim
self.num_blocks = num_blocks
self.max_loops = max_loops
# Shared transformer blocks (reused across loops)
self.transformer_blocks = nn.ModuleList([
self._create_transformer_block(dim, num_heads)
for _ in range(num_blocks)
])
self.conditioner = TrajectoryConditioner(dim)
def _create_transformer_block(
self,
dim: int,
num_heads: int
) -> nn.Module:
"""Create single transformer block with gating."""
class GatedTransformerBlock(nn.Module):
def __init__(self, dim, num_heads):
super().__init__()
self.attn = nn.MultiheadAttention(dim, num_heads)
self.ffn = nn.Sequential(
nn.Linear(dim, 4 * dim),
nn.GELU(),
nn.Linear(4 * dim, dim)
)
# Learned gates for residual connections
self.attn_gate = nn.Parameter(torch.ones(1))
self.ffn_gate = nn.Parameter(torch.ones(1))
def forward(self, x):
# Gated residuals
attn_out, _ = self.attn(x, x, x)
x = x + self.attn_gate * attn_out
ffn_out = self.ffn(x)
x = x + self.ffn_gate * ffn_out
return x
return GatedTransformerBlock(dim, num_heads)
def forward_looped(
self,
embeddings: torch.Tensor, # [batch, seq_len, dim]
num_loops: int,
return_intermediate: bool = False
) -> Tuple[torch.Tensor, list]:
"""
Apply transformer blocks repeatedly in loops.
Args:
embeddings: Input embeddings
num_loops: Number of times to apply block stack
return_intermediate: Return all intermediate states
Returns:
(final_embeddings, intermediate_states)
"""
current_state = embeddings
intermediates = [current_state.clone()]
for loop_idx in range(num_loops):
# Trajectory conditioning
time_emb, step_emb = self.conditioner.encode_trajectory_position(
loop_idx, num_loops
)
# Apply all blocks in stack
for block_idx, block in enumerate(self.transformer_blocks):
current_state = block(current_state)
# Apply trajectory conditioning (modulate residuals for this loop)
current_state = self.conditioner.condition_via_adaLN(
current_state, time_emb, step_emb
)
if return_intermediate:
intermediates.append(current_state.clone())
return current_state, intermediates if return_intermediate else []
Step 3: Shortcut and consistency loss
class ShortcutConsistencyTraining:
"""Train via multi-length trajectories with consistency."""
def __init__(self, model: LoopedTransformer):
self.model = model
def compute_consistency_loss(
self,
embeddings: torch.Tensor,
full_loops: int = 16,
shortcut_loops: int = 4
) -> torch.Tensor:
"""
Compute loss aligning short and full trajectories.
Strategy: Full trajectory as teacher, short trajectory as student.
"""
# Full trajectory (teacher)
full_output, _ = self.model.forward_looped(embeddings, full_loops)
# Shortcut trajectory (student, fewer loops)
short_output, _ = self.model.forward_looped(embeddings, shortcut_loops)
# Consistency loss: MSE between outputs
consistency_loss = torch.nn.functional.mse_loss(short_output, full_output)
return consistency_loss
def compute_multi_trajectory_loss(
self,
embeddings: torch.Tensor,
target: torch.Tensor,
max_loops: int = 16,
sample_loops: list = None
) -> torch.Tensor:
"""
Train on multiple trajectory lengths simultaneously.
Args:
sample_loops: Which loop lengths to train on
(default: [4, 8, 12, 16])
"""
if sample_loops is None:
sample_loops = [4, 8, 12, 16]
total_loss = torch.tensor(0.0, requires_grad=True)
# Main LM loss on full trajectory
full_output, _ = self.model.forward_looped(embeddings, max_loops)
lm_loss = torch.nn.functional.cross_entropy(
full_output.view(-1, full_output.size(-1)),
target.view(-1)
)
total_loss = total_loss + lm_loss
# Shortcut loss: train shorter trajectories to match full
for num_loops in sample_loops:
if num_loops >= max_loops:
continue
short_output, _ = self.model.forward_looped(embeddings, num_loops)
# Stop-gradient on full output (don't update via short trajectory)
shortcut_loss = torch.nn.functional.mse_loss(
short_output,
full_output.detach()
)
total_loss = total_loss + 0.1 * shortcut_loss
return total_loss
Step 4: Inference with elastic depth
class ElasticDepthInference:
"""Perform inference with variable computational budget."""
def __init__(self, model: LoopedTransformer):
self.model = model
def generate_with_budget(
self,
prompt_embeddings: torch.Tensor,
budget_loops: int,
schedule: str = 'linear'
) -> torch.Tensor:
"""
Generate with specified loop budget (variable depth).
Args:
budget_loops: Number of loops available
schedule: How to allocate loops ('linear', 'exponential')
Returns:
output: Generated embeddings
"""
if schedule == 'linear':
num_loops = budget_loops
elif schedule == 'exponential':
# Early layers get fewer loops; later layers refine
num_loops = min(budget_loops, int(math.sqrt(budget_loops)) * 4)
else:
num_loops = budget_loops
output, _ = self.model.forward_looped(prompt_embeddings, num_loops)
return output
def benchmark_depth_vs_quality(
self,
test_embeddings: torch.Tensor,
target: torch.Tensor,
max_budget: int = 16
) -> dict:
"""Profile quality vs. computational budget."""
results = {}
for budget in range(1, max_budget + 1):
output = self.generate_with_budget(test_embeddings, budget)
loss = torch.nn.functional.cross_entropy(
output.view(-1, output.size(-1)),
target.view(-1)
)
results[budget] = {
'loss': loss.item(),
'computation': budget * len(self.model.transformer_blocks)
}
return results
Step 5: Full training loop
def train_loopformer(
model: LoopedTransformer,
train_loader,
optimizer,
num_epochs: int = 10,
full_loops: int = 16,
shortcut_loops_list: list = None,
device: str = 'cuda'
):
"""
Train LoopFormer with shortcut-consistency.
Args:
shortcut_loops_list: Which shortcut depths to train on
"""
if shortcut_loops_list is None:
shortcut_loops_list = [4, 8, 12]
consistency_trainer = ShortcutConsistencyTraining(model)
for epoch in range(num_epochs):
total_loss = 0.0
num_batches = 0
for batch_idx, (embeddings, target) in enumerate(train_loader):
embeddings = embeddings.to(device)
target = target.to(device)
# Multi-trajectory loss
loss = consistency_trainer.compute_multi_trajectory_loss(
embeddings, target,
max_loops=full_loops,
sample_loops=shortcut_loops_list
)
# Update
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
num_batches += 1
avg_loss = total_loss / num_batches
print(f"Epoch {epoch + 1}: Loss={avg_loss:.4f}")
# Validate at different depths
if (epoch + 1) % 5 == 0:
print(" Depth vs Quality:")
inference = ElasticDepthInference(model)
results = inference.benchmark_depth_vs_quality(
embeddings, target, max_budget=full_loops
)
for budget, metrics in results.items():
print(f" {budget} loops: loss={metrics['loss']:.4f}")
return model
Practical Guidance
When to use: Variable-length reasoning; inference with heterogeneous compute budgets; amortized reasoning
Hyperparameters:
- num_blocks: 4-8 (shared blocks)
- max_loops: 12-20 (maximum depth)
- shortcut_loops: [4, 8, 12] (training depths)
- shortcut_loss_weight: 0.05-0.2 (balance main vs. consistency)
- num_heads: 8-16
Key advantages:
- Variable-depth inference without retraining
- Shorter trajectories remain informative
- Shared parameters reduce memory
- Smooth quality scaling with loops
Common pitfalls:
- shortcut_loss too strong → longer trajectories don't improve
- Max loops too large → diminishing returns, instability
- Not validating intermediate trajectories actually work
- Skip-consistency not tight enough
Scaling: Linear in num_loops; shared blocks reduce memory vs. fixed-depth stacks.
Reference
Paper: https://arxiv.org/abs/2602.11451 Related work: Elastic networks, adaptive computation, trajectory-based conditioning Benchmarks: Long reasoning sequences, language modeling