Duo-LLM: A Framework for Studying Adaptive Computation in Large Language Models

TL;DR

This paper introduces Duo-LLM, a framework that enables adaptive computation in large language models by dynamically routing tokens through smaller or larger modules based on their complexity, improving efficiency and understanding of internal routing processes.

Contribution

It proposes a novel framework with auxiliary modules for dynamic token routing in LLMs, providing insights into optimal patterns and the gap between practical routing and theoretical optima.

Findings

Activating a large module in one layer outperforms using it across all layers.

Trained routers differ from oracle solutions, often being suboptimal.

The framework reveals the internal routing dynamics and potential for efficiency gains.

Abstract

Large Language Models (LLMs) typically generate outputs token by token using a fixed compute budget, leading to inefficient resource utilization. To address this shortcoming, recent advancements in mixture of expert (MoE) models, speculative decoding, and early exit strategies leverage the insight that computational demands can vary significantly based on the complexity and nature of the input. However, identifying optimal routing patterns for dynamic execution remains an open challenge, limiting the full potential of these adaptive methods. To address this need, we study adaptive computation in LLMs more systematically. We propose a novel framework that integrates smaller auxiliary modules within each Feed-Forward Network layer of the LLM. This design enables dynamic routing of tokens based on task complexity: tokens can be processed by either the small or big modules at each layer, or even bypass certain layers entirely. This allows us to introduce a novel notion of a token's difficulty, defined by its potential to benefit from additional computational resources. Importantly, by employing oracles to identify optimal patterns of adaptive computations, we gain valuable insights into the internal workings of LLMs and the routing processes in a simplified heterogeneous MoE setup. We show that trained routers operate differently from oracles and often yield suboptimal solutions. Notably, activating a large module in just one layer outperforms models that use large modules across all layers, underscoring the gap between practical implementations of routing in MoE models and theoretical optima for adaptive computation.