Elman, Long Context LLMs Are Reaching New Heights: A Dive into the Latest Architectures and Evaluation Methods
This newsletter explores the cutting-edge advancements in deep learning architectures designed for long-context language modeling. We'll delve into novel attention mechanisms, efficient KV cache management, and robust evaluation metrics, highlighting the key innovations that are pushing the boundaries of LLM capabilities. From tensorized attention to retrieval heads and specialized metrics, this newsletter provides a comprehensive overview of the latest research, offering insights into the evolving landscape of long-context language processing.
Tensorized Attention: Reshaping the Landscape of Long Sequence Modeling
Long Sequence Modeling with Attention Tensorization: From Sequence to Tensor Learning by Aosong Feng, Rex Ying, Leandros Tassiulas https://arxiv.org/abs/2410.20926
Caption: The figure illustrates the process of attention tensorization, showing the transformation of query (Q), key (K), and value (V) vectors into higher-order tensors. Subfigure (a) depicts the reshaping and operations involved in creating tensorized attention, while (b) provides a tensor diagram comparing full attention with the proposed tensorized attention and its low-rank version, highlighting the reduced computational complexity.
Large language models (LLMs) have revolutionized NLP, but processing long sequences remains a challenge due to the quadratic computational cost of traditional attention mechanisms. This paper introduces Tensorized Attention, a novel approach that reshapes long sequences into compact tensor representations, enabling efficient modeling of long-range dependencies.
Instead of attending to all tokens directly, Tensorized Attention folds the one-dimensional interaction into a higher-order tensor. Each dimension of this tensor is then modeled with short-range interactions, effectively reducing the interaction distance between tokens and converting out-of-window interactions into within-window ones. This is achieved by generalizing the input at the attention layer from a sequence to a higher-order tensor, replacing traditional vector operations with their tensor counterparts.
The proposed mechanism employs a multi-hop attention process, sequentially updating the value tensor along each dimension. This is mathematically equivalent to a Kronecker decomposition of full attention and can be expressed as: A = ⊗ₘᵢ₌₁ softmax (QᵢKᵀᵢ/√d) • Mᵢ, O = A × V. Here, Q, K, and V are the tensorized query, key, and value matrices, respectively, and Mᵢ represents the attention mask for the i-th dimension. An efficient custom Triton kernel ensures hardware compatibility and speed. Furthermore, tensorized positional encoding allows for exponential extrapolation of context length along specific dimensions.
Adapting pre-trained LLMs like Llama-8B, Mistral-7B, and OpenLlama-3B with continued pre-training using Tensorized Attention yielded impressive results. Llama-8B, trained with a context length of 32,768, could extrapolate to 128k during inference with an 11x speedup compared to full attention using FlashAttention-2. The method also showed improved perplexity on Proof-pile and strong performance on downstream tasks like HellaSwag, SIQA, and Natural Questions. On LongBench, Tensorized Attention consistently outperformed full attention, achieving an average speedup of 0.61x.
Theoretically, the tensor representation better captures the hierarchical and low-rank structure of attention matrices, enabling more efficient low-rank approximation in tensor space. Empirical analysis of attention spectra from RoBERTa and ViT supports this claim, demonstrating faster decay of singular values in tensor space. This suggests fewer parameters are needed to recover the same information, leading to more efficient computation.
Unraveling the Mysteries of Synthetic Context Extension
Understanding Synthetic Context Extension via Retrieval Heads by Xinyu Zhao, Fangcong Yin, Greg Durrett https://arxiv.org/abs/2410.22316
Caption: This figure visualizes the correlation between retrieval score cosine similarity and F1 scores for three long-context tasks (MDQA, MuSiQue, and SummHay Citation) across different data realism and diversity levels. The correlation coefficients (R) indicate a strong positive relationship between retrieval score similarity and performance, especially for SummHay Citation, suggesting the importance of retrieval heads in long-context LLM performance.
Training long-context LLMs on extensive real-world data is computationally expensive. Synthetic context extension, which involves fine-tuning LLMs with synthetically generated long-context data, offers a more efficient alternative. This paper investigates the effectiveness of this approach and the underlying mechanisms that influence its success. The researchers examined three long-context tasks: multi-document question answering (MDQA), multi-hop question answering (MuSiQue), and document citation (SummHay-Citation).
The study systematically varied the realism of target information ("needles") and the surrounding context ("haystack") diversity within the synthetic datasets. This ranged from using LLMs to generate realistic data to employing templated relations and symbolic datasets.
A key finding is the crucial role of retrieval heads, specialized attention heads responsible for retrieving information from long contexts. Models trained on less effective synthetic data exhibited fewer retrieval heads, often subsets of those learned on realistic or high-quality synthetic data. The similarity between retrieval heads learned on synthetic and real data strongly correlated with downstream performance, indicating that learning specific retrieval heads is necessary, though not sufficient, for effective context extension. For example, on MuSiQue, Llama 3 showed a 2-4% gap between synthetic and real data performance, while on MDQA the gap was much larger at 33%. The recall of retrieval heads on synthetic MuSiQue data correlated strongly with F1 on the real task (R=0.81).
Intervention experiments further validated the importance of retrieval heads. Masking their activations significantly reduced performance, while masking random heads had minimal impact. Patching activations from the intersection of retrieval heads learned on real and poorly performing synthetic data improved the latter's performance. This suggests that while synthetic data can activate the necessary retrieval heads, it doesn't train them as effectively as real data. The retrieval score S<sub>h</sub> for head h is calculated as: S<sub>h</sub> = |G<sub>h</sub> ∩ y| / |y*|*, where G<sub>h</sub> represents the set of tokens retrieved by head h and y** is the answer span. This research provides a mechanistic understanding of synthetic context extension, offering insights into creating better synthetic data for long-context LLM training.
GraphLSS: A Novel Graph-Based Approach for Long Document Summarization
GraphLSS: Integrating Lexical, Structural, and Semantic Features for Long Document Extractive Summarization by Margarita Bugueño, Hazem Abou Hamdan, Gerard de Melo https://arxiv.org/abs/2410.21315
Caption: The top graph depicts the correlation between the optimized class weight for relevant sentences and the ROUGE-1 F1 score, showing how adjusting the weight improves summarization performance. The bottom graph illustrates the decrease in class weight over training epochs as the model learns to identify relevant sentences more effectively, as described in the GraphLSS paper.
Existing graph-based models for long document summarization often rely on external tools or complex architectures. GraphLSS presents a novel heterogeneous graph construction that simplifies this process by integrating lexical, structural, and semantic features. It defines nodes as sentences and semantically rich words (nouns, verbs, and adjectives) and connects them with four edge types: sentence order, sentence semantic similarity, word-sentence association, and word semantic similarity. This intuitive design captures crucial document relationships without external learning models.
Evaluated on PubMed and arXiv datasets using a heterogeneous GAT model, GraphLSS addressed imbalanced extractive labels using a weighted cross-entropy loss with adaptive class weights for relevant sentences, calculated as: λ_(i+1) = λ_i − (1 − log(τ))/T, where τ represents the portion of sentences predicted as relevant. This dynamic weighting proved crucial for optimizing performance on skewed datasets.
GraphLSS outperformed all compared approaches on both datasets in ROUGE-1, -2, and -L scores. Achieving scores of 63.57/36.91/55.32 (R-1/R-2/R-L) on PubMed and 55.14/23.00/50.83 on arXiv, the results highlight the effectiveness of the integrated feature approach. The study also emphasized the impact of labeling strategies on summarization, demonstrating that GraphLSS, even with pre-labeled data, surpasses recent non-graph models.
Removing word-in-sentence edges significantly impacted performance, highlighting the importance of cross-granularity interactions. While sentence edges were also influential, word similarity edges had a lesser impact due to their lower representation in the graph. Analysis of precision and recall revealed balanced performance on PubMed, while recall exceeded precision on arXiv, suggesting potential for refinement in handling additional retrieved text.
Beyond Perplexity: Rethinking Evaluation for Long-Context LLMs
What is Wrong with Perplexity for Long-context Language Modeling? by Lizhe Fang, Yifei Wang, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, Yisen Wang https://arxiv.org/abs/2410.23771
Caption: This image illustrates the calculation of Log Probability Gain (LPG) for tokens in a sentence, comparing probabilities based on long and short contexts. "Buddy" is identified as a key token due to its high LPG (2.08), signifying its dependence on the long context. This concept is central to the LongPPL metric, which focuses on these key tokens to better evaluate long-context language models.
Perplexity (PPL), a standard metric for evaluating language models, falls short when assessing long-context capabilities. This paper proposes LongPPL, a novel metric designed specifically for long-context evaluation. The authors argue that PPL's weakness lies in its equal weighting of all tokens, which obscures the importance of key tokens crucial for long-context understanding.
LongPPL identifies these key tokens using causal intervention on context length, computing the Log Probability Gain (LPG) for each token: LPG(xᵢ) = log Pθ(xᵢ|lᵢ) - log Pθ(xᵢ|sᵢ), where lᵢ represents the full long context and sᵢ is a truncated short context. Tokens with high LPG are considered key tokens as they benefit significantly from the extended context. Additionally, Log Probability Value (LPV), LPV(xᵢ) = log Pθ(xᵢ|lᵢ), filters out tokens that are difficult to predict even with long context. LongPPL is then calculated as the perplexity computed solely on these selected key tokens.
Experiments demonstrated a strong correlation between LongPPL and performance on long-context benchmarks like LongBench, LongEval, and RULER. LongPPL achieved a Pearson correlation of -0.96 with LongBench scores, significantly outperforming traditional PPL, which showed almost no correlation. The authors also introduce LongCE (Long-context Cross-Entropy), a training objective that upweights key tokens during fine-tuning, leading to consistent performance improvements across benchmarks, with a maximum accuracy gain of 22% on LongEval.
SHADOWKV: Optimizing KV Cache for High-Throughput Inference
ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference by Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, Beidi Chen https://arxiv.org/abs/2410.21465
Caption: This diagram illustrates the decoding phase of SHADOWKV. The query interacts with landmarks to select top-k chunk IDs, which are checked against the low-rank pre-ROPE key cache. Cache misses trigger fetches from the value cache, while hits allow for reconstruction of the full key in the current KV cache using the pre-ROPE representation and applying RoPE.
Efficiently serving long-context LLMs is challenging due to the growing memory footprint and access requirements of the key-value (KV) cache. SHADOWKV addresses these limitations by leveraging the low-rank properties of pre-Rotary Position Embedding (ROPE) keys and strategically managing the value cache.
SHADOWKV employs a two-phase approach: pre-filling and decoding. During pre-filling, it performs low-rank decomposition of the pre-ROPE key cache using Singular Value Decomposition (SVD) and stores these compact representations on the GPU. The value cache is offloaded to the CPU. Landmarks (chunk means) of the post-ROPE key cache and outlier chunks are stored on the GPU to facilitate accurate KV selection during decoding. The decoding phase uses these landmarks to select relevant KV chunks. Corresponding values are fetched from the CPU, and the selected key cache is reconstructed on-the-fly. CUDA multi-streams optimize this process, overlapping key cache reconstruction and value fetching.
The system's efficiency is analyzed using equivalent bandwidth: 2SB<sub>GPU</sub> / (S/C + 2(K+O)C + (1-α)KC B<sub>GPU</sub>/B<sub>PCle</sub>), where S is sequence length, C is chunk size, K is selected chunk budget, O is outlier count, α is cache hit rate, B<sub>GPU</sub> is GPU bandwidth, and B<sub>PCle</sub> is PCIe bandwidth. SHADOWKV achieves a theoretical equivalent bandwidth higher than the A100 memory bandwidth, accelerating attention computation.
Experiments across various LLMs and benchmarks demonstrate SHADOWKV's effectiveness. It reduces GPU memory footprint by over 6x while maintaining accuracy with a minimal sparse KV cache budget (1.56%). In large batch serving scenarios on an A100 GPU, it supports up to 6x larger batch sizes and boosts throughput by up to 3.04x compared to full attention, exceeding theoretical infinite batch size performance.
Conclusion: A Converging Focus on Efficiency and Effectiveness in Long-Context LLMs
This newsletter has highlighted the exciting progress being made in addressing the challenges of long-context language modeling. From novel attention mechanisms like Tensorized Attention to efficient KV cache management systems like SHADOWKV, the focus is clearly on improving both computational efficiency and model performance. Furthermore, the development of new evaluation metrics like LongPPL underscores the need for more accurate and insightful assessments of long-context capabilities. The interplay between these advancements promises to unlock the full potential of LLMs for a wide range of applications requiring extended context understanding.