Search papers, labs, and topics across Lattice.
The paper introduces a dual-path transformer block that decouples compute (depth) and capacity (width) scaling, addressing the limitation of looped transformers which sacrifice capacity for compute efficiency. This architecture consists of a deep, looped sublayer and a wide, feed-forward sublayer, combined via learned per-token gates. Experiments demonstrate that the dual-path model outperforms iso-FLOP baseline transformers on language modeling and downstream tasks, achieving better performance with fewer parameters, and reveals interpretable token-level routing decisions.
LLMs can get a free performance boost: decoupling compute and capacity within each layer lets you beat standard transformers at the same FLOPs.
Looped transformers apply a shared block multiple times and have emerged as a parameter-efficient route to scaling compute in language models. However, at fixed FLOPs a looped model has strictly less capacity than a baseline transformer. We propose a novel dual-path block that can flexibly scale compute, the number of sequential operations applied to a hidden state, and capacity, the parameters available at a single step. For this we expose both axes as parallel pathways within a single layer: a deep sublayer re-applied K times with shared parameters, and a wide sublayer with an enlarged feed-forward network applied once. Independent per-token gates combine both axes and allow detailed per-token routing analyses. We show that across two FLOP budgets, our dual-path model surpasses iso-FLOP matched models on language modeling and downstream evaluations, while using fewer parameters than the baseline at matched FLOPs. The learned gates are directly interpretable and show systematic per-token allocation with function words and lexical content trend wide, while punctuation, symbols, and arithmetic tokens trend deep.