Source-linked AI summary
Improving Recursive Transformers with Mixture of LoRAs
Mohammadmahdi Nouriborji, Morteza Rohanian, Omid Rohanian
TL;DR
Parameter sharing makes recursive transformers compact but can constrain expressivity, while existing adapter approaches do not modulate the shared FFN conditionally. The paper introduces MoL within ModernALBERT and reports strong compact-model performance across benchmarks. It also identifies routing complexity as a remaining limitation.
Problem
Parameter sharing reduces memory use, but can limit expressivity and require wider, more computationally costly shared layers to match fully parameterised models.
Source: §1 · p. 1Method
ModernALBERT embeds a sparse pool of LoRA experts in the shared FFN and routes tokens to experts to produce input-dependent behavior while preserving parameter efficiency.
Source: §1 · p. 2Results
Takeaways & Limitations
The MoL layer integrates conditional computation into a shared transformer architecture while maintaining parameter efficiency.
Source: §1 · p. 2, §8 · p. 9Mixture-of-experts routing remains relatively computationally complex compared with fully shared or purely dense alternatives.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Parameter sharing in recursive transformers reduces model size but collapses layer-wise expressivity. We propose Mixture of LoRAs (MoL), a lightweight conditional-computation mechanism that inserts Low-Rank Adaptation (LoRA) experts inside a shared feed-forward network (FFN). MoL enables token-conditional weight-space modulation of the shared FFN without untying backbone parameters, unlike prior approaches that add fixed or externally attached adapters. We pretrain a modernised recursive architecture, ModernALBERT, integrating rotary embeddings, GeGLU, FlashAttention, and a distillation-based initialisation. Across GLUE, SQuAD-v2, and BEIR, ModernALBERT (50M--120M) achieves state-of-the-art performance among compact models and surpasses larger fully parameterised baselines. We also propose an expert-merging procedure that compresses MoL into a single adapter at inference while preserving accuracy, enabling efficient deployment. Our results show that conditional weight-space modulation effectively restores the expressivity lost under aggressive parameter sharing in recursive transformers.
1. Introduction
Parameter sharing makes transformers smaller but can limit expressivity, while sparse experts can increase capacity at a parameter cost. The paper introduces MoL within ModernALBERT to restore token-dependent behavior in a shared FFN, and reports strong compact-model benchmark performance.
1. Introduction
- Parameter sharing reduces memory use, but matching fully parameterised models can require wider shared layers, increasing computational cost.The paper frames preserving expressivity while minimizing both parameters and computation as the motivation for its approach.Source: §1 · p. 1
- MoE can expand representational capacity by activating specialized experts, but its parameter growth conflicts with lightweight shared architectures.The paper identifies this as a limitation of applying MoE directly to parameter-efficient models.Source: §1 · p. 2
- Prior static adapters lack token-level adaptation, while routed adapter approaches apply experts after the FFN; RRT adds static depth-specific LoRAs.The paper identifies conditional computation inside the shared FFN during pretraining as the missing design element.Source: §1 · p. 2
- MoL routes tokens to a sparse subset of LoRA experts embedded in the shared FFN, giving the shared backbone distinct behaviors for different inputs while preserving parameter efficiency.The authors present this as restoring layer-wise functional diversity typically lost through parameter sharing.Source: §1 · p. 2
- ModernALBERT achieves state-of-the-art performance among compact transformer models on GLUE, SQuAD, and BEIR.It integrates GeGLU, RoPE, and FlashAttention, which the authors say improve training stability and runtime efficiency.Source: §1 · p. 2
- The paper introduces a merging strategy that compresses MoL experts into a single adapter for efficient deployment.It also evaluates 50M–120M model variants across GLUE, SQuAD, and BEIR.Source: §1 · p. 2
2. Related Works
Related work describes the efficiency–expressivity trade-off in recursive parameter sharing and approaches that add experts or lightweight modules. ModernALBERT addresses an open direction by combining recursive sharing with token-conditional expert modulation.
2. Related Works
- Static cross-layer sharing reduces memory but can limit expressivity, and matching fully parameterised models may require wider, more computationally costly layers.ALBERT and the Universal Transformer are cited as examples of layer tying and its trade-off.Source: §2 · p. 2
- MoE routes tokens to sparse expert subsets, while many mixture-of-adapter methods use small plug-in modules for fine-tuning rather than integrating conditional computation during pretraining.These frameworks are predominantly applied on top of existing models.Source: §2 · p. 2
- Modern architectural components such as rotary positional embeddings, gated FFNs, and I/O-optimized attention are described as improving efficiency and performance.ModernBERT is cited as demonstrating gains over standard transformer baselines without increased computational cost.Source: §2 · pp. 2–3
- ALBERT shares one attention and FFN block without adaptation; RRT adds depth-specific static LoRA, MoA routes adapters after the FFN, and MoL injects routed low-rank updates into shared FFN weights.MoL thereby enables token-conditional modulation while retaining parameter sharing in the backbone.Source: Figure 2 · p. 3
- RRT and MiniALBERT use static layer-specific LoRA or bottleneck adapters to augment shared backbones, while dynamic conditional computation in compact backbones remains open.ModernALBERT combines recursive parameter sharing with token-conditional expert modulation.Source: §2 · p. 3
3. ModernALBERT
ModernALBERT combines recursive parameter sharing with lightweight modules, using Mixture of LoRAs (MoL) to add conditional flexibility inside shared feed-forward networks. The section outlines its recursive design, MoL routing, distillation-based preparation, training settings, and model variants.
3.3. Effects of Knowledge Distillation and Initialisations
- Knowledge distillation, improved initialization, and architectural refinements are used to stabilize training and improve performance across model sizes.Source: §3 · p. 3
3.1. Parameter-Sharing through Recursions
- Within K groups of G layers, layers in each group share attention and FFN weights; full sharing uses one parameter set and reduces unique parameters by K/N.Source: §3.1 · p. 4
3.2. Improved Flexibility through Mixture of LoRAs
- MoL augments the shared FFN with lightweight LoRA experts and uses sparse top-2 routing to combine expert-modified outputs for conditional computation.The two selected experts' routing weights are renormalized to sum to 1.Source: §3.2 · p. 4
3.4. Training Settings
- 20,000–30,000 RedPajamas-1T warm-up steps precede 70,000–80,000 RefinedWeb steps, with global batch size 384 and maximum sequence length 1024.Source: §3.4 · p. 5
3.5. ModernALBERT Variants
- Pre-Norm, GeGLU, RoPE, and FlashAttention accompany variants with eight-expert top-2 MoL at group ends, except the tiny model's four-expert top-1 setup.Source: §3.5 · p. 5
4. Experimental Results
ModernALBERT performs strongly across GLUE, SQuAD-v2, and BEIR, while ablations compare MoL with alternative conditional-computation designs and initialization strategies.
4.1. GLUE Benchmark Results
- 88.72 GLUE average for ModernALBERT-large (120M parameters) surpasses ModernBERT-base (149M parameters, 88.45).It also reaches base-class state-of-the-art results on RTE (88.44), STS-B (92.1), and MRPC (92.7), and matches ModernBERT on QNLI (93.7) and MNLI (88.9).Source: §4.1 · pp. 5–6, Table 2 · p. 6
- 92.8 SQuAD-v2 F1 for ModernALBERT-base exceeds ModernBERT-base (92.6) and ALBERT-xxlarge (92.5).ModernALBERT also scores 86.1 Exact Match.Source: Table 3 · p. 6, §4.1 · p. 6
4.2. Ablation Studies
- With 8 experts and top-2 routing, MoL scores a GLUE average of 77.24 versus 76.87 for equivalent MoA, while scaling from one expert to eight gains 1.16 points.With the same stepwise initialization and MLM training, ModernALBERT reaches a GLUE average of 81.94 versus 80.95 for RRT; ablation models used 30k steps, batch size 128, and sequence length 1024.Source: §4.2 · pp. 6–7
4.1. GLUE Benchmark Results
- 48.82 on BEIR ArguAna for ModernALBERT exceeds ModernBERT (35.7).ModernALBERT has the best average performance among compact models on TREC-COVID, FiQA, and ArguAna.Source: §4.1 · p. 6, Table 4 · p. 7
5. Optimising MoL for Inference
MoL’s conditional routing adds inference latency, so the paper compresses its experts into a static adapter. It compares uniform averaging with EMA-based merging, which preserves nearly all unmerged accuracy.
5.1. Expert Merging Strategies
- Expert merging synthesizes one dense adapter from the E LoRA experts, and fine-tuning disables routing while updating that merged adapter directly.The procedure converts sparse MoL into a standard static LoRA module.Source: §5.1 · p. 7
5.1.1. UNIFORM INITIALIZATION
- Uniform initialization averages the pretrained experts with w_j = 1/E, then updates the merged adapter through standard backpropagation.This uses a neutral starting point without token-specific routing.Source: §5.1.1 · p. 7
5.1.2. DYNAMIC EMA MERGING
- 86.28 versus 84.83 on RTE: EMA-based merging outperforms uniform initialization and preserves nearly all unmerged-model accuracy.EMA uses routing statistics to adjust expert contributions; the router is trained jointly for data-rich MNLI but frozen for smaller STS-B, RTE, and SST-2 datasets.Source: §5.1.2 · pp. 7–8
5. Optimising MoL for Inference
- 9.46 ms latency, 106,527 tokens/s throughput, and 0.196 GB GPU memory are reported for merged ModernALBERT-tiny.The merged-adapter deployment reduces inference cost to that of a dense model.Source: §5.1 · p. 7, §5 · p. 7, Table 7 · p. 8
6. Discussion
MoL experts specialize across inputs, but dynamic routing increases inference cost. Expert merging suggests the learned expressivity can often be compressed for inference, with the footprint of compact dense models.
6. Discussion
- MoL routers specialize experts for different data aspects, including syntactic cues, semantic relations, and domain-specific patterns.This division of responsibility yields richer internal representations and helps mitigate expressivity limitations of fully shared architectures.Source: §6 · p. 8
- Dynamic routing increases inference cost because several experts may be activated per token.Source: §6 · p. 8
- Expert merging suggests learned expressivity can often be compressed into a static inference representation with the footprint of compact dense models.Table 8 concerns expert merging across datasets, but the supplied text does not report its individual results.Source: §6 · pp. 8–9
7. Limitations
The limitations include costly mixture-of-experts routing and global attention without locality-aware mechanisms. The authors identify expert selection, load balancing, and attention patterns as areas for further work.
7. Limitations
- Mixture-of-experts routing remains relatively costly compared with fully shared or purely dense alternatives.Source: §7 · p. 9
- Lightweight routing and optional merging mitigate dispatching overhead, but efficient expert selection and load balancing still need improvement.Source: §7 · p. 9
- Global attention lacks locality-aware mechanisms, and the approach may underperform on extended-context retention or fine-grained reasoning tasks.The authors suggest local or block-sparse attention, or hybrid schemes, as possible future directions.Source: §7 · p. 9
8. Conclusion and Future Work
The paper introduces MoL as conditional computation for recursive transformers, reporting improved performance while maintaining parameter efficiency. It proposes future applications in multimodal settings and large autoregressive language models.
8. Conclusion and Future Work
- The MoL layer integrates conditional computation into shared recursive transformers and demonstrated performance improvements while maintaining parameter efficiency.Source: §8 · p. 9
- Extending the framework to multimodal settings is proposed as a direction for studying conditional representations across diverse input types.Source: §8 · p. 9
- Applying parameter sharing with MoL to large autoregressive language models could reduce parameter count and memory footprint while preserving or enhancing representational capacity.The authors suggest this could contribute to more efficient large-scale generative models.Source: §8 · p. 9