Source-linked AI summary
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
Long-context sequence modeling requires architectures that avoid Transformer-style quadratic complexity and weak length extrapolation. Gecko extends the Mega and Megalodon design with timestep decay normalization, sliding chunk attention, and adaptive working memory, achieving stronger controlled 7B comparisons and inherent long-context processing without explicit context extension.
Problem
Efficiently processing sequential data with arbitrary context lengths remains challenging because Transformer designs have quadratic complexity and weak length extrapolation.
Method
Gecko builds on Megalodon by adding timestep decay normalization, sliding chunk attention, and adaptive working memory for long-context pretraining and inference.
Results
Gecko-7B reaches a training loss of 1.68, versus Llama2-7B at 1.75 and Megalodon-7B at 1.70, while handling sequences up to 4 million tokens and retrieving from contexts 4× longer than its attention window.
Takeaways & Limitations
Without explicit context-extension mechanisms, Gecko supports efficient long-context processing and retrieval while improving training perplexity and downstream benchmark performance over Llama2 and Megalodon.
Takeaways & Limitations
Megalodon’s chunk-wise attention exhibits boundary losses and limited access to historical information outside its context, motivating Gecko’s long-term memory design.
Abstract
from arXiv · showhide
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
1. Introduction
Gecko targets efficient, inherently long-context sequence processing by extending Mega and Megalodon with mechanisms for normalization, local context, and long-term memory. Controlled evaluations report stronger language-modeling performance and long-context behavior than the compared baselines.
- Transformer-based long-sequence modeling is limited by quadratic computational complexity and weak length extrapolation.
- Gecko is designed to process sequences with unlimited context length and outperform Transformer architectures on real-world language modeling.
- Table 1 compares Gecko at 1.3B and 7B scales with open-source base models while reporting model size, context length, and pretraining tokens.
- Gecko builds on the Megalodon backbone with timestep decay normalization, sliding chunk attention, and adaptive working memory.These components address cumulative-statistics control, chunked context access, and information outside sliding chunks.
- Gecko-7B outperforms Llama2-7B and Megalodon-7B on training perplexity and downstream benchmarks under controlled comparisons.
2. Background: Backbone Architectures in Mega
Mega and Megalodon combine exponential-moving-average sequence processing with gated attention, but their normalization and chunk-wise attention designs constrain long-context utilization. Megalodon’s chunking provides linear complexity while leaving boundary and historical-context limitations.
- 2.1 CEMA: Complex Multi-dimensional Damped EMA: CEMA expands each input dimension into h dimensions and applies damped exponential moving average in a complex hidden space.The hidden state is projected back to a one-dimensional output for each input dimension.
- 2.2 Timestep Normalization: Megalodon’s timestep normalization computes cumulative means and variances from autoregressive input groups across previous timesteps.
- 2.2 Timestep Normalization: As timestep t increases, the current mean and variance contribute less to cumulative statistics, limiting Megalodon’s inherent scaling to long sequences.
- 2.3 Chunk-wise Normalized Gated Attention: Chunk-wise gated attention partitions queries, keys, and values into chunks and applies attention independently with linear complexity O(nc).
- 2.3 Chunk-wise Normalized Gated Attention: Megalodon’s chunk-wise attention has higher losses near chunk boundaries and poorly captures historical information outside the attention context.The latter limitation particularly affects in-context retrieval-oriented tasks.
3. Gecko
Gecko combines timestep decay normalization, sliding chunk attention, and adaptive working memory to improve efficient long-context sequence modeling. Its mechanisms address normalization control, chunk-boundary context loss, and retention of information beyond local attention windows.
- 3.1 Timestep Decay Normalization: Timestep decay normalization fixes the ratio of current mean and variance contributions in cumulative statistics through decay hyperparameters β1 and β2.The method applies initialization bias correction to both cumulative statistics and uses hardware-friendly GPU implementation.
- 3.2 Sliding Chunk Attention: Sliding chunk attention computes attention over each current chunk and its preceding chunk, propagating local context across chunk boundaries.Unlike sliding window attention’s per-token shifting, SCA performs chunk-level matrix multiplications for more efficient GPU execution.
- 3.2 Sliding Chunk Attention: SCA provides lossless short-term memory through previous-chunk keys and values, but its constrained receptive field still limits long-range dependency capture.The paper therefore motivates a separate long-term memory mechanism for information beyond sliding chunks.
- 3.3 Adaptive Working Memory: Linear-attention compressive memory has bounded capacity and can suffer memory collisions, while gated alternatives selectively forget historical information.Gecko’s AWM is designed to retain long-term memory rather than deliberately forgetting historical information.
- 3.3 Adaptive Working Memory: Adaptive working memory compresses information from current and all previous chunks into fixed-size memory through a position-aware online softmax kernel.The kernel integrates historical information into its denominator, while memory updates operate chunk by chunk.
- 3.3 Adaptive Working Memory: Gecko removes the normalization term by constraining it to one when using the position-aware online softmax feature kernel.The supplied formulation states τ_s = 1 for every chunk.
4. Experiments
Gecko is evaluated against matched baselines on pretraining, short-context benchmarks, long-sequence modeling, retrieval, and long-context tasks. Across these evaluations, it shows stronger language-modeling and long-context behavior, including retrieval beyond its nominal attention window.
- Pretraining Results: 1.68 training loss: Gecko-7B outperforms Llama2-7B (1.75) and Megalodon-7B (1.70), nearly matching Llama2-13B (1.67).The comparison uses the same number of processed tokens and reports lower NLL throughout training.
- Short-Context Evaluation: Under identical 2T-token pretraining conditions, Gecko-7B consistently outperforms Llama2-7B and Megalodon-7B across all reported academic benchmarks.On several tasks, its performance is comparable to or exceeds Llama2-13B.
- Long-Context Modeling: Gecko shows better context utilization than Megalodon as validation context length increases from 4K to 4M tokens.Figure 4a measures perplexity over 1,500 books, each containing sequences of at least 4M tokens.
- Long-Context Modeling: Gecko’s loss steadily decreases with increasing context length, whereas Megalodon’s loss rises near chunk boundaries.Figure 4b reports average NLL by token position up to 32K.
- Long-Range Retrieval: With only 4K attention context, Gecko achieves near-perfect passkey retrieval up to 16K sequence length and 100% NIAH performance at 8K across all depths.The reported retrieval results do not use context-extension methods.
- Long-Context QA: Table 2 evaluates Gecko and comparable long-context models on NarrativeQA, Qasper, and QMSum using task-specific F1 or geometric ROUGE metrics.The evaluation uses 0-shot, 2-shot, and 1-shot settings respectively, with a unified prompt format.
5. Conclusion
Gecko extends the Megalodon backbone with mechanisms for efficient long-context pretraining and inference. The architecture reports stronger language modeling and robust processing and retrieval over contexts far beyond its nominal attention window.
- Conclusion: Gecko combines timestep decay normalization, sliding chunk attention, and adaptive working memory to support efficiently processing effectively unbounded contexts.These components are integrated into an architecture built on the Megalodon backbone.
- Conclusion: Without explicit context-extension mechanisms, Gecko handles sequences up to 4 million tokens and retrieves information from contexts 4× longer than its nominal attention context.The conclusion also reports consistent improvements in training perplexity and downstream benchmark performance over Llama2 and Megalodon.
Appendix A. Adaptive Working Memory
The appendix presents recurrence and output expressions for adaptive working memory when combined with sliding chunk processing.
- Adaptive Working Memory: The adaptive working-memory derivation applies the delta rule and defines the output of each chunk.The displayed derivation begins with an initial memory state M0 = 0.
- Adaptive Working Memory: When incorporating sliding chunks, the derivation resets the auxiliary memory state before defining the chunk output.The passage specifies M′1 = 0.
B.1 Efficient Fused CUDA Operators Implementation
The implementation uses fused CUDA kernels and parallel associative scans to accelerate recurrent sequence processing while reducing long-context activation memory. Chunk-level rematerialization and hardware-aware execution make the operators more scalable on GPUs.
- Parallel Associative Scan: The recurrent computation is reformulated as a parallel associative scan, reducing sequence-modeling span from O(L) to O(log L) on parallel hardware.The scan operates on complex affine tuples whose associative operator composes multiplicative and additive terms.
- Boundary Handling: Boundary masks encode document resets algebraically, allowing uniform execution when the effective multiplicative term is zero.When Mt = 0, the state resets to ht = pt.
- Kernel Architecture: The CUDA kernel maps features to the grid, assigns one warp per feature, and processes sequence chunks with warp-level scans.This mapping reuses time-invariant parameters and avoids shared-memory round trips for intra-chunk dependencies.
- Memory Hierarchy: Coalesced float4 loads and shared-memory tiling improve memory throughput for complex affine tuples and prefetched input chunks.The vectorized representation stores the real and imaginary components of q and p in aligned 128-bit transactions.
- Backward Pass: 32× activation-memory reduction: chunk-level rematerialization stores only chunk boundaries during the forward pass and reconstructs intermediate states during backpropagation.The backward pass reruns the prefix scan from the nearest stored boundary before the adjoint scan.
Appendix C. Experimental Details
Table 3 compares Gecko-1B with OLMo1-1B on standard academic benchmarks at matched training checkpoints of approximately 2T tokens.
- Table 3 compares the performance of Gecko-1B and OLMo1-1B on standard academic benchmarks.
- The comparison uses matched training checkpoints at approximately 2T tokens.
- The table reports model size, context length (CTX), and total data tokens used during pretraining.
C.1 Experiment on Gecko-1B
The Gecko-1B experiment compares against OLMo1-1B using matched training conditions, while noting that Gecko was trained on fewer total tokens.
- Gecko-1B was trained for 2T tokens on Dolma v1.7, whereas OLMo1-1B was trained for 3T tokens.
- The comparison matched OLMo1-1B’s global batch size, defined as the number of tokens per optimizer step.
- Despite less training, Gecko-1B achieved comparable results to published OLMo1-1B benchmark performance.
- Gecko-1B obtained better results on PIQA, ARC-Easy, and ARC-Challenge.
- Like OLMo1-1B, Gecko-1B fully anneals between 3T and 4T tokens.