Low-rank momentum approximations can reduce the memory cost of training surrogate models used in amortized structural optimization without degrading solution quality.

Applying low-rank approximations (rank r << full rank) to the first and second momentum buffers in adaptive optimizers (e.g., Adam, AdamW) during training of neural surrogate models for amortized structural optimization reduces peak GPU memory consumption by ≥20% relative to full-rank momentum storage, while maintaining solution quality metrics (compliance, stress, volume fraction error) within ±5% of the full-rank baseline across standard benchmark topologies.

Minimum viable test: Train a U-Net or graph neural network surrogate for 2D topology optimization (e.g., 64×64 or 128×128 grid) using Adam optimizer with full-rank momentum (baseline) versus low-rank momentum approximations at ranks r ∈ {4, 8, 16, 32}. Evaluate on 500 held-out boundary condition/load configurations. Measure peak GPU memory, final compliance error, volume fraction error, and training convergence curves. Run 3 independent seeds per condition.

# Low-Rank Adam Optimizer for Surrogate Training
# Core data structure: factored momentum buffers

import torch
from torch.optim import Optimizer

class LowRankAdam(Optimizer):
    """
    Adam optimizer with low-rank momentum approximation.
    Momentum buffers stored as M ≈ U @ V.T
    where U: (m, r), V: (n, r) instead of full (m, n) matrix.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
                 eps=1e-8, rank=16, update_freq=10):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                       rank=rank, update_freq=update_freq)
        super().__init__(params, defaults)

    def step(self, closure=None):
        for group in self.param_groups:
            rank = group['rank']
            update_freq = group['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
                    # Full first moment (can also be low-rank)
                    state['exp_avg'] = torch.zeros_like(p.data)
                    # Low-rank second moment: diagonal + low-rank
                    # For 2D weight matrices: U (m×r), V (n×r)
                    if grad.dim() == 2:
                        m, n = grad.shape
                        r = min(rank, min(m, n))
                        state['exp_avg_sq_U'] = torch.randn(m, r,
                            device=p.device) * 0.01
                        state['exp_avg_sq_V'] = torch.randn(n, r,
                            device=p.device) * 0.01
                        state['exp_avg_sq_diag'] = torch.zeros(
                            min(m,n), device=p.device)
                        state['is_2d'] = True
                    else:
                        # Fallback to full for 1D (bias, BN params)
                        state['exp_avg_sq'] = torch.zeros_like(p.data)
                        state['is_2d'] = False

                state['step'] += 1
                beta1, beta2 = group['betas']

                # Update first moment (standard)
                exp_avg = state['exp_avg']
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)

                if state['is_2d']:
                    # Update low-rank second moment via Oja's rule
                    # or periodic SVD truncation
                    U = state['exp_avg_sq_U']
                    V = state['exp_avg_sq_V']

                    if state['step'] % update_freq == 0:
                        # Reconstruct approximate second moment
                        approx_sq = U @ V.T
                        # Add current grad^2 contribution
                        approx_sq.mul_(beta2).add_(
                            grad ** 2, alpha=1 - beta2)
                        # Re-factorize via truncated SVD
                        try:
                            # Use randomized SVD for efficiency
                            U_new, S_new, Vh_new = torch.linalg.svd(
                                approx_sq, full_matrices=False)
                            r = U.shape[1]
                            state['exp_avg_sq_U'] = (
                                U_new[:, :r] *
                                S_new[:r].sqrt().unsqueeze(0))
                            state['exp_avg_sq_V'] = (
                                Vh_new[:r, :].T *
                                S_new[:r].sqrt().unsqueeze(0))
                        except Exception:
                            # Fallback: Oja's rule update
                            grad_outer_U = (grad @ V) * (1 - beta2)
                            U.mul_(beta2).add_(grad_outer_U)
                            grad_outer_V = (grad.T @ U) * (1 - beta2)
                            V.mul_(beta2).add_(grad_outer_V)

                    # Reconstruct for parameter update
                    exp_avg_sq_approx = (U @ V.T).abs().clamp(min=1e-30)

                    # Bias correction
                    bias_correction1 = 1 - beta1 ** state['step']
                    bias_correction2 = 1 - beta2 ** state['step']

                    step_size = (group['lr'] *
                                (bias_correction2 ** 0.5) /
                                bias_correction1)

                    denom = exp_avg_sq_approx.sqrt().add_(group['eps'])
                    p.data.addcdiv_(exp_avg, denom, value=-step_size)

                else:
                    # Standard Adam for 1D params
                    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_correction2 ** 0.5) /
                                bias_correction1)
                    denom = exp_avg_sq.sqrt().add_(group['eps'])
                    p.data.addcdiv_(exp_avg, denom, value=-step_size)

# ─── Surrogate Model (U-Net for topology optimization) ───────────────────────

class TopOptSurrogate(torch.nn.Module):
    """
    U-Net surrogate: maps (boundary_conditions, loads, vf) -> density_field
    Input: (B, C_in, H, W) where C_in = load channels + BC mask
    Output: (B, 1, H, W) density field in [0, 1]
    """
    def __init__(self, in_channels=4, base_features=64):
        super().__init__()
        # Encoder
        self.enc1 = self._block(in_channels, base_features)
        self.enc2 = self._block(base_features, base_features * 2)
        self.enc3 = self._block(base_features * 2, base_features * 4)
        # Bottleneck
        self.bottleneck = self._block(base_features * 4, base_features * 8)
        # Decoder
        self.dec3 = self._block(base_features * 8 + base_features * 4,
                                base_features * 4)
        self.dec2 = self._block(base_features * 4 + base_features * 2,
                                base_features * 2)
        self.dec1 = self._block(base_features * 2 + base_features,
                                base_features)
        self.final = torch.nn.Conv2d(base_features, 1, 1)
        self.pool = torch.nn.MaxPool2d(2)
        self.up = torch.nn.Upsample(scale_factor=2, mode='bilinear',
                                    align_corners=True)

    def _block(self, in_c, out_c):
        return torch.nn.Sequential(
            torch.nn.Conv2d(in_c, out_c, 3, padding=1),
            torch.nn.BatchNorm2d(out_c),
            torch.nn.ReLU(inplace=True),
            torch.nn.Conv2d(out_c, out_c, 3, padding=1),
            torch.nn.BatchNorm2d(out_c),
            torch.nn.ReLU(inplace=True)
        )

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        b = self.bottleneck(self.pool(e3))
        d3 = self.dec3(torch.cat([self.up(b), e3], dim=1))
        d2 = self.dec2(torch.cat([self.up(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up(d2), e1], dim=1))
        return torch.sigmoid(self.final(d1))

# ─── Training Loop with Memory Profiling ─────────────────────────────────────

def train_and_profile(rank=None, n_epochs=100, seed=42):
    torch.manual_seed(seed)
    model = TopOptSurrogate().cuda()

    if rank is None:
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        condition = 'full_rank'
    else:
        optimizer = LowRankAdam(model.parameters(), lr=1e-3, rank=rank)
        condition = f'rank_{rank}'

    memory_log = []
    loss_log = []

    for epoch in range(n_epochs):
        for batch in dataloader:  # dataloader yields (inputs, targets)
            inputs, targets = batch
            inputs, targets = inputs.cuda(), targets.cuda()

            optimizer.zero_grad()
            preds = model(inputs)
            loss = torch.nn.functional.binary_cross_entropy(preds, targets)
            loss.backward()
            optimizer.step()

            # Profile memory
            mem = torch.cuda.memory_stats()
            peak_mem_gb = mem['reserved_bytes.all.peak'] / 1e9
            memory_log.append(peak_mem_gb)

        loss_log.append(loss.item())
        torch.cuda.reset_peak_memory_stats()

    return {
        'condition': condition,
        'peak_