Research map / Architectures and efficiency / Pruning and efficient tuning
LaX: Boosting Low-Rank Training of Foundation Models via Latent Crossing
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Latent Crossing (LaX), a plug-and-play residual module that passes latent features between consecutive low-rank layers (with optional gates) to recover performance lost in low-rank training of ViTs, LLaMA-style LLMs, and LoRA fine-tuning.
Problem
Low-rank factorized models (SVD, tensor train, CoLA, LoRA) save parameters and compute but lose capacity, degrading performance at small ranks, while larger ranks erode the efficiency gains.
Approach
Instead of increasing the physical rank, add residual connections that route the latent representation of the previous low-rank layer into the up-projection of the current layer (Inter-Layer LaX), or across tensor-train contraction steps (Intra-Layer LaX), with learnable or identity gates and layer normalization; this adds almost no parameters or FLOPs.
Key results
- On ImageNet-1K with ViT-B, LaX raises Tensor Train accuracy from 71.11% to 75.43% (+4.32%) and SVD to 77.20%, with parameter overhead of at most 0.2%.
- On C4 pre-training at 60M, SVD+LaX reaches 33.54 PPL (vs 36.25 for SVD), and CoLA+LaX surpasses LORO on all scales from 60M to 350M.
- LoRA+LaX improves average accuracy on LLaMA-7B commonsense reasoning from 74.7% to 77.3% and on arithmetic reasoning from 61.6% to 63.4%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add LaX to low-rank ViT pre-training | SVD, Tensor Train, and CoLA · ImageNet-1k, ViT-B | accuracy (%): SVD 75.20 -> 77.20; Tensor Train 71.11 -> 75.43; CoLA 76.04 -> 77.84 | Helped |
| Add LaX to low-rank ViT pre-training | SVD, Tensor Train, and CoLA · ImageNet-1k, ViT-L | accuracy (%): SVD 76.81 -> 78.60; Tensor Train 75.21 -> 77.77; CoLA 77.63 -> 79.07 | Helped |
| Add LaX to low-rank pre-training | SVD and CoLA · C4 validation, 60M–1B LLaMA-like models | perplexity: SVD: 36.25 -> 33.54; 26.84 -> 24.63; 21.18 -> 18.90; 16.54 -> 15.51; CoLA: 34.04 -> 33.21; 24.48 -> 24.21; 19.40 -> 18.5 | Helped |
| Add LaX to LoRA | LLaMA-7B and LLaMA-13B · Math10K arithmetic reasoning average | accuracy (%): 61.6 -> 63.4 (+1.8); 65.2 -> 67.3 (+2.1) | Helped |
| Add LaX to LoRA | LLaMA-7B · Commonsense reasoning average | accuracy (%): 74.7 -> 77.3 (+2.6) | Helped |
| Add LaX to LoRA | LLaMA-13B · Commonsense reasoning tasks | accuracy (%): Average +1.1%; BoolQ -0.8% | Mixed |
| Use Tensor Gate for LaX instead of Identity Gate (ablation) | SVD ViT-B rank=256 · ImageNet-1k | accuracy (%): 75.81 -> 77.20 | Helped |
| Add CoLA activation to Tensor Train + LaX (ablation) | Tensor Train + LaX · ImageNet-1k, ViT-B | test accuracy (%): 75.43 -> 77.73 | Helped |
| Add Inter-LaX and Intra-LaX (ablation) | Tensor Train · ImageNet-1k, ViT-B | test accuracy (%): 71.11 -> 75.43 | Helped |
| Use Inter-LaX instead of Intra-LaX only (ablation) | Tensor Train · ImageNet-1k, ViT-B | test accuracy (%): 73.25 -> 73.87 | Helped |
| Remove LayerNorm from LaX (ablation) | SVD + LaX · ImageNet-1k, ViT-B | test accuracy (%): 77.20 -> 0.10 | Hurt |
| Apply LaX to V projection only rather than no layers (ablation) | SVD ViT-B · ImageNet-1k | test accuracy (%): 75.20 -> 75.91 (+0.71) | Helped |
Surprising findings
- Removing LayerNorm from LaX causes training to diverge immediately (NaN loss), showing normalization is essential.
- For SVD, the Dense Gate does not consistently help over the Identity Gate in LLM pre-training, unlike CoLA.
Limitations
- Gating variant benefit is mixed across methods in LLM pre-training; the paper recommends a conservative Identity Gate.
- Pre-training runs are too costly for error bars, so no statistical significance is reported.
- LLM pre-training results for baselines other than SVD and CoLA are cited from prior work rather than re-run.
Experimental setup
Datasets: ImageNet-1K C4 Math10K MultiArith GSM8K AddSub AQuA SingleEq SVAMP BoolQ PIQA SIQA
Benchmarks: ImageNet-1K C4 MultiArith GSM8K AddSub AQuA SingleEq SVAMP BoolQ PIQA SIQA HellaSwag
Metrics: top-1 accuracy validation perplexity parameter count memory FLOPs
Baselines: Full-rank SVD Tensor Train CoLA ReLoRA GaLore SLTrain LORO LoRA
Builds on: LoRA CoLA ResNet tensor train decomposition GaLore
Closest papers by meaning
- Asymmetry in Low-Rank Adapters of Foundation Models (ICML 2024 · similarity 0.87)
- A Stronger Mixture of Low-Rank Experts for Fine-Tuning Foundation Models (ICML 2025 · similarity 0.85)
- From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications (ICML 2025 · similarity 0.85)
- LoRA-X: Bridging Foundation Models with Training-Free Cross-Model Adaptation (ICLR 2025 · similarity 0.85)
- AdaRankGrad: Adaptive Gradient Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning (ICLR 2025 · similarity 0.84)
- The Expressive Power of Low-Rank Adaptation (ICLR 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.