Mixture of In-Context Experts Enhance LLMs' Long Context Awareness

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Lin, Hongzhan, Lv, Ang, Chen, Yuhan, Zhu, Chen, Song, Yang, Zhu, Hengshu, Yan, Rui
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866912074902274048
author Lin, Hongzhan
Lv, Ang
Chen, Yuhan
Zhu, Chen
Song, Yang
Zhu, Hengshu
Yan, Rui
author_facet Lin, Hongzhan
Lv, Ang
Chen, Yuhan
Zhu, Chen
Song, Yang
Zhu, Hengshu
Yan, Rui
contents Many studies have revealed that large language models (LLMs) exhibit uneven awareness of different contextual positions. Their limited context awareness can lead to overlooking critical information and subsequent task failures. While several approaches have been proposed to enhance LLMs' context awareness, achieving both effectiveness and efficiency remains challenging. In this paper, for LLMs utilizing RoPE as position embeddings, we introduce a novel method called "Mixture of In-Context Experts" (MoICE) to address this challenge. MoICE comprises two key components: a router integrated into each attention head within LLMs and a lightweight router-only training optimization strategy: (1) MoICE views each RoPE angle as an `in-context' expert, demonstrated to be capable of directing the attention of a head to specific contextual positions. Consequently, each attention head flexibly processes tokens using multiple RoPE angles dynamically selected by the router to attend to the needed positions. This approach mitigates the risk of overlooking essential contextual information. (2) The router-only training strategy entails freezing LLM parameters and exclusively updating routers for only a few steps. When applied to open-source LLMs including Llama and Mistral, MoICE surpasses prior methods across multiple tasks on long context understanding and generation, all while maintaining commendable inference efficiency.
format Preprint
id arxiv_https___arxiv_org_abs_2406_19598
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Mixture of In-Context Experts Enhance LLMs' Long Context Awareness
Lin, Hongzhan
Lv, Ang
Chen, Yuhan
Zhu, Chen
Song, Yang
Zhu, Hengshu
Yan, Rui
Computation and Language
Many studies have revealed that large language models (LLMs) exhibit uneven awareness of different contextual positions. Their limited context awareness can lead to overlooking critical information and subsequent task failures. While several approaches have been proposed to enhance LLMs' context awareness, achieving both effectiveness and efficiency remains challenging. In this paper, for LLMs utilizing RoPE as position embeddings, we introduce a novel method called "Mixture of In-Context Experts" (MoICE) to address this challenge. MoICE comprises two key components: a router integrated into each attention head within LLMs and a lightweight router-only training optimization strategy: (1) MoICE views each RoPE angle as an `in-context' expert, demonstrated to be capable of directing the attention of a head to specific contextual positions. Consequently, each attention head flexibly processes tokens using multiple RoPE angles dynamically selected by the router to attend to the needed positions. This approach mitigates the risk of overlooking essential contextual information. (2) The router-only training strategy entails freezing LLM parameters and exclusively updating routers for only a few steps. When applied to open-source LLMs including Llama and Mistral, MoICE surpasses prior methods across multiple tasks on long context understanding and generation, all while maintaining commendable inference efficiency.
title Mixture of In-Context Experts Enhance LLMs' Long Context Awareness
topic Computation and Language
url https://arxiv.org/abs/2406.19598