💬 NLP

Attention Expansion: Enhancing Keyphrase Extraction from Long Documents with Attention-Augmented Contextualized Embeddings

This paper proposes an attention expansion mechanism that augments pre-trained language model token representations with information from surrounding out-of-context chunks, effectively enhancing keyphrase extraction performance on long documents without the computational cost of full-document attention or large language models.

Original authors: Roberto Martínez-Cruz, Alvaro J. López-López, José Portela

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Roberto Martínez-Cruz, Alvaro J. López-López, José Portela

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Tunnel Vision" of AI

Imagine you are trying to understand a massive 100-page mystery novel. You want to pull out the most important clues (keyphrases) that tell you what the story is about.

Current AI models (like the ones that power search engines or chatbots) are incredibly smart, but they have a "tunnel vision" problem. They can only read a small chunk of text at a time—say, 512 words—before they have to stop and reset. If the novel is 10,000 words long, the AI reads the first 512 words, forgets them, reads the next 512, and so on.

The issue: Important clues are often scattered. A character might be introduced on page 1, but their true importance isn't revealed until page 50. If the AI reads page 1 in isolation, it misses the connection to page 50. It's like trying to solve a puzzle while only looking at one piece at a time, never seeing how the pieces fit together.

The Old Solutions (and why they are too expensive)

To fix this, scientists have tried two main things:

Make the AI's "eyes" bigger: Build models that can read the whole book at once. This works, but it's like trying to fit a giant telescope into a bicycle. It requires massive amounts of computer power and memory, making it too slow and expensive for everyday use.
Use "Super-AIs" (LLMs): Use huge, general-purpose models that can read long texts. But these are like using a sledgehammer to crack a nut. They are slow and costly for simple tasks like just listing the main topics of a document.

The New Solution: "Attention Expansion"

The authors of this paper propose a clever middle ground. They call it Attention Expansion.

Think of the AI as a detective reading a specific room in a large mansion (the current chunk of text).

The Standard Way: The detective looks only at the furniture in that room.
The Attention Expansion Way: The detective still focuses on the current room, but they also have a quick, low-resolution sketch of the rooms immediately before and after it.

They don't re-read the whole mansion (which is slow). Instead, they use a lightweight, pre-made summary (called "Pre-trained Word Embeddings") of the surrounding text. Then, they use a special "magnifying glass" (a cross-attention layer) to glance at those sketches while reading the current room.

The Result: The detective can now say, "Ah, this chair in the current room makes sense because I saw a matching rug in the sketch of the previous room." The AI gets the benefit of seeing the whole picture without the cost of reading the whole book again.

How They Tested It

The researchers tested this idea on five different types of AI "brains" (models), ranging from general-purpose ones to ones specifically trained on science papers. They tested them on:

Long Scientific Papers: Where the main idea is often buried deep in the text.
News Articles: Where context shifts quickly.
Short Abstracts: To make sure the new method didn't break anything when the text was already short.

What They Found

It Works Everywhere: In almost every test (48 out of 50 scenarios), adding this "glance at the neighbors" made the AI better at finding the right keyphrases.
It Helps Even the "Smart" Models: Even models that were already designed to read long texts (like ModernBERT) got better with this trick. This proves the method isn't just fixing a broken model; it's adding extra helpful information that the model didn't have before.
It's Fast and Cheap: The "sketches" of the surrounding text are very lightweight. Adding this feature only increased the computer's workload by about 3.6%. It's a tiny price to pay for a big jump in performance.
It's Not Magic for Everything: While it worked great for science and news, it didn't help in every single case when switching between very different types of documents (like going from science to news), but it was still a strong improvement overall.

The Bottom Line

This paper introduces a way to give AI models "long-range memory" without making them slow or expensive. By letting the AI peek at a lightweight summary of the text surrounding what it's currently reading, it can understand long documents much better. It's a simple, efficient upgrade that makes existing AI tools smarter at finding the most important parts of a text.

Technical Summary: Attention Expansion for Long-Document Keyphrase Extraction

Problem Statement
Keyphrase extraction (KPE) is a critical task for document understanding, yet it faces significant challenges when applied to long documents. While Pre-trained Language Models (PLMs) have achieved strong performance in KPE due to their ability to generate rich contextualized representations, they are constrained by limited context windows (typically 512 tokens). In long documents, salient evidence for a keyphrase is often scattered across distant sections. Standard approaches split documents into non-overlapping windows and process them independently, preventing the model from capturing dependencies that span multiple windows.

While Long-Context Large Language Models (LLMs) and extended-context encoders (e.g., Longformer, ModernBERT) can process broader contexts, they incur high computational costs and memory usage, making them less practical for high-throughput KPE pipelines. Conversely, replacing compact specialized models with generic LLMs is often inefficient. The core problem addressed is how to provide specialized KPE models with access to long-document evidence without incurring the cost of full long-context attention or LLM-based inference.

Methodology: Attention Expansion
The authors propose an Attention Expansion mechanism that augments standard PLM token representations with information from surrounding "out-of-context" document chunks. The approach operates as follows:

Sliding Window Encoding: The input document is partitioned into chunks ( $c_i$ ) of length $L$ (e.g., 512 tokens). Each chunk is processed independently by the PLM to produce contextualized token representations ( $H_i$ ).
Out-of-Context Representation: Surrounding chunks ( $c_{i-w/2}, \dots, c_{i-1}, c_{i+1}, \dots, c_{i+w/2}$ ) are encoded using a lightweight Pre-trained Word Embedding (PWE) model (specifically Model2Vec). This generates static vector representations ( $S_i$ ) for the surrounding text without requiring a second forward pass through the transformer.
Cross-Attention Integration: A cross-attention layer allows tokens in the current chunk ( $H_i$ ) to query the PWE representations of the surrounding chunks ( $S_i$ ). The mechanism uses a single cross-attention head where queries are derived from $H_i$ , and keys/values are derived from $S_i$ . Notably, the key and value projections are tied (using the same matrix) to reduce parameters.
Enriched Classification: The output of the cross-attention is concatenated with the original PLM hidden states. This enriched representation is fed into a linear classification head for BIO (Begin, Inside, Outside) sequence tagging.
Multi-Head Variant: A multi-head variant (MH-AE) splits the hidden states and PWE representations into slices, applying independent cross-attention blocks to each slice before concatenation.

This design expands the effective context of the model (e.g., allowing access to up to 5x the native window size for 512-token backbones) without increasing the sequence length processed by the PLM itself.

Key Contributions

Novel Mechanism: Introduction of an attention expansion mechanism that bridges in-context PLM embeddings with out-of-context lexical representations via cross-attention, avoiding full long-context attention or LLM inference.
Efficiency: The method is parameter-light and computationally efficient. Parameter growth is bounded by the static embedding dimension, and the additional computation per token is linear. It adds minimal overhead (approx. 3.6% average forward-pass time increase) while significantly expanding the accessible context.
Complementarity: The mechanism provides complementary information rather than merely compensating for limited input length. It improves performance even on domain-specialized (SciBERT, KBIR) and native long-context (ModernBERT) encoders.

Experimental Results
The approach was evaluated across five encoder backbones (DistilBERT, SciBERT, KBIR, DeBERTa-v3, ModernBERT), two training regimes (SemEval-2010 and LDKP3K), and five benchmark corpora (Inspec, DUC-2001, SemEval-2010, LDKP3K, NUS).

Long Documents: Attention expansion consistently improved KPE performance across all evaluation settings. In 48 out of 50 backbone-dataset-training comparisons, the attention-expanded variants outperformed their baselines.
- General & Specialized Models: Improvements were observed for general-purpose models and those specialized for scientific domains (SciBERT) or keyphrase tasks (KBIR).
- Long-Context Models: Even ModernBERT, which natively supports 8,192 tokens, showed performance gains, indicating the mechanism adds value beyond just extending the window.
- Cross-Domain: Improvements were largely observed when transferring from scientific training data to news domains (DUC-2001), though results were not uniform across all backbones.
Short Documents: On the Inspec dataset (abstracts fitting within native limits), the mechanism did not degrade performance and, in some cases (e.g., ModernBERT), yielded significant gains, suggesting the component influences fine-tuning dynamics beneficially even without out-of-context evidence.
Cost: The multi-head variant (MH-AE) generally performed slightly better or comparably to the single-head variant (AE) with negligible additional parameter overhead (approx. 0.05% increase for MH-AE vs. 0.21% for AE).

Significance and Claims
The paper claims that attention expansion establishes an efficient and effective strategy for long-document KPE. It demonstrates that specialized models can be enhanced with distributed document evidence without the prohibitive costs of full long-context attention or LLM-based inference. The authors position this as a practical solution for high-throughput pipelines (e.g., large-scale indexing, retrieval-augmented generation) where efficiency is critical.

The results suggest that the mechanism provides complementary information that enriches token representations, rather than simply acting as a patch for short context windows. The work reinforces the argument that fast, efficient, task-specific components should complement, rather than be replaced by, larger generic reasoners. The authors modestly note that while the method improves transfer in many cases, improved cross-domain transfer is not automatic and depends on the interaction between the backbone and the target distribution. Future work is suggested to explore adaptive context selection and richer static representations.