Unlocking Long-Sequence LLMs: Position Embeddings, Scaling, and Efficient Attention

Position Embedding

Early long‑sequence models used either ALIBI or RoPE for positional encoding. Recent models (e.g., LLaMA, Mistral, Cohere) adopt RoPE as the default because it has a solid mathematical basis, derives relative positions from absolute indices, and integrates smoothly with flash‑attention kernels when numerical precision is handled carefully.

ALIBI was initially claimed to enable lossless extrapolation, but empirical studies show over‑fitting once training reaches a token count on the order of 1 T. ALIBI also lacks fine‑tuning tricks such as RoPE‑NTK and is incompatible with flash attention without materializing bias masks, creating a performance bottleneck for ultra‑long sequences.

RoPE Scaling / Extrapolation

Industry practice for extending context length relies on RoPE scaling , which is closely related to NTK scaling. Typical implementations increase the RoPE angle θ during long‑sequence fine‑tuning. Notable extensions include Dynamic NTK, Yarn, LongLoRA, PoSE, LLM‑infinite, Focus Transformer, ReRoPE, and LogN scaling. Empirical experience indicates that plain RoPE scaling remains the most stable and easy‑to‑apply method.

https://arxiv.org/html/2405.14591v1

https://zhuanlan.zhihu.com/p/717174366

Thoughts on Position Encoding

RoPE‑NTK‑based variants often suffer performance drops on very long sequences, suggesting a bottleneck in extrapolation. Benchmarks such as Ruler show that models like Gemini and Jamba‑1.5B perform comparatively better, likely due to architectural variations or hybrid designs.

Recent work proposes Contextual Position Encoding , which blends contextual semantics with positional information, reminiscent of early DeBERTa and Mamba selective‑SSM approaches. The main challenges are memory efficiency and compatibility with flash‑attention.

https://arxiv.org/abs/2405.18719

Attention

Long‑sequence attention research focuses on two complementary goals:

Reducing attention compute cost (e.g., linear attention, state‑space models).

Preserving attention entropy on very long inputs (e.g., LogN scaling).

Cheaper Attention

Early linear‑attention methods such as Linformer, Linear Transformer, and Sparse Transformer were popular before 2022 but exhibited stability issues at scale. Dense attention has become more cost‑effective thanks to flash attention and ring attention.

Effective dense variants include GQA, MQA, MLA, and hybrid designs like GAU and its derivatives Mega/Mega2, which retain quadratic complexity while lowering compute overhead. Hybrid architectures that combine sparse and dense attention, or integrate Mamba with dense layers, demonstrate strong extrapolation beyond 1 M tokens.

Entropy

As token length grows, attention distributions flatten, increasing entropy. Multiplying attention logits by a scalar makes the softmax output sparser, helping the model focus on salient tokens.

LogN scaling, described in a blog by “苏神”, has been applied in QWEN models. Experiments show that LogN scaling and RoPE scaling are not fully additive; combined use does not always outperform each method alone, likely due to hyper‑parameter interactions.

https://spaces.ac.cn/archives/9444

The Differential Transformer reduces attention noise by subtracting two softmax outputs, improving performance on both short and long sequences and on many‑shot evaluations.

https://arxiv.org/abs/2410.05258

Long‑Sequence Data & Evaluation

During pre‑training, careful proportioning of long‑context data and efficient training pipelines are essential.

https://arxiv.org/abs/2402.10171

Instruction‑tuning for long contexts often relies on synthetic data, as exemplified by the LLaMA 3 approach.

https://ai.meta.com/research/publications/effective-long-context-scaling-of-foundation-models/

Some studies suggest concatenating short texts yields more reliable evaluation than using raw long sequences.

https://arxiv.org/abs/2410.02660

Traditional metrics such as perplexity and simple needle‑in‑haystack tests do not capture true long‑context capability. The Ruler benchmark, which uses partially synthetic data, provides a more informative measure, though performance still degrades on lengths like Ruler‑128k. Retrieval‑augmented generation shows promise for answering with very long contexts.

https://arxiv.org/html/2410.03227v1

Long‑Sequence Training Frameworks

Flash and Ring Attention

State‑of‑the‑art long‑sequence training relies on Flash Attention and Ring Attention kernels. For contexts around 128 k tokens, sequence parallelism is often unnecessary because batch sizes are small.

Sequence Parallelism

When sequence length exceeds ~256 k tokens, sequence parallelism may be required. DeepSpeed Ulysses can toggle between sequence and tensor parallelism, but sharding is limited by the number of attention heads (e.g., max sharding = 32 when num_heads = 32). Llama 3 proposes a context‑parallel + all‑gather KV scheme that splits the attention matrix by rows; the reduced KV size from GQA/MQA keeps communication costs manageable.

Other Considerations

Combining these frameworks with variable‑length masks, sparse attention, or multimodal prefix‑LMs adds engineering complexity. In JAX environments, flash‑attention support lags behind PyTorch, and custom kernels such as Pallas may underperform native GPU kernels, requiring extensive debugging.

Inference‑Related Topics

The primary inference bottleneck for long contexts is KV‑cache memory consumption. Reducing KV size while preserving attention quality is critical.

Key works on KV compression:

https://arxiv.org/abs/2306.14048 (H2O)

https://arxiv.org/abs/2404.14469 (SnapKV)

https://arxiv.org/abs/2404.11912 (speculative decoding / cache offloading)

https://arxiv.org/abs/2408.11049 (dynamic KV selection)

Static KV selection during the prefilling stage can degrade performance when conversation topics shift dramatically. Dynamic selection and offloading strategies aim to mitigate this issue.