Low-rank EMA reformulation from Taming Momentum can reduce the memory footprint of optimizer states when training multi-agent LLM systems for financial applications.

Applying a low-rank exponential moving average (EMA) reformulation — as introduced in the Taming Momentum framework — to optimizer state matrices (specifically the first- and second-moment tensors in Adam-family optimizers) reduces peak optimizer-state memory consumption by ≥40% relative to full-rank Adam baselines, without degrading final validation loss by more than 2% (relative), when training multi-agent large language model (LLM) systems (≥7B total parameters across agents) on financial time-series or financial NLP tasks. The claim is falsifiable: if memory reduction is <40% OR validation loss degrades >2% relative under matched compute budgets, the hypothesis is refuted.

MINIMUM VIABLE TEST (MVT): Train two 7B-parameter two-agent LLM systems (actor + critic or two collaborative agents) on a financial corpus for 20,000 steps:

• Condition A: Full-rank AdamW (baseline)
• Condition B: Low-rank EMA AdamW (Taming Momentum, rank r=64)
• Condition C: Low-rank EMA AdamW (rank r=128) as ablation Measure: (1) peak optimizer-state VRAM, (2) validation perplexity on FinText held-out set, (3) tokens/sec throughput, (4) Frobenius approximation error per layer. Run 3 seeds per condition. Statistical test: paired Wilcoxon signed-rank on memory measurements across layers; two-sided t-test on final validation loss.

FULL VALIDATION: Extend to 13B and 70B two-agent systems. Add financial downstream tasks (FinQA, FiNER-139, synthetic portfolio optimization RL environment). Include wall-clock-matched comparisons (not just step-matched). Add GradNorm and loss landscape sharpness (Hessian trace approximation via Hutchinson estimator) as secondary metrics.

# ============================================================
# LOW-RANK EMA ADAMW — CORE IMPLEMENTATION SKETCH
# ============================================================

import torch
from torch.optim import Optimizer
import torch.linalg as LA

class LowRankEMAAdamW(Optimizer):
    """
    AdamW with low-rank EMA reformulation (Taming Momentum).
    Stores moment tensors as rank-r SVD factors instead of full matrices.
    
    Memory: O(r*(d+k)) per weight matrix W in R^{d x k}
    vs O(d*k) for standard AdamW.
    Reduction ratio: r*(d+k) / (d*k) = r*(1/k + 1/d)
    For d=k=4096, r=64: 64*8192 / (4096*4096) = 524288/16777216 = 3.1%
    → 96.9% reduction per layer (before accounting for embeddings/biases)
    """
    
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0.01, rank=64, svd_update_freq=1):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                       weight_decay=weight_decay, rank=rank,
                       svd_update_freq=svd_update_freq)
        super().__init__(params, defaults)
    
    def _get_lowrank_factors(self, tensor, rank):
        """Compute truncated SVD, return low-rank factors."""
        if tensor.dim() < 2:
            return None  # scalars/biases: use full storage
        
        # Reshape to 2D if needed (e.g., conv filters)
        orig_shape = tensor.shape
        mat = tensor.reshape(tensor.shape[0], -1)
        
        # Truncated SVD — top-r components only
        # torch.linalg.svd is O(min(d,k)^2 * max(d,k))
        # For large matrices, use randomized SVD for speed
        U, S, Vh = LA.svd(mat, full_matrices=False)
        U_r = U[:, :rank]      # (d, r)
        S_r = S[:rank]         # (r,)
        Vh_r = Vh[:rank, :]    # (r, k)
        
        return U_r, S_r, Vh_r, orig_shape
    
    def _reconstruct_from_lowrank(self, U_r, S_r, Vh_r, orig_shape):
        """Reconstruct approximate moment tensor from low-rank factors."""
        mat = U_r @ torch.diag(S_r) @ Vh_r
        return mat.reshape(orig_shape)
    
    @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:
            rank = group['rank']
            beta1, beta2 = group['betas']
            lr = group['lr']
            eps = group['eps']
            wd = group['weight_decay']
            svd_freq = group['svd_update_freq']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                
                grad = p.grad.data
                state = self.state[p]
                
                # Initialize state
                if len(state) == 0:
                    state['step'] = 0
                    state['use_lowrank'] = (p.dim() >= 2 and 
                                           min(p.shape) > rank)
                    
                    if state['use_lowrank']:
                        # Initialize low-rank moment factors
                        # m_t (first moment) low-rank factors
                        state['m_U'], state['m_S'], state['m_Vh'], \
                            state['orig_shape'] = \
                            self._get_lowrank_factors(
                                torch.zeros_like(p.data), rank)
                        # v_t (second moment) low-rank factors  
                        state['v_U'], state['v_S'], state['v_Vh'], _ = \
                            self._get_lowrank_factors(
                                torch.zeros_like(p.data), rank)
                    else:
                        # Fallback: full storage for small tensors
                        state['exp_avg'] = torch.zeros_like(p.data)
                        state['exp_avg_sq'] = torch.zeros_like(p.data)
                
                state['step'] += 1
                t = state['step']
                
                # Weight decay
                p.data.mul_(1 - lr * wd)
                
                if state['use_lowrank']:
                    # Reconstruct current moments from low-rank factors
                    m_t = self._reconstruct_from_lowrank(
                        state['m_U'], state['m_S'], state['m_Vh'],
                        state['orig_shape'])
                    v_t = self._reconstruct_from_lowrank(
                        state['v_U'], state['v_S'], state['v_Vh'],
                        state['orig_shape'])
                    
                    # Standard EMA updates
                    m_t.mul_(beta1).add_(grad, alpha=1 - beta1)
                    v_t.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                    
                    # Re-project to low-rank (every svd_freq steps)
                    if t % svd_freq == 0:
                        state['m_U'], state['m_S'], state['m_Vh'], _ = \
                            self._get_lowrank_factors(m_t, rank)
                        state['v_U'], state['v_S'], state['v_Vh'], _ = \
                            self._get_lowrank_factors(v_t, rank)
                    
                    # Bias correction
                    bc1 = 1 - beta1 ** t
                    bc2 = 1 - beta2 ** t
                    m_hat = m_t / bc1
                    v_hat = v_t / bc2
                    
                    # Parameter update
                    p.data.addcdiv_(m_hat, v_hat.sqrt().add_(eps),
                                   value=-lr)
                else:
                    # Standard AdamW update for small tensors
                    exp_avg = state['exp_avg']
                    exp_avg_sq = state['exp_avg_sq']
                    exp_avg.mul_(beta1).add_(grad, alpha=1-beta1)
                    exp_avg_sq.mul_(beta2).addcmul_(
                        grad, grad, value=1-beta2)
                    bc1 = 1 - beta1 ** t
                    bc2 = 1 - beta2 ** t
                    step_size = lr / bc1
                    denom = (exp_avg_sq.sqrt() / (bc2**0.5)).add_(eps)
                    p.data.addcdiv_(exp_avg, denom, value=-step_size)
        
        return loss


# ============================================================
# MULTI-AGENT TRAINING LOOP SKETCH
# ============================================================

class TwoAgentFinancialSystem:
    """
    Two-agent LLM system: Analyst + Critic
    Both trained with LowRankEMAAdamW
    """
    def __init__(self, analyst_model, critic_model, rank=64):
        self.analyst = analyst_model   # 7B LLaMA-3
        self.critic = critic_model     # 7B LLaMA-3
        
        # Shared optimizer with low-rank EMA
        all_params = (list(analyst_model.parameters()) + 
                     list(critic_model.parameters()))
        self.optimizer = LowRankEMAAdamW(
            all_params, lr=1e-4, rank=rank,
            svd_update_freq=1  # ablate: try 5, 10
        )
    
    def training_step(self, financial_text_batch):
        # Agent 1: Generate financial analysis
        analyst_output = self.analyst(financial_text_batch)
        analyst_loss = compute_lm_loss(analyst_output, financial_text_batch)
        
        # Agent 2: Critique the analysis
        critic_input = concat(financial_text_batch, analyst_output.detach())
        critic_output = self.critic(critic_input)
        critic_loss = compute_critique_loss(critic_output)
        
        # Joint loss
        total_loss = analyst_loss + 0.1 * critic_loss
        
        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(
            list(self.analyst.parameters()) + 
            list(self.critic.parameters()), 1.0)
        self.optimizer.step()
        
        return total_loss.item()


# ============================================================
# MEMORY PROFILING HARNESS
# ============================================================

def profile_optimizer_memory(model, optimizer_class, rank=None, 
                              n_steps=100):
    """
    Measure optimizer state memory footprint.
    Returns: dict with memory breakdown in GB
    """
    torch.cuda.reset_peak_memory_stats()
    
    if rank is not None:
        opt = optimizer_class(model.parameters(), lr=1e-4, rank=rank)
    else:
        opt = optimizer_class(model.parameters(), lr=1e-4)
    
    # Measure after optimizer init (states allocated after first step)
    dummy_input = torch.randn(4, 512, device='cuda')
    
    mem_before_step = torch.cuda.memory_allocated() / 1e9
    
    for step in range(n_steps):
        loss = model(dummy_input).sum()
        loss.backward()
        opt.step()
        opt.zero_grad()
    
    mem_after_steps = torch.cuda.memory_allocated() / 1e9
    peak_mem = torch.cuda.max_memory_allocated() / 1e9
    
    # Isolate optimizer state memory
    param_mem = sum(p.numel() * p.element_size() 
                   for p in model.parameters()) / 1e9
    optimizer_state_mem = mem_after_steps - param_mem
    
    return {
        'total_peak_gb': peak_mem,
        'optimizer_state_gb': optimizer_state_mem,
        'param_gb': param_mem,
        'n_steps': n_steps
    }


# ============================================================
# FROBENIUS ERROR DIAGNOSTIC
# ============================================================

def compute_approximation_error(optimizer):
    """
    For each low-rank state, compute ||M - M_r||_F / ||M||_F
    Returns per-layer errors for quality monitoring.
    """
    errors = {}
    for name, state in optimizer.state.items():
        if 'use_lowrank' in state and state['use_lowrank']:
            # Reconstruct approximate moment
            m_approx = optimizer._reconstruct_from_lowrank(
                state['m_U'], state['m_S'], state['m_Vh'],
                state['orig_shape'])
            # We don't have ground truth full-rank moment here
            # Proxy: compare to a full-rank AdamW run (offline analysis)
            # During training: track singular value decay as proxy
            sv_decay = state['m_S'][-1] / state['m_S'][0]  # last/first SV ratio
            errors[str(name)] = sv_decay.item()
    return errors