Taming Momentum can improve the efficiency of training LLMs for agentic AI systems.

Applying "Taming Momentum" — a gradient update modification that clips or rescales the momentum term (e.g., via bounded exponential moving average of gradients, momentum norm capping, or sign-based momentum dampening) in first-order optimizers (Adam, AdamW, Lion) — reduces total compute-to-convergence (measured in FLOPs and wall-clock GPU hours) by ≥15% relative to vanilla AdamW baseline, without degrading final task performance (perplexity, downstream benchmark accuracy) by more than 1.5% absolute, when training transformer-based LLMs (≥1B parameters) on agentic AI task distributions (multi-step reasoning, tool-use, instruction-following corpora). The effect is hypothesized to arise from reduced oscillatory gradient dynamics in deep attention layers, yielding more stable loss landscapes and permitting larger effective learning rates or fewer warmup steps.

MINIMUM VIABLE TEST (MVT): Train a 1.3B parameter decoder-only transformer (GPT-NeoX architecture or LLaMA-style) on a 10B-token agentic corpus (mixture: 40% OpenHermes-2.5, 30% ToolBench instructions, 30% GSM8K-extended chain-of-thought) for 20,000 gradient steps. Compare four optimizer conditions: (A) AdamW baseline, (B) Taming Momentum as specified, (C) Lion optimizer, (D) AdamW with matched effective LR as ablation. Use 3 random seeds per condition (12 total runs). Primary metric: validation perplexity at matched FLOP budget. Secondary metrics: wall-clock time, gradient norm statistics, loss landscape curvature proxy (sharpness via SAM-style perturbation at 5 checkpoints).

FULL VALIDATION: Extend to 7B and 13B models, 100B-token training, 5 seeds, full agentic benchmark suite evaluation post-training.

# ============================================================
# Taming Momentum Experimental Validation — Implementation Sketch
# ============================================================

import torch
import torch.nn as nn
from torch.optim import Optimizer
import math
from dataclasses import dataclass
from typing import Optional, List, Dict, Tuple
import wandb
import numpy as np

# ---- 1. TAMING MOMENTUM OPTIMIZER -------------------------
class TamingMomentum(Optimizer):
    """
    Taming Momentum optimizer — AdamW variant with bounded momentum.
    Formulation (reconstruct from preprint):
      m_t = β1 * m_{t-1} + (1 - β1) * g_t          [standard EMA]
      m_t_tamed = m_t / max(1, ||m_t|| / τ)          [momentum norm clipping]
      v_t = β2 * v_{t-1} + (1 - β2) * g_t^2         [standard variance EMA]
      m_hat = m_t_tamed / (1 - β1^t)                 [bias correction]
      v_hat = v_t / (1 - β2^t)                        [bias correction]
      θ_t = θ_{t-1} - lr * m_hat / (sqrt(v_hat) + ε) [parameter update]
      θ_t = θ_t - lr * λ * θ_{t-1}                   [decoupled weight decay]
    
    τ (tau): momentum norm clipping threshold — KEY hyperparameter
    """
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0.1, tau=1.0):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                       weight_decay=weight_decay, tau=tau)
        super().__init__(params, defaults)
    
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        
        for group in self.param_groups:
            lr = group['lr']
            beta1, beta2 = group['betas']
            eps = group['eps']
            wd = group['weight_decay']
            tau = group['tau']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                
                grad = p.grad
                state = self.state[p]
                
                # Initialize state
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p)
                    state['exp_avg_sq'] = torch.zeros_like(p)
                
                state['step'] += 1
                t = state['step']
                m = state['exp_avg']
                v = state['exp_avg_sq']
                
                # Standard EMA updates
                m.mul_(beta1).add_(grad, alpha=1 - beta1)
                v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                
                # TAMING: clip momentum norm
                m_norm = m.norm(2)
                taming_scale = torch.clamp(tau / m_norm, max=1.0)
                m_tamed = m * taming_scale
                
                # Bias correction
                bc1 = 1 - beta1 ** t
                bc2 = 1 - beta2 ** t
                m_hat = m_tamed / bc1
                v_hat = v / bc2
                
                # Parameter update
                denom = v_hat.sqrt().add_(eps)
                p.addcdiv_(m_hat, denom, value=-lr)
                
                # Decoupled weight decay
                p.mul_(1 - lr * wd)
                
                # Logging hook
                state['momentum_norm'] = m_norm.item()
                state['taming_scale'] = taming_scale.item()
        
        return loss


# ---- 2. EXPERIMENT CONFIGURATION --------------------------
@dataclass
class ExperimentConfig:
    # Model
    model_size: str = "1.3B"          # "1.3B", "7B", "13B"
    n_layers: int = 24
    n_heads: int = 16
    d_model: int = 2048
    d_ff: int = 8192
    vocab_size: int = 32000
    max_seq_len: int = 2048
    
    # Training
    total_steps: int = 20000
    batch_size: int = 1024            # sequences per step
    gradient_accumulation: int = 4
    
    # Optimizer
    optimizer_type: str = "taming"    # "adamw", "taming", "lion", "adamw_lrmatch"
    lr: float = 3e-4
    beta1: float = 0.9
    beta2: float = 0.999
    weight_decay: float = 0.1
    tau: float = 1.0                  # Taming Momentum specific
    
    # Schedule
    warmup_steps: int = 2000
    lr_decay: str = "cosine"
    min_lr_ratio: float = 0.1
    
    # Experiment
    seed: int = 42
    n_seeds: int = 3
    log_interval: int = 100
    eval_interval: int = 500
    checkpoint_steps: List[int] = None
    
    def __post_init__(self):
        if self.checkpoint_steps is None:
            self.checkpoint_steps = [1000, 5000, 10000, 15000, 20000]


# ---- 3. FLOP COUNTER --------------------------------------
class FLOPCounter:
    """Track FLOPs using Chinchilla formula: 6ND per forward-backward."""
    def __init__(self, n_params: int, tokens_per_step: int):
        self.n_params = n_params
        self.tokens_per_step = tokens_per_step
        self.total_flops = 0
        self.flops_per_step = 6 * n_params * tokens_per_step
    
    def step(self):
        self.total_flops += self.flops_per_step
        return self.total_flops
    
    def flops_to_convergence(self, convergence_step: int) -> int:
        return convergence_step * self.flops_per_step


# ---- 4. TRAINING LOOP -------------------------------------
def train(config: ExperimentConfig, model, train_loader, val_loader,
          optimizer, scheduler, flop_counter, run_name: str):
    
    torch.manual_seed(config.seed)
    model.train()
    
    metrics = {
        'train_loss': [], 'val_perplexity': [], 'flops': [],
        'grad_norm': [], 'momentum_norm': [], 'wall_clock': [],
        'taming_scale': []
    }
    
    convergence_step = None
    target_perplexity = None  # Set after AdamW baseline run
    
    import time
    t0 = time.time()
    
    for step, batch in enumerate(train_loader):
        if step >= config.total_steps:
            break
        
        # Forward + backward
        loss = model(batch['input_ids'], labels=batch['labels']).loss
        loss = loss / config.gradient_accumulation
        loss.backward()
        
        if (step + 1) % config.gradient_accumulation == 0:
            # Gradient norm
            grad_norm = nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            
            total_flops = flop_counter.step()
            
            # Collect optimizer internals
            momentum_norms = []
            taming_scales = []
            for group in optimizer.param_groups:
                for p in group['params']:
                    if p in optimizer.state:
                        s = optimizer.state[p]
                        if 'momentum_norm' in s:
                            momentum_norms.append(s['momentum_norm'])
                        if 'taming_scale' in s:
                            taming_scales.append(s['taming_scale'])
            
            if step % config.log_interval == 0:
                metrics['train_loss'].append(loss.item() * config.gradient_accumulation)
                metrics['flops'].append(total_flops)
                metrics['grad_norm'].append(grad_norm.item())
                metrics['wall_clock'].append(time.time() - t0)
                if momentum_norms:
                    metrics['momentum_norm'].append(np.mean(momentum_norms))
                if taming_scales:
                    metrics['taming_scale'].append(np.mean(taming_scales))
                
                wandb.log({
                    'train/loss': metrics['train_loss'][-1],
                    'train/grad_norm': grad_norm.item(),
                    'train/flops': total_flops,
                    'optimizer/momentum_norm_mean': np.mean(momentum_norms) if momentum_norms else 0,
                    'optimizer/taming_scale_mean': np.mean(taming_scales) if taming_scales else 1,
                }, step=step)
            
            if step % config.eval_interval == 0:
                val_ppl = evaluate(model, val_loader)
                metrics['val_perplexity'].append(val_ppl)
                
                # Check convergence
                if target_perplexity is not None and val_ppl <= target_perplexity:
                    if convergence_step is None:
                        convergence_step = step
                
                wandb.log({'val/perplexity': val_ppl}, step=step)
            
            if step in config.checkpoint_steps:
                save_checkpoint(model, optimizer, step, metrics, run_name)
    
    return metrics, convergence_step


# ---- 5. SHARPNESS MEASUREMENT -----------------------------
def measure_sharpness(model, val_loader, n_perturbations=10, epsilon=0.01):
    """SAM-style sharpness: mean perplexity increase under random weight perturbation."""
    base_ppl = evaluate(model, val_loader)
    sharpness_values = []
    
    original_params = {n: p.clone() for n, p in model.named_parameters()}
    
    for _ in range(n_perturbations):
        # Add random perturbation
        with torch.no_grad():
            for p in model.parameters():
                noise = torch.randn_like(p) * epsilon
                p.add_(noise)
        
        perturbed_ppl = evaluate(model, val_loader)
        sharpness_values.append(perturbed_ppl - base_ppl)
        
        # Restore original params
        with torch.no_grad():
            for n, p in model.named_parameters():
                p.copy_(original_params[n])
    
    return np.mean(sharpness_values), np.std(sharpness_values)


# ---- 6. STATISTICAL ANALYSIS ------------------------------
def compute_efficiency_gain(adamw_flops_to_convergence: List[float],
                            taming_flops_to_convergence: List[float]) -> Dict:
    """
    Compute efficiency gain with statistical tests.
    Returns: mean gain, 95% CI, p-value, Cohen's d
    """
    from scipy import stats
    
    gains = [(a - t) / a * 100 for a, t in 
             zip(adamw_flops_to_convergence, taming_flops_to_convergence)]
    
    mean_gain = np.mean(gains)
    std_gain = np.std(gains, ddof=1)
    n = len(gains)
    
    # Paired t-test
    t_stat, p_value = stats.ttest_rel(adamw_flops_to_convergence,
                                       taming_flops_to_convergence)
    
    # Cohen's d
    pooled_std = np.sqrt((np.std(adamw_flops_to_convergence, ddof=1)**2 +
                          np.std(taming_flops_to_convergence, ddof=1)**2) / 2)
    cohens_d = (np.mean(adamw_flops_to_convergence) - 
                np.mean(taming_flops_to_convergence)) / pooled_std
    
    # 95% CI
    ci = stats.t.interval(0.95, df=n-1, loc=mean_gain,
                          scale=stats.sem(gains))
    
    return {
        'mean_gain_pct': mean_gain,
        'std_gain_pct': std_gain,
        'ci_95': ci,
        'p_value': p_value,
        'cohens_d': cohens_d,
        'n_seeds': n,
        'significant': p_value < 0.05 and mean_gain >= 15.0
    }


# ---- 7. MAIN EXPERIMENT RUNNER ----------------------------
def run_full_experiment():
    conditions = [
        ('adamw',        {'optimizer_type': 'adamw'}),
        ('taming',       {'optimizer_type': 'taming', 'tau': 1.0}),
        ('lion',         {'optimizer_type': 'lion'}),
        ('adamw_lrmatch',{'optimizer_type': 'adamw_lrmatch'}),  # ablation
    ]
    
    seeds = [42, 123, 777]
    results = {}
    
    for cond_name, cond_kwargs in conditions:
        results[cond_name] = {'flops_to_convergence': [], 'final_metrics': []}
        
        for seed in seeds:
            run_name = f"{cond_name}_seed{seed}_1.3B"
            config = ExperimentConfig(seed=seed, **cond_kwargs)
            
            wandb.init(project="taming-momentum-validation",
                      name=run_name, config=vars(config))
            
            model = build_model(config)
            optimizer = build_optimizer(model, config)
            scheduler = build_scheduler(optimizer, config)
            train_loader, val_loader = build_dataloaders(config)
            flop_counter = FLOPCounter(
                n_params=count_params(model),
                tokens_per_step=config.batch_size * config.max_seq_len
            )
            
            metrics, convergence_step = train(
                config, model, train_loader, val_loader,
                optimizer, scheduler, flop_counter, run_name
            )
            
            flops_conv = flop_counter.flops_to_convergence(convergence_step or config.total_steps)
            results[cond_name]['flops_to_convergence'].append(flops_conv)
            
            # Benchmark evaluation
            bench_results = run_benchmarks(model, config)
            results[cond_name]['final_metrics'].append(bench_results)
            
            wandb.finish()
    
    # Statistical analysis
    efficiency_stats = compute_efficiency_gain(
        results['adamw']['flops_to_convergence'],
        results['taming']['flops_to_convergence']
    )
    
    print(f"Efficiency gain: {efficiency_stats['mean_gain_pct']:.1f}% "
          f"(95% CI: {efficiency_stats['ci_95'][0]:.1f}–{efficiency_stats['ci_95'][1]:.1f}%)")
    print(f"p-value: {efficiency_stats['p_value']:.4f}, "
          f"Cohen's d: {efficiency_stats['cohens_d']:.2f}")
    print(f"Hypothesis supported: {efficiency_stats['significant']}")
    
    return results, efficiency_stats


# ---- 8. ABORT CHECKPOINT LOGIC ----------------------------
def check_abort_conditions(step: int, metrics: Dict, config: ExperimentConfig) -> Tuple[bool, str]:
    """Return (should_abort, reason)."""
    
    if step >= 1000:
        recent_loss = metrics['train_loss'][-10:]
        if min(recent_loss) > 0.9 * metrics['train_loss'][0]:
            return True, f"ABORT: No convergence by step {step} — loss not decreasing"
    
    if step >= 5000:
        if len(metrics['val_perplexity']) >= 2:
            ppl_trend = metrics['val_perplexity'][-1] - metrics['val_perplexity'][-5]
            if ppl_trend > 5.0:
                return True, f"ABORT: Perplexity diverging at step {step}"
    
    if step >= 2000:
        recent_grad_norms = metrics['grad_norm'][-20:]
        if np.mean(recent_grad_norms) > 10.0:
            return True, f"ABORT: Gradient explosion (mean norm {np.mean(recent_grad_norms):.1f})"
    
    return False, ""


if __name__ == "__main__":
    results, stats = run_full_experiment()