Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

Technical Summary: Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

Problem Statement
The pre-training of Large Language Models (LLMs) is constrained by the memory costs associated with storing full-rank weights, gradients, and optimizer states. Low-rank pre-training has emerged as a solution to reduce these costs, with methods evolving into three categories: optimizer-based (e.g., GaLore, Fira), architecture-based (e.g., CoLA, SLTrain), and adapter-based (e.g., ReLoRA). However, the evaluation of these methods remains limited. Current practices rely almost exclusively on validation perplexity from single-seed runs. The authors argue that perplexity is a poor proxy for solution quality; models with identical perplexity may converge to fundamentally different regions of the loss landscape with distinct internal representations. Furthermore, existing comparisons often fail to distinguish between how well a method fine-tunes a pre-trained model versus how well a model pre-trained with that method generalizes downstream.

Methodology
The authors introduce a diagnostic framework to characterize the solutions found by five low-rank pre-training methods (GaLore, Fira, CoLA, SLTrain, ReLoRA) against full-rank training across three model scales (60M, 130M, 350M). The study evaluates these methods using 16 metrics across four dimensions:

1. 1-D Loss Landscape (Random Directions): The authors perturb model parameters along 100 random directions to measure the mean and variance of the validation loss. They derive scalar metrics for Sharpness (the rate of loss growth away from the checkpoint) and Direction Variance (anisotropy of the loss surface).
2. 1-D Loss Landscape (Top-K PCA Directions): Perturbations are applied along the top singular directions of the weight matrices to analyze the geometry along the most significant axes.
3. 1-D Interpolation: The authors measure the loss along linear paths between checkpoints. They introduce two specific metrics:
   • Consecutive Checkpoints Barrier Height (CCBH): Measures the barrier between sequential checkpoints of the same method.
   • Inter-Method Barrier Height (IMBH): A novel metric quantifying the geometric separation between different training methods at the same training step. This reveals whether different methods converge to the same basin or geometrically distinct solutions.
4. Spectral and Rank Metrics: The study tracks the spectral structure of weights and updates, including Effective Rank, Stable Rank, Spectral Gap, and Threshold Rank.
5. Activation Similarity: The authors compare hidden states of low-rank models against full-rank baselines using Activation L2 distance, Cosine Similarity, and Linear CKA (Centered Kernel Alignment).

Finally, the authors test whether these geometric and spectral features can predict downstream zero-shot performance better than perplexity alone, using linear regression with leave-one-size-out (LOSO) and leave-one-method-out (LOMO) cross-validation.

Key Contributions

1. Diagnostic Framework: The paper proposes a comprehensive evaluation suite including the novel IMBH metric, which quantifies geometric separation between training algorithms rather than just between seeds or epochs.
2. Cross-Family Empirical Findings: The study establishes that low-rank methods are not equivalent to full-rank training nor to each other. Key findings include:
   • Geometric Separation: Methods converge to geometrically distinct basins, with IMBH growing monotonically as training progresses.
   • Spectral Signatures: GaLore closely matches the spectral structure of full-rank updates, whereas CoLA does not.
   • Representation Divergence: Low-rank activations increasingly diverge from full-rank in later layers. CoLA exhibits orthogonal activations, while Fira recovers alignment only in the final layer.
3. Improved Downstream Prediction: Validation perplexity alone is an incomplete predictor of downstream performance. Augmenting perplexity with geometric and spectral features significantly improves prediction accuracy (LOSO Pearson correlation increases from 0.873 to 0.913).
4. Unified Codebase: The authors release an open-source pipeline integrating pre-training, metric computation, and downstream evaluation to facilitate further research.

Results

• Loss Landscape Geometry: Full-rank training settles into sharper basins along random directions but flatter basins along the top-1 PCA direction. Conversely, low-rank methods like Fira and SLTrain converge to flatter basins along random directions, while CoLA and ReLoRA reach the sharpest basins. GaLore exhibits a mixed profile: moderately flat along random directions but sharp along its top-1 PCA direction.
• Inter-Method Barriers: All methods occupy geometrically separated basins. At the 60M scale, the barrier between full-rank and any low-rank method is higher than barriers between low-rank methods. Fira converges to the basin geometrically closest to full-rank, while SLTrain remains the most isolated.
• Spectral Analysis: Full-rank dominates all rank metrics. GaLore tracks full-rank closely in effective and stable rank but exhibits a smaller spectral gap. Fira and GaLore most faithfully reproduce the spectral distribution of full-rank weight updates.
• Activation Alignment: ReLoRA tracks full-rank representations most closely across all scales. CoLA shows the largest deviation, with activations becoming orthogonal to the baseline in later layers. Fira shows high layer-averaged drift but surprisingly low deviation in the final layer.
• Downstream Performance: Validation perplexity does not consistently predict downstream ranking. For instance, Fira achieves the lowest perplexity at all scales, yet CoLA leads in multiple downstream categories at 130M and 350M despite not ranking in the top two for perplexity.

Significance
The paper concludes that low-rank pre-training methods are not interchangeable approximations of full-rank training; they fundamentally alter the solutions reached. The study demonstrates that validation perplexity provides an incomplete picture of model quality. By incorporating geometric and spectral metrics, researchers can better predict downstream performance and understand the distinct inductive biases introduced by different low-rank strategies. The work highlights that while some methods (like Fira) converge to basins geometrically close to full-rank, others (like CoLA) occupy distinct valleys, suggesting that the choice of low-rank method dictates the specific region of the loss landscape the model inhabits.