Low-rank approximation of optimizer momentum states (as in Taming Momentum) can be applied to reduce memory overhead in amortized optimization surrogates without degrading solution quality.

Applying low-rank approximation (rank r << d) to the momentum buffer states of optimizers used within amortized optimization surrogates (e.g., learned optimizers, meta-learned initializations, or neural surrogate solvers) reduces peak memory consumption by at least 30% relative to full-rank momentum storage, while maintaining solution quality within 5% relative degradation on standard benchmark tasks, measured by final objective value or constraint satisfaction rate.

Minimum viable test: Train an amortized optimization surrogate (e.g., a learned optimizer or amortized variational inference network) on a standard benchmark task distribution, comparing full-rank Adam momentum against low-rank approximated momentum (SVD-truncated or factored representation) at ranks r ∈ {4, 8, 16, 32}. Measure peak GPU memory, wall-clock time, and solution quality across 5 seeds per configuration. Two domains tested: (1) combinatorial/continuous optimization surrogates (CS domain) and (2) physics-informed neural network surrogates (Physics domain).

# Low-Rank Momentum Adam Optimizer for Amortized Optimization Surrogates
# Core implementation sketch

import torch
import torch.nn as nn
from typing import Optional, Tuple

class LowRankMomentumAdam(torch.optim.Optimizer):
    """
    Adam optimizer with low-rank approximation of momentum states.
    Applies truncated SVD to momentum buffers m1 and m2 every K steps.
    """
    
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 rank: int = 16, truncate_every: int = 10,
                 compress_m2: bool = True):
        defaults = dict(lr=lr, betas=betas, eps=eps, rank=rank,
                       truncate_every=truncate_every, compress_m2=compress_m2)
        super().__init__(params, defaults)
        self.step_count = 0
    
    def _low_rank_approx(self, tensor: torch.Tensor, rank: int) -> torch.Tensor:
        """Truncated SVD approximation of a 2D tensor."""
        if tensor.dim() == 1:
            return tensor  # Skip 1D tensors (biases)
        if tensor.dim() > 2:
            orig_shape = tensor.shape
            tensor = tensor.view(tensor.shape[0], -1)
            approx = self._low_rank_approx_2d(tensor, rank)
            return approx.view(orig_shape)
        return self._low_rank_approx_2d(tensor, rank)
    
    def _low_rank_approx_2d(self, tensor: torch.Tensor, rank: int) -> torch.Tensor:
        """Core SVD truncation for 2D tensors."""
        effective_rank = min(rank, min(tensor.shape) - 1)
        try:
            U, S, Vh = torch.linalg.svd(tensor, full_matrices=False)
            # Truncate to effective_rank
            U_r = U[:, :effective_rank]
            S_r = S[:effective_rank]
            Vh_r = Vh[:effective_rank, :]
            return (U_r * S_r.unsqueeze(0)) @ Vh_r
        except torch.linalg.LinAlgError:
            return tensor  # Fallback to full rank on SVD failure
    
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        
        self.step_count += 1
        should_truncate = (self.step_count % self.defaults['truncate_every'] == 0)
        
        for group in self.param_groups:
            rank = group['rank']
            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
                    state['exp_avg'] = torch.zeros_like(p)      # m1
                    state['exp_avg_sq'] = torch.zeros_like(p)   # m2
                
                state['step'] += 1
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                
                # Standard Adam momentum updates
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                
                # Apply low-rank truncation periodically
                if should_truncate and p.dim() >= 2:
                    state['exp_avg'] = self._low_rank_approx(exp_avg, rank)
                    if group['compress_m2']:
                        state['exp_avg_sq'] = self._low_rank_approx(
                            exp_avg_sq, rank)
                
                # Bias correction
                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']
                
                step_size = group['lr'] / bias_correction1
                denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(
                    group['eps'])
                
                p.addcdiv_(exp_avg, denom, value=-step_size)
        
        return loss


# Amortized Surrogate Training Loop
class AmortizedSurrogate(nn.Module):
    """Example surrogate network for optimization problems."""
    def __init__(self, input_dim: int, hidden_dim: int = 512, output_dim: int = 1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
    
    def forward(self, x):
        return self.net(x)


def train_amortized_surrogate(
    surrogate: AmortizedSurrogate,
    problem_distribution,  # Iterator over (context, target) pairs
    rank: int = 16,
    n_steps: int = 2000,
    truncate_every: int = 10,
    use_low_rank: bool = True
):
    """
    Training loop comparing full-rank vs low-rank momentum.
    """
    if use_low_rank:
        optimizer = LowRankMomentumAdam(
            surrogate.parameters(), lr=1e-3,
            rank=rank, truncate_every=truncate_every
        )
    else:
        optimizer = torch.optim.Adam(surrogate.parameters(), lr=1e-3)
    
    metrics = {'loss': [], 'memory_mb': [], 'step_time': []}
    
    for step in range(n_steps):
        context, target = next(problem_distribution)
        
        optimizer.zero_grad()
        pred = surrogate(context)
        loss = nn.MSELoss()(pred, target)
        loss.backward()
        optimizer.step()
        
        # Memory tracking
        if step % 100 == 0:
            mem = torch.cuda.max_memory_allocated() / 1e6  # MB
            metrics['loss'].append(loss.item())
            metrics['memory_mb'].append(mem)
            print(f"Step {step}: Loss={loss.item():.4f}, Memory={mem:.1f}MB")
    
    return metrics


# Evaluation: Rank structure analysis
def analyze_momentum_rank(optimizer: LowRankMomentumAdam, top_k: int = 32):
    """Analyze singular value decay of momentum states."""
    variance_explained = {}
    for i, group in enumerate(optimizer.param_groups):
        for j, p in enumerate(group['params']):
            if p.dim() < 2:
                continue
            state = optimizer.state[p]
            if 'exp_avg' not in state:
                continue
            m1 = state['exp_avg'].view(p.shape[0], -1)
            _, S, _ = torch.linalg.svd(m1, full_matrices=False)
            total_var = (S ** 2).sum()
            top_k_var = (S[:top_k] ** 2).sum()
            variance_explained[f'param_{i}_{j}'] = (top_k_var / total_var).item()
    return variance_explained


# Experiment runner
def run_comparison_experiment(benchmark_name: str, ranks=[4, 8, 16, 32], n_seeds=5):
    results = {}
    for rank in ranks:
        results[rank] = {'quality': [], 'memory': [], 'time': []}
        for seed in range(n_seeds):
            torch.manual_seed(seed)
            surrogate = AmortizedSurrogate(input_dim=64, hidden_dim=512)
            surrogate.cuda()
            # ... load benchmark-specific problem distribution ...
            metrics = train_amortized_surrogate(
                surrogate, problem_distribution=None,  # inject benchmark
                rank=rank, use_low_rank=True
            )
            results[rank]['memory'].append(max(metrics['memory_mb']))
            results[rank]['quality'].append(metrics['loss'][-1])
    return results