Skip to main content Home Creators adu2021 skillxiv looped-elastic-depth-transformers
looped-elastic-depth-transformers Enable budget-conditioned reasoning by repeatedly applying a shared transformer block stack with trajectory-based conditioning on time and step size. Train via shortcut-consistency loss to align shorter and full-length trajectories, enabling variable-depth inference without retraining.
Jump to install Skills Marketplace Discover and explore AI skills built by the community.
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
Show prompt details A direct command skips the review prompt. Inspect the source before running it.
npx skills add https://github.com/ADu2021/skillXiv --skill looped-elastic-depth-transformersThe command stays on one line. Scroll horizontally to inspect it before copying.
Prefer a local copy? Download the files currently available to SkillsMP.
Download Zip Downloading... More from this repository meaningful-kebab-case-name Convert arXiv papers into ready-to-use agent skills using category-aware extraction. First classifies the paper into one or more of 11 research categories, then applies a specialized extraction pipeline for each category — because different types of papers produce different types of usable knowledge. A single paper can yield multiple skills if it spans categories. Use this skill whenever the user wants to turn a paper into a skill, extract practical techniques from research, build a skill library from papers, convert arXiv papers into reusable agent instructions, or batch-process multiple papers into skills. Also trigger when someone asks about extracting actionable knowledge from papers, making research practical for LLM agents, or systematically converting academic contributions into structured agent capabilities.
action-quantization-behavior-cloning Establish regret bounds for behavior cloning with discretized actions combining statistical error and quantization error terms. Prove smoothness requirements for safe quantizer design, show that learning-based quantizers fail these requirements, and propose model-based augmentation to reduce error dependence from H² to H.
adaptive-lora-personalized-ranks Dynamically allocate LoRA ranks per-layer during fine-tuning instead of using fixed uniform ranks. Learn optimal rank for each layer and subject via variational framework with discretized exponential distribution, reducing memory footprint while maintaining fidelity and text-alignment.
Related occupations SOC
Based on SOC occupation classification
name looped-elastic-depth-transformers title LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning version 0.0.2 engine skillxiv-v0.0.2-claude-opus-4.6 license MIT url https://arxiv.org/abs/2602.11451 keywords ["Transformer Architecture","Elastic Depth","Looped Processing","Budget-Conditioned Reasoning","Self-Distillation"] description Enable budget-conditioned reasoning by repeatedly applying a shared transformer block stack with trajectory-based conditioning on time and step size. Train via shortcut-consistency loss to align shorter and full-length trajectories, enabling variable-depth inference without retraining.
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
"""
t = iteration_idx / max (1 , total_iterations - 1
if
is
None
1.0
self
self
return
def
_sinusoidal_encode
self, value: float
"""Sine-cosine encoding of scalar value."""
self
for
in
range
0
self
2
1
10000
self
if
self
if
1
self
1
return
def
condition_via_adaLN
self,
hidden_states: torch.Tensor,
time_embedding: torch.Tensor,
step_embedding: torch.Tensor
"""
Apply trajectory conditioning via learned modulation (AdaLN-style).
Learns: scale and bias that depend on time and step.
"""
1.0
0.1
0.05
return
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
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)
)
self .attn_gate = nn.Parameter(torch.ones(1 ))
self .ffn_gate = nn.Parameter(torch.ones(1 ))
def forward (self, x ):
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,
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):
time_emb, step_emb = self .conditioner.encode_trajectory_position(
loop_idx, num_loops
)
for block_idx, block in enumerate (self .transformer_blocks):
current_state = block(current_state)
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_output, _ = self .model.forward_looped(embeddings, full_loops)
short_output, _ = self .model.forward_looped(embeddings, shortcut_loops)
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 )
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
for num_loops in sample_loops:
if num_loops >= max_loops:
continue
short_output, _ = self .model.forward_looped(embeddings, num_loops)
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' :
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)
loss = consistency_trainer.compute_multi_trajectory_loss(
embeddings, target,
max_loops=full_loops,
sample_loops=shortcut_loops_list
)
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:.4 f} " )
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' ]:.4 f} " )
return model
Practical Guidance When to use : Variable-length reasoning; inference with heterogeneous compute budgets; amortized reasoning
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
Variable-depth inference without retraining
Shorter trajectories remain informative
Shared parameters reduce memory
Smooth quality scaling with loops
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
