Low-rank approximation of optimizer states can reduce memory overhead in agent-based economic simulations.

Applying low-rank approximation (rank r << d) to the second-moment (variance) and first-moment (momentum) optimizer state tensors in gradient-based learning agents within agent-based economic simulations (ABES) reduces peak GPU/CPU memory consumption by ≥20% relative to full-rank optimizer states, while maintaining simulation convergence quality (policy loss degradation ≤5%) and agent behavioral fidelity (Gini coefficient deviation ≤0.02, market clearing error ≤1%) across simulations with ≥1,000 heterogeneous agents over ≥500 time steps.

Minimum Viable Test (MVT): Single-commodity exchange economy with N=1,000 Adam-optimized agents, comparing full-rank vs. rank-r (r=4,8,16) optimizer states across 500 training episodes, measuring peak memory, convergence speed, and behavioral metrics. Full Validation: Scale to N=10,000 agents, three economic model types, five optimizer variants, ablation over rank values r ∈ {2,4,8,16,32}, five random seeds each, with profiling of memory and compute at each checkpoint.

# LowRankAdam: Core Implementation Sketch
import torch
from torch.optim import Optimizer

class LowRankAdam(Optimizer):
    """
    Adam optimizer with low-rank factored moment states.
    Stores m_t ≈ U_m @ V_m (rank-r) instead of full m_t.
    Stores v_t ≈ U_v @ V_v (rank-r) instead of full v_t.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
                 eps=1e-8, rank=8, update_freq=1):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                       rank=rank, update_freq=update_freq)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure else None

        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]
                original_shape = grad.shape

                # Reshape to 2D matrix for low-rank ops
                # G: (m, n) where m*n = numel(grad)
                m = grad.shape[0]
                n = grad.numel() // m
                G = grad.reshape(m, n)

                # Initialize factored states
                if len(state) == 0:
                    state['step'] = 0
                    # First moment factors: m_t ≈ U1 @ V1
                    state['U1'] = torch.zeros(m, rank, device=p.device)
                    state['V1'] = torch.zeros(rank, n, device=p.device)
                    # Second moment factors: v_t ≈ U2 @ V2
                    state['U2'] = torch.zeros(m, rank, device=p.device)
                    state['V2'] = torch.zeros(rank, n, device=p.device)

                state['step'] += 1
                U1, V1 = state['U1'], state['V1']
                U2, V2 = state['U2'], state['V2']

                # Reconstruct current moments
                m_prev = U1 @ V1  # (m, n)
                v_prev = U2 @ V2  # (m, n)

                # Update moments (standard Adam equations)
                m_new = beta1 * m_prev + (1 - beta1) * G
                v_new = beta2 * v_prev + (1 - beta2) * G.pow(2)

                # Re-factorize updated moments via randomized SVD
                # torch.svd_lowrank: O(m*n*rank) vs O(m*n*min(m,n))
                U1_new, S1, V1h_new = torch.svd_lowrank(
                    m_new, q=rank, niter=2
                )
                U2_new, S2, V2h_new = torch.svd_lowrank(
                    v_new, q=rank, niter=2
                )

                # Store factored states (memory: 2*rank*(m+n) vs m*n)
                state['U1'] = U1_new * S1.unsqueeze(0)  # absorb S into U
                state['V1'] = V1h_new.T
                state['U2'] = U2_new * S2.unsqueeze(0)
                state['V2'] = V2h_new.T

                # Bias correction
                bc1 = 1 - beta1 ** state['step']
                bc2 = 1 - beta2 ** state['step']
                m_hat = (state['U1'] @ state['V1']) / bc1
                v_hat = (state['U2'] @ state['V2']) / bc2

                # Parameter update
                update = m_hat / (v_hat.sqrt() + group['eps'])
                p.add_(update.reshape(original_shape), alpha=-group['lr'])

        return loss


# Memory comparison utility
def compare_memory_usage(model, N_agents, rank):
    """
    Estimate memory savings ratio for given model and rank.
    Returns: (full_rank_bytes, low_rank_bytes, reduction_ratio)
    """
    full_rank_bytes = 0
    low_rank_bytes = 0

    for name, param in model.named_parameters():
        if param.requires_grad:
            numel = param.numel()
            m = param.shape[0]
            n = numel // m

            # Full Adam: 2 state tensors (m_t, v_t) of size numel
            full_rank_bytes += 2 * numel * 4  # float32

            # Low-rank Adam: 4 factor matrices
            # U1,U2: (m, rank); V1,V2: (rank, n)
            low_rank_bytes += 2 * (m * rank + rank * n) * 4

    # Scale by N_agents
    full_total = full_rank_bytes * N_agents
    low_total = low_rank_bytes * N_agents
    reduction = (full_total - low_total) / full_total

    return full_total, low_total, reduction


# Simulation integration example (AI Economist style)
class EconomicSimulation:
    def __init__(self, N_agents, policy_net_width, rank):
        self.agents = [PolicyNet(width=policy_net_width)
                      for _ in range(N_agents)]
        self.optimizers = [
            LowRankAdam(agent.parameters(), lr=1e-3, rank=rank)
            for agent in self.agents
        ]

    def training_step(self, observations, rewards):
        total_loss = 0
        for i, (agent, opt) in enumerate(
                zip(self.agents, self.optimizers)):
            opt.zero_grad()
            loss = agent.compute_loss(observations[i], rewards[i])
            loss.backward()
            opt.step()
            total_loss += loss.item()
        return total_loss / len(self.agents)


# Evaluation metrics
def compute_behavioral_metrics(wealth_distribution, prices, demands, supplies):
    """Compute Gini, market clearing error, price volatility."""
    # Gini coefficient
    w = torch.sort(wealth_distribution).values
    n = len(w)
    gini = (2 * torch.arange(1, n+1, dtype=torch.float) * w).sum()
    gini = gini / (n * w.sum()) - (n + 1) / n

    # Market clearing error
    mc_error = torch.abs(supplies - demands) / (demands + 1e-8)
    mc_error_mean = mc_error.mean().item()

    # Price volatility
    price_vol = prices.std() / prices.mean()

    return {
        'gini': gini.item(),
        'market_clearing_error': mc_error_mean,
        'price_volatility': price_vol.item()
    }