Benyamin Jamialahmadi
Author directory2026
DTRNet: Dynamic Token Routing Network to Reduce Quadratic Costs in Transformers
Aman Sharma | Saeed Najafi | Parsa Farinneya | Benyamin Jamialahmadi | Marzieh S. Tahaei | Yuhe Fan | Mehdi Rezagholizadeh | Boxing Chen | Aref Jafari
Findings of the Association for Computational Linguistics: EMNLP 2026
Aman Sharma | Saeed Najafi | Parsa Farinneya | Benyamin Jamialahmadi | Marzieh S. Tahaei | Yuhe Fan | Mehdi Rezagholizadeh | Boxing Chen | Aref Jafari
Findings of the Association for Computational Linguistics: EMNLP 2026
Transformers achieve state-of-the-art results across many tasks, but their uniform application of quadratic self-attention to every token at every layer makes them computationally expensive. We introduce DTRNet (Dynamic Token Routing Network), an improved Transformer architecture that allows tokens to dynamically skip the quadratic cost of cross-token mixing while still receiving lightweight linear updates. By preserving the MLP module and reducing the attention cost for most tokens to linear, DTRNet ensures that every token is explicitly updated while significantly lowering overall computation. This design offers an efficient and effective alternative to standard dense attention. Once trained, DTRNet blocks routes only ~10% of tokens through attention at each layer while maintaining performance comparable to a full Transformer. It outperforms routing-based layer skipping methods such as MoD and D-LLM in both accuracy and memory at matched FLOPs, while routing fewer tokens to full attention. Its efficiency gains, scales with sequence length, offering significant reduction in FLOPs for long-context inputs. By decoupling token updates from attention mixing, DTRNet substantially reduces the quadratic share of computation, providing a simple, efficient, and scalable alternative to Transformers.
2025
Balcony: A Lightweight Approach to Dynamic Inference of Generative Language Models
Benyamin Jamialahmadi | Parsa Kavehzadeh | Mehdi Rezagholizadeh | Parsa Farinneya | Hossein Rajabzadeh | Aref Jafari | Boxing Chen | Marzieh S. Tahaei
Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing
Benyamin Jamialahmadi | Parsa Kavehzadeh | Mehdi Rezagholizadeh | Parsa Farinneya | Hossein Rajabzadeh | Aref Jafari | Boxing Chen | Marzieh S. Tahaei
Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing
Deploying large language models (LLMs) in real-world applications is often hindered by strict computational and latency constraints. While dynamic inference offers the flexibility to adjust model behavior based on varying resource budgets, existing methods are frequently limited by hardware inefficiencies or performance degradation. In this paper, we introduce Balcony, a simple yet highly effective framework for depth-based dynamic inference. By freezing the pretrained LLM and inserting additional transformer layers at selected exit points, Balcony maintains the full model’s performance while enabling real-time adaptation to different computational budgets. These additional layers are trained using a straightforward self-distillation loss, aligning the sub-model outputs with those of the full model. This approach requires significantly fewer training tokens and tunable parameters, drastically reducing computational costs compared to prior methods. When applied to the LLaMA3-8B model, using only 0.2% of the original pretraining data, Balcony achieves minimal performance degradation while enabling significant speedups. Remarkably, we show that Balcony outperforms state-of-the-art methods such as Flextron and Layerskip on multiple models at various scales, as well as other leading compression techniques across a variety of benchmarks.