Phase-aware Oja preconditioner

Integrate this into the optimizer's second-moment path for one or more large parameter blocks, such as the query or key projection of a Transformer layer or a convolutional channel block. Flatten each minibatch gradient into g_k in R^d, maintain an Oja vector v_k, and use it to identify the dominant gradient-covariance direction without storing a d by d matrix. For every minibatch: first compute g_k; update v by v <- v + (delta/d) (v^T g_k) g_k; normalize v by v <- v / (||v||_2 + 1e-8); and increment the sample counter n. Estimate the top Rayleigh quotient with lambda_hat_raw = (v^T g_k)^2 and smooth it using an exponential moving average. Estimate the isotropic residual variance with sigma_hat_squared = EMA((||g_k||_2^2 - (v^T g_k)^2) / max(d - 1, 1)), then set theta_hat_squared = max(lambda_hat - sigma_hat_squared, 0). If 2 theta_hat_squared is larger than delta, compute n_crit = d log(d) / [delta (2 theta_hat_squared - delta)]; otherwise disable the spectral correction. Define alpha_n = sigmoid((n - n_crit) / (tau n_crit)), with tau between 0.1 and 0.5. Replace the raw gradient by g'_n = g_n - alpha_n (g_n^T v_n) v_n / (lambda_hat + epsilon), then pass g'_n to AdamW, or apply this correction to the AdamW update itself. Clip lambda_hat and theta_hat_squared to robust percentile ranges because minibatch gradients are non-Gaussian. The paper supplies the Oja recursion and phase-threshold formula; all covariance, spike, and overlap quantities are estimated empirically. First test on synthetic gradients g_k = theta z_k v_0 + xi_k, where z_k is standard normal and xi_k is standard isotropic Gaussian noise, using d equal to 512 and 2048. Compare AdamW, always-on Oja correction, and the gated method. Then test a small Transformer on TinyStories or CIFAR-10. Success requires fewer early loss spikes, equal or better final loss, and lower memory than a dense covariance preconditioner. Ablate the log(d) threshold and the sigmoid gate to check whether the predicted transition explains the benefit rather than generic warmup.