Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

Dynamic Rank Reinforcement Learning for Adaptive Low-Rank Multi-Head Self Attention in Large Language Models

About

Dynamic Rank Reinforcement Learning (DR-RL) approximations rely on static rank assumptions, limiting their flexibility across diverse linguistic contexts. Our method dynamically modulates ranks based on real-time sequence dynamics, layer-specific sensitivities, and hardware constraints. The core innovation is a deep reinforcement learning agent that formulates rank selection as a sequential policy optimization problem, strictly balancing attention fidelity against computational latency. To ensure stability during inference, we derive and employ online matrix perturbation bounds, enabling incremental rank updates without the prohibitive cost of full decomposition. Furthermore, the integration of a lightweight Transformer-based policy network and batched Singular Value Decomposition (SVD) operations ensures scalable deployment on modern architectures. Extensive experiments demonstrate that DR-RL significantly reduces Floating Point Operations (FLOPs) by over 40% in long-sequence regimes (L > 4096) while maintaining downstream accuracy statistically equivalent to full-rank attention. Beyond standard language modeling benchmarks, we validate the real-world applicability of DR-RL on the GLUE benchmark. Specifically, our method achieves 92.78% accuracy on the SST-2 sentiment analysis task, matching the performance of full-rank baselines and outperforming static low-rank methods, such as Performer and Nystr\"omformer, by a significant margin.

Caner Erden• 2025

Related benchmarks

TaskDatasetResultRank
Language ModelingPTB
Perplexity46.5
650
Language ModelingWikiText-103 (test)
Perplexity24.7
524
Language ModelingPTB (test)
Perplexity46.5
471
Language ModelingWikiText-103
PPL24.7
146
Language ModelingBookCorpus (test)
PPL29.8
15
Natural Language UnderstandingGLUE
SST-2 Accuracy92.8
5
Showing 6 of 6 rows

Other info

Follow for update