Table of Contents
Fetching ...

Gecko: An Efficient Neural Architecture Inherently Processing Sequences with Arbitrary Lengths

Xuezhe Ma, Shicheng Wen, Linghao Jin, Bilge Acun, Ruihang Lai, Bohan Hou, Will Lin, Hao Zhang, Songlin Yang, Ryan Lee, Mengxi Wu, Jonathan May, Luke Zettlemoyer, Carole-Jean Wu

TL;DR

Gecko tackles the fundamental challenge of long-context sequence modeling by building on the Megalodon backbone and introducing three core innovations: timestep decay normalization to stabilize long-horizon statistics, sliding chunk attention to maintain efficient local and cross-chunk coherence, and adaptive working memory with a position-aware online softmax to compress and retrieve information from very long histories. These components are tightly integrated to enable inherently long-context processing without explicit context-extension techniques, achieving superior training efficiency and long-context performance compared to Llama2 and Megalodon while handling sequences up to millions of tokens. Empirically, Gecko demonstrates lower training perplexity and stronger performance on short-context academic benchmarks, and substantial gains on long-context tasks such as long-range retrieval and Scrolls QA, indicating robust implicit memory and retrieval capabilities. Overall, Gecko represents a scalable, memory-efficient approach to long-context LLMs with practical implications for real-world, multi-turn and reasoning-heavy applications, validated on 2 trillion-token pretraining and up to 4 million-token contexts.

Abstract

Designing a unified neural network to efficiently and inherently process sequential data with arbitrary lengths is a central and challenging problem in sequence modeling. The design choices in Transformer, including quadratic complexity and weak length extrapolation, have limited their ability to scale to long sequences. In this work, we propose Gecko, a neural architecture that inherits the design of Mega and Megalodon (exponential moving average with gated attention), and further introduces multiple technical components to improve its capability to capture long range dependencies, including timestep decay normalization, sliding chunk attention mechanism, and adaptive working memory. In a controlled pretraining comparison with Llama2 and Megalodon in the scale of 7 billion parameters and 2 trillion training tokens, Gecko achieves better efficiency and long-context scalability. Gecko reaches a training loss of 1.68, significantly outperforming Llama2-7B (1.75) and Megalodon-7B (1.70), and landing close to Llama2-13B (1.67). Notably, without relying on any context-extension techniques, Gecko exhibits inherent long-context processing and retrieval capabilities, stably handling sequences of up to 4 million tokens and retrieving information from contexts up to $4\times$ longer than its attention window. Code: https://github.com/XuezheMax/gecko-llm

Gecko: An Efficient Neural Architecture Inherently Processing Sequences with Arbitrary Lengths

TL;DR

Gecko tackles the fundamental challenge of long-context sequence modeling by building on the Megalodon backbone and introducing three core innovations: timestep decay normalization to stabilize long-horizon statistics, sliding chunk attention to maintain efficient local and cross-chunk coherence, and adaptive working memory with a position-aware online softmax to compress and retrieve information from very long histories. These components are tightly integrated to enable inherently long-context processing without explicit context-extension techniques, achieving superior training efficiency and long-context performance compared to Llama2 and Megalodon while handling sequences up to millions of tokens. Empirically, Gecko demonstrates lower training perplexity and stronger performance on short-context academic benchmarks, and substantial gains on long-context tasks such as long-range retrieval and Scrolls QA, indicating robust implicit memory and retrieval capabilities. Overall, Gecko represents a scalable, memory-efficient approach to long-context LLMs with practical implications for real-world, multi-turn and reasoning-heavy applications, validated on 2 trillion-token pretraining and up to 4 million-token contexts.

Abstract

Designing a unified neural network to efficiently and inherently process sequential data with arbitrary lengths is a central and challenging problem in sequence modeling. The design choices in Transformer, including quadratic complexity and weak length extrapolation, have limited their ability to scale to long sequences. In this work, we propose Gecko, a neural architecture that inherits the design of Mega and Megalodon (exponential moving average with gated attention), and further introduces multiple technical components to improve its capability to capture long range dependencies, including timestep decay normalization, sliding chunk attention mechanism, and adaptive working memory. In a controlled pretraining comparison with Llama2 and Megalodon in the scale of 7 billion parameters and 2 trillion training tokens, Gecko achieves better efficiency and long-context scalability. Gecko reaches a training loss of 1.68, significantly outperforming Llama2-7B (1.75) and Megalodon-7B (1.70), and landing close to Llama2-13B (1.67). Notably, without relying on any context-extension techniques, Gecko exhibits inherent long-context processing and retrieval capabilities, stably handling sequences of up to 4 million tokens and retrieving information from contexts up to longer than its attention window. Code: https://github.com/XuezheMax/gecko-llm
Paper Structure (39 sections, 24 equations, 5 figures, 3 tables)

This paper contains 39 sections, 24 equations, 5 figures, 3 tables.

Figures (5)

  • Figure 1: Negative log-likelihood for Gecko-7B, Megalodon-7B, Llama2-7B and Llama2-13B, w.r.t processed tokens during training.
  • Figure 2: Comparison of various sparse attention patterns. (a) Chunk-wise Attention: attention is restricted to separated chunks; (b) Sliding Window Attention: attention is restricted to fixed-size windows; (c) Sliding Chunk Attention: attention is assigned to both current and previous chunks.
  • Figure 3: Adaptive working memory in Gecko. (a) illustrates how the memory is updated by compressing information from all previous and current chunks. (b) depicts how Gecko stores contextual information in short- and long-term memories in different components.
  • Figure 4: PPL/NLL over long sequences. (a) shows the perplexity (PPL) in various context lengths. (b) plots averaged negative log-likelihood (NLL) broken down by token positions.
  • Figure 5: Evaluate Gecko-7B Long context on Passkey Retrieval and Needle in A Haystack.