Low-rank approximations of optimizer states can improve the training of machine learning surrogates in structural optimization.

Applying low-rank approximations (e.g., rank-r SVD truncation or LoRA-style factorization) to optimizer state tensors (momentum buffers, second-moment estimates in Adam/AdamW) during training of ML surrogate models for structural optimization problems reduces memory footprint and/or wall-clock training time while achieving surrogate prediction accuracy (e.g., compliance, stress field, displacement field) within ≤5% relative error compared to full-rank optimizer baselines, measured on at least two distinct structural topology optimization benchmarks.

Minimum Viable Test (MVT): Train a graph neural network (GNN) or CNN-based surrogate on a 2D topology optimization dataset (e.g., 64×64 or 128×128 grid) using (a) standard Adam optimizer and (b) Adam with low-rank approximated second-moment estimates. Compare final surrogate accuracy, peak memory usage, and training time. Run 3 seeds per condition. Total MVT scope: 6 training runs × ~4 GPU-hours each = ~24 GPU-hours.

Full Validation: Extend to 3D structural problems, multiple surrogate architectures (CNN, GNN, FNO), multiple rank values (r ∈ {4, 8, 16, 32, full}), and 5 seeds per condition across 2 benchmark datasets.

# Low-Rank Adam Optimizer Implementation Sketch

import torch
from torch.optim import Optimizer

class LowRankAdam(Optimizer):
    """Adam optimizer with low-rank approximation of second moment."""
    
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 rank=16, svd_freq=100):
        self.rank = rank
        self.svd_freq = svd_freq  # recompute SVD every K steps
        defaults = dict(lr=lr, betas=betas, eps=eps)
        super().__init__(params, defaults)
    
    def step(self, closure=None):
        for group in self.param_groups:
            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['exp_avg'] = torch.zeros_like(p.data)  # m_t (full rank)
                    # Initialize low-rank second moment: v_t ≈ U_r @ diag(s_r) @ V_r^T
                    if grad.dim() >= 2:
                        m, n = grad.view(grad.shape[0], -1).shape
                        r = min(self.rank, m, n)
                        state['U'] = torch.zeros(m, r, device=p.device)
                        state['S'] = torch.zeros(r, device=p.device)
                        state['V'] = torch.zeros(n, r, device=p.device)
                        state['use_lowrank'] = True
                    else:
                        state['exp_avg_sq'] = torch.zeros_like(p.data)
                        state['use_lowrank'] = False
                
                state['step'] += 1
                beta1, beta2 = group['betas']
                
                # Update first moment (full rank - standard)
                exp_avg = state['exp_avg']
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                
                if state['use_lowrank']:
                    # Reshape gradient for matrix operations
                    g_mat = grad.view(grad.shape[0], -1)
                    
                    # Update second moment in low-rank form
                    # Reconstruct current v_t approximation
                    U, S, V = state['U'], state['S'], state['V']
                    v_approx = (U * S.unsqueeze(0)) @ V.T  # m×n
                    
                    # EMA update: v_t = beta2 * v_{t-1} + (1-beta2) * g^2
                    v_new = beta2 * v_approx + (1 - beta2) * (g_mat ** 2)
                    
                    # Recompute SVD periodically
                    if state['step'] % self.svd_freq == 0:
                        # Randomized SVD for efficiency
                        U_new, S_new, Vh_new = torch.linalg.svd(
                            v_new, full_matrices=False
                        )
                        r = min(self.rank, S_new.shape[0])
                        state['U'] = U_new[:, :r].contiguous()
                        state['S'] = S_new[:r].contiguous()
                        state['V'] = Vh_new[:r, :].T.contiguous()
                    
                    # Reconstruct v_t for parameter update
                    v_approx_current = (state['U'] * state['S'].unsqueeze(0)) @ state['V'].T
                    v_approx_current = v_approx_current.view_as(grad)
                    
                    # Bias correction
                    bias_correction1 = 1 - beta1 ** state['step']
                    bias_correction2 = 1 - beta2 ** state['step']
                    
                    # Adam update
                    step_size = group['lr'] / bias_correction1
                    denom = (v_approx_current.sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
                    p.data.addcdiv_(exp_avg, denom, value=-step_size)
                
                else:
                    # Standard Adam for 1D parameters (biases, etc.)
                    exp_avg_sq = state['exp_avg_sq']
                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                    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.data.addcdiv_(exp_avg, denom, value=-step_size)

# Surrogate Model Architecture (U-Net for TopOpt)
class TopOptSurrogate(torch.nn.Module):
    def __init__(self, in_channels=3, out_channels=1, base_channels=64):
        super().__init__()
        # Encoder
        self.enc1 = ConvBlock(in_channels, base_channels)
        self.enc2 = ConvBlock(base_channels, base_channels*2)
        self.enc3 = ConvBlock(base_channels*2, base_channels*4)
        # Bottleneck
        self.bottleneck = ConvBlock(base_channels*4, base_channels*8)
        # Decoder with skip connections
        self.dec3 = ConvBlock(base_channels*8 + base_channels*4, base_channels*4)
        self.dec2 = ConvBlock(base_channels*4 + base_channels*2, base_channels*2)
        self.dec1 = ConvBlock(base_channels*2 + base_channels, base_channels)
        self.output = torch.nn.Conv2d(base_channels, out_channels, 1)
    
    def forward(self, x):
        # Standard U-Net forward pass
        e1 = self.enc1(x)
        e2 = self.enc2(F.max_pool2d(e1, 2))
        e3 = self.enc3(F.max_pool2d(e2, 2))
        b = self.bottleneck(F.max_pool2d(e3, 2))
        d3 = self.dec3(torch.cat([F.interpolate(b, scale_factor=2), e3], dim=1))
        d2 = self.dec2(torch.cat([F.interpolate(d3, scale_factor=2), e2], dim=1))
        d1 = self.dec1(torch.cat([F.interpolate(d2, scale_factor=2), e1], dim=1))
        return self.output(d1)

# Training Loop
def train_surrogate(model, dataset, optimizer_type='lowrank_adam', rank=16, 
                    epochs=200, seed=42):
    torch.manual_seed(seed)
    
    if optimizer_type == 'lowrank_adam':
        optimizer = LowRankAdam(model.parameters(), lr=1e-3, rank=rank)
    else:
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    
    metrics = {'train_loss': [], 'val_loss': [], 'memory_mb': [], 'time_per_epoch': []}
    
    for epoch in range(epochs):
        t0 = time.time()
        model.train()
        for batch in dataset.train_loader:
            x, y = batch
            pred = model(x)
            loss = F.mse_loss(pred, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        metrics['memory_mb'].append(torch.cuda.max_memory_allocated() / 1e6)
        metrics['time_per_epoch'].append(time.time() - t0)
        
        # Validation
        val_loss = evaluate(model, dataset.val_loader)
        metrics['val_loss'].append(val_loss)
        
        # Early abort checkpoint
        if epoch == 20 and val_loss > 10 * metrics['val_loss'][0]:
            print(f"ABORT: Loss diverging at epoch {epoch}")
            return metrics, 'diverged'
    
    return metrics, 'completed'

# Experiment Runner
def run_experiment():
    results = {}
    for rank in [4, 8, 16, 32, 'full']:
        for seed in [42, 123, 456, 789, 1024]:
            opt_type = 'lowrank_adam' if rank != 'full' else 'adam'
            r = rank if rank != 'full' else 512
            metrics, status = train_surrogate(
                model=TopOptSurrogate(),
                dataset=TopOptDataset('2d_64x64'),
                optimizer_type=opt_type,
                rank=r, seed=seed
            )
            results[f'rank_{rank}_seed_{seed}'] = {'metrics': metrics, 'status': status}
    return results