Low-rank approximation of optimizer states can reduce memory overhead in mRNA sequence design optimization.

Applying low-rank approximation (rank r << d) to the first and second moment optimizer states (e.g., Adam's m_t and v_t) during gradient-based mRNA sequence design optimization reduces peak memory consumption by ≥30% relative to full-rank optimizer states, while maintaining sequence optimization quality (measured by predicted MFE, CAI, or codon adaptation score) within 5% of the full-rank baseline across at least 3 distinct mRNA target sequences of length ≥500 nucleotides.

Minimum viable test: Implement Adam optimizer with low-rank state approximation (GaLore-style or LoRA-adapted state compression) for a differentiable mRNA sequence optimization pipeline. Compare against full-rank Adam baseline on 3 mRNA targets (spike protein, GFP, luciferase) across sequence lengths 500–1000 nt. Measure peak GPU memory, optimization trajectory (loss vs. steps), and final sequence quality metrics. Run 5 random seeds per condition. Total: 6 conditions × 5 seeds = 30 runs.

# Low-Rank Adam for mRNA Sequence Optimization
# Architecture Overview

import torch
import torch.nn as nn
from torch.optim import Optimizer
import numpy as np

# === 1. Sequence Parameterization ===
class mRNASequenceParam(nn.Module):
    def __init__(self, seq_length: int, init_sequence: str = None):
        super().__init__()
        # Shape: (L, 4) for A/U/G/C at each position
        self.logits = nn.Parameter(torch.randn(seq_length, 4) * 0.1)
        if init_sequence:
            self._init_from_sequence(init_sequence)
    
    def forward(self):
        # Gumbel-softmax for differentiable discrete sampling
        return torch.nn.functional.gumbel_softmax(
            self.logits, tau=1.0, hard=False, dim=-1
        )  # Shape: (L, 4)
    
    def decode(self):
        # Argmax decoding for final sequence
        indices = self.logits.argmax(dim=-1)
        return ''.join(['AUGC'[i] for i in indices.tolist()])

# === 2. Low-Rank Adam Optimizer ===
class LowRankAdam(Optimizer):
    def __init__(self, params, lr=1e-2, rank=8, 
                 betas=(0.9, 0.999), eps=1e-8):
        defaults = dict(lr=lr, rank=rank, betas=betas, eps=eps)
        super().__init__(params, defaults)
    
    def step(self, closure=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.data  # Shape: (L, 4)
                state = self.state[p]
                
                # Initialize low-rank state
                if len(state) == 0:
                    state['step'] = 0
                    L, d = grad.shape
                    r = min(rank, min(L, d))
                    # Low-rank first moment: U (L×r), V (d×r)
                    state['m_U'] = torch.zeros(L, r, device=grad.device)
                    state['m_V'] = torch.zeros(d, r, device=grad.device)
                    # Scalar second moment (Adafactor-style)
                    state['v_row'] = torch.zeros(L, device=grad.device)
                    state['v_col'] = torch.zeros(d, device=grad.device)
                
                state['step'] += 1
                t = state['step']
                
                # Project gradient to low-rank subspace via online SVD
                U, S, Vh = torch.linalg.svd(grad, full_matrices=False)
                r = group['rank']
                U_r = U[:, :r]      # (L, r)
                S_r = S[:r]         # (r,)
                V_r = Vh[:r, :].T   # (d, r)
                
                # Low-rank gradient approximation
                grad_lr = U_r @ torch.diag(S_r) @ V_r.T  # (L, d)
                
                # Update low-rank first moment
                state['m_U'] = beta1 * state['m_U'] + (1-beta1) * U_r
                state['m_V'] = beta1 * state['m_V'] + (1-beta1) * V_r
                m_hat = state['m_U'] @ state['m_V'].T / (1 - beta1**t)
                
                # Update factored second moment (row/col factors)
                state['v_row'] = beta2 * state['v_row'] + \
                    (1-beta2) * (grad_lr**2).mean(dim=1)
                state['v_col'] = beta2 * state['v_col'] + \
                    (1-beta2) * (grad_lr**2).mean(dim=0)
                v_hat = torch.outer(state['v_row'], state['v_col'])
                v_hat = v_hat / (1 - beta2**t)
                
                # Parameter update
                p.data -= group['lr'] * m_hat / (v_hat.sqrt() + group['eps'])

# === 3. Optimization Objective ===
class mRNAObjective:
    def __init__(self, target_protein: str, codon_table: dict):
        self.codon_table = codon_table
        self.target_protein = target_protein
    
    def cai_loss(self, seq_probs: torch.Tensor) -> torch.Tensor:
        # Differentiable CAI approximation using codon frequency weights
        # seq_probs: (L, 4) soft nucleotide probabilities
        # Returns scalar loss (1 - CAI_approx)
        # [Implementation: weighted sum over codon triplets]
        pass
    
    def gc_penalty(self, seq_probs: torch.Tensor) -> torch.Tensor:
        gc_content = seq_probs[:, [1, 2]].sum()  # G+C channels
        gc_frac = gc_content / seq_probs.sum()
        return (gc_frac - 0.5).pow(2)  # Target 50% GC
    
    def total_loss(self, seq_probs, alpha=1.0, beta=0.5, gamma=0.2):
        return (alpha * self.cai_loss(seq_probs) + 
                gamma * self.gc_penalty(seq_probs))

# === 4. Main Optimization Loop ===
def run_optimization(target_seq: str, rank: int, n_steps: int = 1000,
                     seed: int = 42):
    torch.manual_seed(seed)
    
    model = mRNASequenceParam(seq_length=len(target_seq))
    optimizer = LowRankAdam(model.parameters(), lr=1e-2, rank=rank)
    objective = mRNAObjective(target_seq, codon_table=load_human_codon_table())
    
    memory_log = []
    loss_log = []
    
    for step in range(n_steps):
        optimizer.zero_grad()
        seq_probs = model()
        loss = objective.total_loss(seq_probs)
        loss.backward()
        optimizer.step()
        
        # Log memory every 10 steps
        if step % 10 == 0:
            mem = torch.cuda.max_memory_allocated() / 1e9  # GB
            memory_log.append((step, mem))
            loss_log.append((step, loss.item()))
        
        # Compute MFE every 50 steps (expensive subprocess)
        if step % 50 == 0:
            seq_str = model.decode()
            mfe = compute_rnafold_mfe(seq_str)  # subprocess call
    
    return {
        'final_sequence': model.decode(),
        'peak_