🤖 AI Summary
This study addresses the limitation of conventional low-rank compression, which allocates a fixed computational budget to all tokens and lacks dynamic adaptability. We propose LRCC, a method that trains lightweight routers to select nested low-rank paths for each Transformer block, thereby enabling token-level conditional computation. By freezing the low-rank factors and optimizing only the routers, LRCC introduces an input-dependent dynamic computation mechanism with minimal overhead. Under an identical average active parameter budget, LRCC improves the average accuracy on both upstream and downstream tasks of Llama-2-7B by 7.6 percentage points. Furthermore, it significantly enhances perplexity and accuracy under matched latency constraints, establishing a new paradigm for efficient model fine-tuning.
📝 Abstract
Low-rank compression reduces the cost of pretrained language models by replacing linear transformations with low-rank factorizations. However, conventional methods use a fixed rank allocation during inference, assigning the same amount of compute regardless of the input token. We introduce Low-Rank Conditional Computation (LRCC), which adds token-dependent computation to pretrained models by training one lightweight router per Transformer block to select among a small set of nested low-rank paths. During training, the low-rank factors remain frozen, and only the routers are optimized. We evaluate LRCC on Llama and Qwen models for language modeling and zero-shot downstream tasks. Within the same average active-parameter budget, LRCC improves the predictive performance over static low-rank compression, including a 7.6 percentage-point gain in average downstream accuracy on Llama-2-7B over static methods. At matched batch-size-1 decoding latency, LRCC improves both perplexity and downstream accuracy on Llama-3.2-1B and remains competitive on Llama-2-7B, without specialized kernels. Finally, we assess the usefulness of assigning a token-wise path by analyzing the routers' path choices.