Fira (Full-Rank Training under Low-Rank Optimizer Constraint)

Given full-rank gradient G_t in R^{m x n} (m <= n) and rank-r SVD projection matrix P_t in R^{m x r} updated every T steps, Fira computes the low-rank projected gradient R_t = P_t^T G_t in R^{r x n} and its adaptive optimizer correction psi_t(R_t) using low-rank states M_t, V_t in R^{r x n}, defines the matrix-level scaling factor phi_t(R_t) = ||psi_t(R_t)|| / ||R_t|| (or column-wise factors phi_t(R_t)_i = ||psi_t(R_{t,:,i})|| / (||R_{t,:,i}|| + epsilon) for i = 1, ..., n), and rescales the orthogonal residual gradient S_t = phi_t(R_t) * (G_t - P_t R_t) so the parameter update W_{t+1} = W_t - eta * (P_t psi_t(R_t) + S_t) preserves full-rank gradient information with only mr + 2nr + 1 optimizer states.

To prevent training loss spikes triggered by periodic projection-subspace switching (t mod T = 0) and steep curvature along unsmoothed residual directions, Fira applies a relative norm-growth limiter to the scaled residual gradient S_t = phi_t(R_t)(G_t - P_t R_t): whenever ||S_t|| / ||S_{t-1}|| > gamma (default gamma = 1.01), S_t is rescaled to (S_t / ||S_t||) * gamma * ||S_{t-1}||, converting abrupt gradient norm jumps into bounded geometric growth while preserving direction.

Under bounded component-wise adaptive scaling c_min <= |psi_i(g_i)/g_i| <= c_max, Theorem 1 proves that the squared scaling-factor approximation error satisfies kappa(r) = |phi_t^2(G_t) - phi_t^2(R_t)| <= (c_max^2 - c_min^2) * (1 - ||G_r||^2 / ||G||^2), and Theorem 2 proves under early-stage Gaussian gradients and inverse-chi-squared second moments (rho > 4) that Dirichlet(1/2, ..., 1/2) weighting yields Var[phi^2] = (3 / (n + 2)) * Var[psi^2] <= Var[psi^2].

In C4 pre-training of LLaMA-60M/130M/350M/1B/7B and commonsense-reasoning fine-tuning of LLaMA-7B (BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC-e, ARC-c, OBQA), Fira reduces optimizer-state memory by 61.1% on LLaMA-1B and uses 8x smaller optimizer state memory (rank 64 vs 512) on LLaMA-7B while achieving lower validation perplexity and higher average accuracy (76.9% vs 62.7% for GaLore and 74.7% for LoRA).

Unlike LoRA, which stores only low-rank adapter gradients (mr + nr), Fira computes full-rank weight gradients (mn) during backpropagation and incurs a slight throughput overhead relative to GaLore and LoRA from column-wise norm-based scaling and norm-growth limiting.