Taming Momentum can reduce the memory footprint of optimizers used in training LLMs for financial trading agent systems.

Applying the "Taming Momentum" technique (bounded/clipped momentum updates) to standard adaptive optimizers (e.g., Adam, AdamW) reduces peak GPU memory consumption by ≥15% during training of LLMs (≥1B parameters) used as financial trading agents, without degrading downstream trading performance (Sharpe ratio, PnL) by more than 5% relative to the baseline optimizer.

Minimum Viable Test (MVT): Fine-tune a 1.3B parameter causal LM (e.g., GPT-2 XL or OPT-1.3B) on a financial text+time-series dataset using (a) AdamW baseline and (b) Taming Momentum variant. Profile peak GPU memory, convergence speed, and downstream trading metrics on a fixed compute budget of 4× A100 80GB GPUs for 72 hours.

Full Validation: Scale to 7B and 13B parameter models, sweep momentum clipping thresholds τ ∈ {0.1, 0.5, 1.0, 2.0, 5.0}, compare against memory-efficient baselines (Adafactor, 8-bit Adam, CAME), and evaluate on live paper-trading simulation over 30 days.

# Taming Momentum AdamW Implementation
import torch
from torch.optim import Optimizer

class TamingMomentumAdamW(Optimizer):
    """
    AdamW with momentum taming (clipped momentum buffer).
    Memory savings come from bounded momentum values reducing
    the effective dynamic range of the m_t buffer,
    enabling potential quantization or sparse storage.
    
    Args:
        params: model parameters
        lr: learning rate (default 1e-4)
        betas: (beta1, beta2) momentum coefficients
        eps: numerical stability term
        weight_decay: L2 regularization
        tau: momentum clipping threshold (KEY HYPERPARAMETER)
    """
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.999),
                 eps=1e-8, weight_decay=0.01, 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 = closure() if closure is not None else None
        
        for group in self.param_groups:
            tau = group['tau']
            beta1, beta2 = group['betas']
            
            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
                    # First moment (momentum) - TAMED
                    state['exp_avg'] = torch.zeros_like(p)
                    # Second moment (variance)
                    state['exp_avg_sq'] = torch.zeros_like(p)
                
                state['step'] += 1
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                
                # Standard second moment update
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                
                # TAMING MOMENTUM: clip before EMA update
                # This bounds the momentum buffer's dynamic range
                tamed_grad = torch.clamp(grad, -tau, tau)
                exp_avg.mul_(beta1).add_(tamed_grad, alpha=1 - beta1)
                
                # Bias correction
                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']
                
                # Compute step size
                step_size = group['lr'] / bias_correction1
                denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
                
                # Weight decay (decoupled)
                p.mul_(1 - group['lr'] * group['weight_decay'])
                
                # Parameter update
                p.addcdiv_(exp_avg, denom, value=-step_size)
        
        return loss


# Memory Profiling Harness
def profile_optimizer_memory(model, optimizer_class, optimizer_kwargs,
                              dataloader, n_steps=100):
    """
    Profile peak GPU memory for a given optimizer.
    Returns: dict with peak_memory_mb, avg_step_time_ms, final_loss
    """
    import time
    
    optimizer = optimizer_class(model.parameters(), **optimizer_kwargs)
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()
    
    losses = []
    step_times = []
    
    for step, batch in enumerate(dataloader):
        if step >= n_steps:
            break
        
        t0 = time.perf_counter()
        
        optimizer.zero_grad()
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        
        losses.append(loss.item())
        step_times.append((t1 - t0) * 1000)
    
    peak_memory_mb = torch.cuda.max_memory_allocated() / (1024 ** 2)
    
    return {
        'peak_memory_mb': peak_memory_mb,
        'avg_step_time_ms': sum(step_times) / len(step_times),
        'final_loss': losses[-1],
        'loss_curve': losses
    }


# Experiment Runner
def run_comparison_experiment(model_name, tau_values, seeds, n_steps=5000):
    """
    Full comparison: AdamW vs TamingMomentumAdamW across tau values and seeds.
    """
    results = {}
    
    for seed in seeds:
        torch.manual_seed(seed)
        
        # Load model and financial dataset
        model = load_financial_llm(model_name)  # OPT-1.3B, etc.
        dataloader = load_financial_dataloader(seed=seed)
        
        # Baseline: AdamW
        baseline_results = profile_optimizer_memory(
            model, torch.optim.AdamW,
            {'lr': 1e-4, 'weight_decay': 0.01},
            dataloader, n_steps
        )
        results[f'adamw_seed{seed}'] = baseline_results
        
        # Taming Momentum sweep
        for tau in tau_values:
            taming_results = profile_optimizer_memory(
                model, TamingMomentumAdamW,
                {'lr': 1e-4, 'weight_decay': 0.01, 'tau': tau},
                dataloader, n_steps
            )
            results[f'taming_tau{tau}_seed{seed}'] = taming_results
    
    # Statistical analysis
    memory_reduction = compute_memory_reduction_stats(results)
    return results, memory_reduction


# Trading Agent Evaluation
def evaluate_trading_agent(model, test_env, n_episodes=252):
    """
    Evaluate trained LLM trading agent on holdout financial data.
    Returns Sharpe ratio, max drawdown, annualized return.
    """
    portfolio_returns = []
    
    for episode in range(n_episodes):
        obs = test_env.reset()
        done = False
        episode_return = 0
        
        while not done:
            # LLM generates trading action from market context
            action = model.generate_trading_action(obs)
            obs, reward, done, info = test_env.step(action)
            episode_return += reward
        
        portfolio_returns.append(episode_return)
    
    sharpe = compute_sharpe_ratio(portfolio_returns, risk_free_rate=0.05)
    max_dd = compute_max_drawdown(portfolio_returns)
    ann_return = compute_annualized_return(portfolio_returns)
    
    return {'sharpe': sharpe, 'max_drawdown': max_dd, 'ann_return': ann_return}