LaX: Boosting Low-Rank Training of Foundation Models via Latent Crossing

TL;DR

LaX is a simple module that enhances low-rank models of foundation models, enabling them to perform on par with full-rank models while using significantly fewer parameters.

Contribution

Introduces Latent Crossing (LaX), a plug-and-play module that improves low-rank model capacity by facilitating information flow across subspaces.

Findings

LaX boosts low-rank model performance to match or exceed full-rank baselines.

LaX reduces parameter count by 2-3 times compared to full models.

LaX improves fine-tuned LLaMA performance on reasoning tasks.

Abstract

Training foundation models such as ViTs and LLMs requires tremendous computing cost. Low-rank matrix or tensor factorization offers a parameter-efficient alternative, but often downgrades performance due to the restricted parameter space. In this work, we introduce {\textbf{Latent Crossing (LaX)}} -- a simple yet effective plug-and-play module that enhances the capacity of low-rank models by enabling information flow across low-rank subspaces. We extensively validate the benefits of LaX on pre-training tasks with ViT-Base/Large and LLaMA-like models ranging from 60M to 1B parameters. LaX boosts low-rank model performance to match or exceed the full-rank baselines while using 2-3\(\times\) fewer parameters. When equipped with low-rank adapters (i.e., LoRA) for fine-tuning LLaMA-7/13B, LaX consistently improves performance on arithmetic and common sense reasoning tasks with negligible cost.