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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Alizadeh, Keivan, Mirzadeh, Iman, Shahrokhi, Hooman, Belenko, Dmitry, Sun, Frank, Cho, Minsik, Sekhavat, Mohammad Hossein, Nabi, Moin, Farajtabar, Mehrdad
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910649389416448
author Alizadeh, Keivan
Mirzadeh, Iman
Shahrokhi, Hooman
Belenko, Dmitry
Sun, Frank
Cho, Minsik
Sekhavat, Mohammad Hossein
Nabi, Moin
Farajtabar, Mehrdad
author_facet Alizadeh, Keivan
Mirzadeh, Iman
Shahrokhi, Hooman
Belenko, Dmitry
Sun, Frank
Cho, Minsik
Sekhavat, Mohammad Hossein
Nabi, Moin
Farajtabar, Mehrdad
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_10846
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Duo-LLM: A Framework for Studying Adaptive Computation in Large Language Models
Alizadeh, Keivan
Mirzadeh, Iman
Shahrokhi, Hooman
Belenko, Dmitry
Sun, Frank
Cho, Minsik
Sekhavat, Mohammad Hossein
Nabi, Moin
Farajtabar, Mehrdad
Machine Learning
Computation and Language
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.
title Duo-LLM: A Framework for Studying Adaptive Computation in Large Language Models
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2410.10846