Low-rank approximation of Adam optimizer momentum matrices, analogous to matrix-interpolatory reduced-order models, will reduce memory overhead without degrading convergence on transformer architectures.

Replacing the full-rank first-moment (m) and second-moment (v) matrices in the Adam optimizer with rank-r approximations (r << d, where d is the parameter dimension) will reduce optimizer state memory by ≥40% while maintaining final validation loss within 2% relative degradation and convergence speed within 10% additional wall-clock time, compared to standard Adam, when training transformer architectures with ≥100M parameters on standard NLP/vision benchmarks.

Minimum Viable Test (MVT): Train GPT-2 Small (117M params) and GPT-2 Medium (345M params) on OpenWebText for 10,000 steps using: (A) Standard Adam, (B) Low-rank Adam with r=8, (C) Low-rank Adam with r=16, (D) Low-rank Adam with r=32. Measure validation perplexity, peak GPU memory, and wall-clock time per step. Full Validation: Extend to LLaMA-7B on C4 dataset for 50,000 steps, plus BERT-Base fine-tuning on GLUE benchmark suite.

import torch
from torch.optim import Optimizer
import torch.nn.functional as F

class LowRankAdam(Optimizer):
    """
    Adam optimizer with low-rank approximation of momentum matrices.
    For weight matrices W in R^{m x n}, maintains:
      m_U in R^{m x r}, m_V in R^{r x n}  (first moment factors)
      v_U in R^{m x r}, v_V in R^{r x n}  (second moment factors)
    For vectors/scalars: standard Adam state.
    """
    
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 rank=16, svd_freq=10, weight_decay=0.01,
                 min_dim_for_lowrank=64, update_strategy='randomized'):
        defaults = dict(lr=lr, betas=betas, eps=eps, rank=rank,
                       svd_freq=svd_freq, weight_decay=weight_decay,
                       min_dim_for_lowrank=min_dim_for_lowrank,
                       update_strategy=update_strategy)
        super().__init__(params, defaults)
    
    def _is_matrix_param(self, p, min_dim):
        """Check if parameter qualifies for low-rank treatment."""
        return (p.dim() == 2 and 
                min(p.shape) >= min_dim and
                max(p.shape) / min(p.shape) <= 100)  # avoid extreme aspect ratios
    
    def _init_low_rank_state(self, state, p, rank):
        """Initialize low-rank factor matrices."""
        m, n = p.shape
        r = min(rank, min(m, n) // 2)  # safety clamp
        state['rank'] = r
        state['step'] = 0
        # First moment factors
        state['m_U'] = torch.zeros(m, r, device=p.device, dtype=p.dtype)
        state['m_V'] = torch.zeros(r, n, device=p.device, dtype=p.dtype)
        # Second moment factors  
        state['v_U'] = torch.zeros(m, r, device=p.device, dtype=p.dtype)
        state['v_V'] = torch.zeros(r, n, device=p.device, dtype=p.dtype)
        state['is_low_rank'] = True
    
    def _randomized_svd(self, G, rank, n_oversampling=5):
        """Randomized SVD for efficient low-rank approximation."""
        m, n = G.shape
        k = rank + n_oversampling
        # Random projection
        Omega = torch.randn(n, k, device=G.device, dtype=G.dtype)
        Y = G @ Omega  # m x k
        Q, _ = torch.linalg.qr(Y)  # m x k orthonormal
        B = Q.T @ G  # k x n
        U_hat, S, Vh = torch.linalg.svd(B, full_matrices=False)
        U = Q @ U_hat[:, :rank]  # m x rank
        S = S[:rank]
        V = Vh[:rank, :]  # rank x n
        return U, S, V
    
    def _update_low_rank_momentum(self, state, grad, beta1, beta2, step):
        """Update low-rank momentum factors."""
        r = state['rank']
        strategy = self.defaults['update_strategy']
        svd_freq = self.defaults['svd_freq']
        
        if strategy == 'randomized' or (step % svd_freq == 0):
            # Recompute SVD of gradient
            U, S, V = self._randomized_svd(grad, r)
            sqrt_S = S.sqrt()
            
            # EMA update on factor matrices
            beta1_t = beta1 ** step
            beta2_t = beta2 ** step
            
            # First moment: m = beta1 * m + (1-beta1) * G_lowrank
            # Represented as m_U @ m_V where G_lowrank = U @ diag(S) @ V
            new_m_U = U * sqrt_S.unsqueeze(0)  # m x r
            new_m_V = V * sqrt_S.unsqueeze(1)  # r x n
            state['m_U'].mul_(beta1).add_(new_m_U, alpha=1-beta1)
            state['m_V'].mul_(beta1).add_(new_m_V, alpha=1-beta1)
            
            # Second moment: v = beta2 * v + (1-beta2) * G^2_lowrank
            # Approximate: (U @ diag(S) @ V)^2 elementwise ~ U @ diag(S^2) @ V
            new_v_U = U * S.unsqueeze(0)  # m x r
            new_v_V = V * S.unsqueeze(1)  # r x n  
            state['v_U'].mul_(beta2).add_(new_v_U, alpha=1-beta2)
            state['v_V'].mul_(beta2).add_(new_v_V, alpha=1-beta2)
        
        # Bias correction
        bias_correction1 = 1 - beta1 ** step
        bias_correction2 = 1 - beta2 ** step
        
        # Reconstruct corrected moments
        m_hat = (state['m_U'] @ state['m_V']) / bias_correction1
        v_hat = (state['v_U'] @ state['v_V']) / bias_correction2
        
        return m_hat, v_hat
    
    @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:
            beta1, beta2 = group['betas']
            lr = group['lr']
            eps = group['eps']
            rank = group['rank']
            min_dim = group['min_dim_for_lowrank']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                
                grad = p.grad
                state = self.state[p]
                
                # Initialize state
                if len(state) == 0:
                    if self._is_matrix_param(p, min_dim):
                        self._init_low_rank_state(state, p, rank)
                    else:
                        # Standard Adam for non-matrix params
                        state['step'] = 0
                        state['exp_avg'] = torch.zeros_like(p)
                        state['exp_avg_sq'] = torch.zeros_like(p)
                        state['is_low_rank'] = False
                
                state['step'] += 1
                step = state['step']
                
                # Weight decay
                if group['weight_decay'] != 0:
                    grad = grad.add(p, alpha=group['weight_decay'])
                
                if state['is_low_rank']:
                    # Low-rank Adam update
                    m_hat, v_hat = self._update_low_rank_momentum(
                        state, grad, beta1, beta2, step)
                    
                    # Adam update step
                    denom = v_hat.abs().sqrt().add_(eps)
                    p.addcdiv_(m_hat, denom, value=-lr)
                    
                else:
                    # Standard Adam update
                    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)
                    
                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2 ** step
                    
                    step_size = lr / bias_correction1
                    denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)
                    p.addcdiv_(exp_avg, denom, value=-step_size)
        
        return loss


# ============================================================
# EXPERIMENT RUNNER
# ============================================================

def run_experiment(model, train_loader, val_loader, optimizer_class, 
                   optimizer_kwargs, n_steps=10000, eval_every=500):
    """
    Standardized experiment runner with memory and timing profiling.
    Returns: dict of metrics
    """
    import time
    
    optimizer = optimizer_class(model.parameters(), **optimizer_kwargs)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=n_steps)
    
    metrics = {
        'train_loss': [], 'val_perplexity': [],
        'peak_memory_mb': [], 'step_time_ms': [],
        'grad_norm': []
    }
    
    model.train()
    step = 0
    
    for batch in train_loader:
        if step >= n_steps:
            break
        
        t0 = time.perf_counter()
        torch.cuda.reset_peak_memory_stats()
        
        # Forward + backward
        loss = model(**batch).loss
        loss.backward()
        
        # Gradient clipping
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
        
        step_time = (time.perf_counter() - t0) * 1000
        peak_mem = torch.cuda.max_memory_allocated() / 1e6
        
        metrics['train_loss'].append(loss.item())
        metrics['step_time_ms'].append(step_time)
        metrics['peak_memory_mb'].append(peak_mem)
        metrics['grad_norm'].append(grad_norm.item())
        
        if step % eval_every == 0:
            val_ppl = evaluate_perplexity(model, val_loader)
            metrics['val_perplexity'].append((step, val_ppl))
            print(f"Step {step}: loss={loss.item():.4f}, "
                  f"val_ppl={val_ppl:.2f}, "
                  f"mem={peak_mem:.0f}MB, "
                  f"time={step_time:.1f}ms")
        
        step += 1
    
    return metrics


# ============================================================
# MEMORY ANALYSIS UTILITY
# ============================================================

def compute_optimizer_state_memory(model, rank):
    """
    Estimate optimizer state memory for standard Adam vs LowRankAdam.
    Returns: (standard_mb, lowrank_mb, reduction_pct)
    """
    standard_params = 0
    lowrank_params = 0
    
    for name, p in model.named_parameters():
        n_elements = p.numel()
        
        if p.dim() == 2 and min(p.shape) >= 64:
            m, n = p.shape
            r = min(rank, min(m, n) // 2)
            # Standard Adam: 2 copies (m, v) of full matrix
            standard_params += 2 * n_elements
            # LowRankAdam: 2 copies of (U, V) factors each
            lowrank_params += 2 * (m * r + r * n)
        else:
            # Both use standard Adam for non-matrix params
            standard_params += 2 * n_elements
            lowrank_params += 2 * n_elements
    
    bytes_per_param = 4  # FP32
    standard_mb = standard_params * bytes_per_param / 1e6
    lowrank_mb = lowrank_params * bytes_per_param / 1e6
    reduction_pct = (1 - lowrank_mb / standard_mb) * 100
    
    return standard_mb, lowrank_mb, reduction_pct


# ============================================================
# MAIN EXPERIMENT CONFIGURATION
# ============================================================

EXPERIMENT_CONFIG = {
    'models': ['gpt2', 'gpt2-medium', 'gpt2-large'],
    'ranks': [4, 8, 16, 32, 64],
    'svd_frequencies': [1, 10, 50, 100],
    'update_strategies': ['randomized', 'periodic', 'online'],
    'seeds': [42, 123, 456],
    'n_steps': 10000,
    'eval_every': 500,
    'optimizer_baseline': {
        'class': torch.optim.AdamW,
        'kwargs': {'lr': 3e-4, 'betas': (0.9, 0.999), 
                   'eps': 1e-8, 'weight_decay': 0.01}
    },
    'optimizer_lowrank': {
        'class': LowRankAdam,
        'kwargs': {'lr': 3e-4, 'betas': (0.9, 0.999),
                   'eps': 1e-8, 'weight_decay': 0.01,
                   'rank': 16, 'svd_freq': 10,
                   'update_strategy': 'randomized'}
    }
}