Tagging-Augmented Generation: Assisting Language Models in Finding Intricate Knowledge In Long Contexts

AI summary

Overview

• Research area: Natural Language Processing / long-context large language model (LLM) question answering and prompting.
• Technical level: Intermediate. The paper assumes familiarity with LLMs, retrieval-augmented generation (RAG), context windows, named entity recognition (NER), and long-context benchmarks.
• Scope: The paper introduces Tagging-Augmented Generation (TAG), a lightweight prompting and input-augmentation technique that injects structured semantic tags into long documents, and evaluates it on two extended benchmarks (NoLiMa+ and NovelQA+) using two Anthropic Claude models.

What This Paper Is About

Even the strongest LLMs struggle to answer questions and reason over long, complex contexts, with accuracy dropping sharply well below their advertised maximum context lengths. Existing fixes such as retrieval-augmented generation (RAG) and chunk-based re-ranking are sensitive to chunking, embedding and retrieval choices, and require heavy pre-processing, knowledge acquisition and indexing. The authors propose TAG, which instead enriches the input text itself with semantic annotations so the model's attention is guided to the right spans, without changing the document's integrity or adding retrieval infrastructure.

Key Contributions

1. Semantic tagging and context-enhancement framework (TAG). A flexible, method-agnostic framework that generates semantic tags using different tagging mechanisms — LLM-based tagging, traditional NLP tools, or agentic systems — without model retraining and without architectural changes. The authors state this is the first framework that guides LLMs to perform better in long-context and complex reasoning tasks by semantically enriching the input.
2. Benchmark extension. The authors create NoLiMa+ and NovelQA+, tag-enriched versions of the NoLiMa and NovelQA benchmarks, purpose-built for evaluating whether LLMs can exploit structural cues in long contexts.
3. Empirical validation across two settings. Experiments span varying context lengths (250, 500, 16K, 32K tokens) and complexity (single-hop, multi-hop, detail) with two Anthropic Claude models, showing semantic tagging improves accuracy by over 17% for 32K token contexts and 2.9% in complex reasoning question answering for multi-hop queries relative to untagged baselines.
4. Comparison of tagging strategies. The paper contrasts tag definitions alone (TD) against tag definitions plus tagged context (TD+TC), and LLM-based privileged tagging against spaCy-based NER tagging, quantifying their effect on both peak accuracy and performance degradation rates.

Main Findings

• Long-context degradation is severe in the baseline. Without tagging, Claude 3.5 Sonnet fell from 81.19% accuracy at 250-token context to 32.67% at 32K tokens, a 59.77% extremum drop rate. Claude 3.7 Sonnet fell from 94.66% to 45.56%, a 51.87% extremum drop rate.
• Tag definitions in the prompt lift short-context accuracy. The TD setting raised Claude 3.5 Sonnet to 91.34% at 250 tokens (+10.15 gain) and Claude 3.7 Sonnet to 97.12% (+2.46 gain). Adding spaCy tagging yielded 88.77% and 95.11% respectively at 250 tokens.
• Tag definitions alone do not fix degradation. Claude 3.5 Sonnet still showed a 60.1% decline with TD, spaCy tagging showed 60.99%, and Claude 3.7 Sonnet 48.59% from peak performance to 32K tokens.
• Privileged tagging trades peak accuracy for long-context robustness. TD+TC privileged tagging scored slightly lower at 250 tokens (87.53% vs 91.34% for Claude 3.5; 95.21% vs 97.12% for Claude 3.7) but at 32K tokens Claude 3.7 Sonnet reached 52.39%, compared with 49.93% for TD alone and 45.56% for the baseline. The authors conclude tagging becomes increasingly effective as context length extends.
• Tagging helps complex narrative reasoning. On NovelQA+, Claude 3.7 Sonnet with spaCy TD+TC achieved the best single-hop (91.09%) and multi-hop (56.97%) scores, and Claude 3.5 Sonnet reached 78.97% on detail questions versus its 73.36% baseline. The paper reports a 2.9% gain in multi-hop complex reasoning.
• Both reasoning and non-reasoning models benefit. The two Claude models (3.7 Sonnet with the thinking feature enabled, 3.5 Sonnet v2 at temperature 0) both improved, answering the paper's RQ3.
• Leaderboard context. The authors state that, at the time of publication, their best-performing model ranks #2 on the NoLiMa leaderboard for 32K context.
• Post-processing cost of truncation. Because many novels exceeded the Claude context-window limit, the authors truncated books and filtered out questions from omitted sections, leaving 1,035 questions across complexity levels.

Methodology in Plain English

The TAG pipeline starts by splitting an input document into multi-sentence chunks using a configurable strategy — sentence-level, paragraph-level, or semantic chunking. Chunks are de-duplicated, then passed to tagging modules. Tags are inserted directly into the text using XML-style markup, for example <Person>Marie Curie</Person>, so the document keeps its original structure while carrying explicit attention cues.

Two families of taggers are used. The first is LLM-based tagging, split into an information-extraction (IE) approach, where the model extracts entities and assigns each to a predefined semantic category and writes the tag inline, and a classification-based approach, where each tag category is treated as a class and the model outputs matching tags for later post-processing. The two are merged into a unified set of tagged chunks to improve coverage and robustness. The second family is traditional NER using spaCy, which identifies 18 standard entity types and applies a nested XML tagging scheme for segments containing multiple entity types.

For evaluation, the authors build NoLiMa+, where haystack chunks matching candidate semantic categories are marked up using privileged tags derived from NoLiMa's needle keyword list, and NovelQA+, where novels are chunked and tagged with spaCy entities using broader narrative-oriented semantic categories. Three prompt settings are compared: a baseline vanilla RAG setting with no tagged context and no tag definitions, tag definitions only (TD), and tag definitions with tagged context (TD+TC). Models are invoked through Amazon Bedrock. NoLiMa scoring gives 1 if the generated response contains the golden answer and 0 otherwise; NovelQA uses exact-match accuracy on the model's single-character multiple-choice output (A, B, C or D).

Why This Matters

The work targets a practical failure mode: LLMs that claim very long context windows but degrade badly on real long documents, especially when the query and the relevant passage share little or no literal wording. TAG offers an interpretable, low-latency, infrastructure-free alternative to retrieval pipelines, and tags can be cached so previously tagged context can be retagged without additional compute overhead.

Real-world applications the paper's framing supports:

• Technical support and document analysis, where models must handle previously unseen, unstructured, extensive documents without a pre-indexed retrieval system.
• Technical document understanding, where fine-grained reasoning over long specifications or manuals is required.
• Complex narrative comprehension, such as reasoning about character relationships, plot development and thematic connections distributed across full-length novels.
• Dialogue summarization and other multi-document tasks that demand reasoning over increasingly long and complex inputs.

Industry relevance: the technique requires no model fine-tuning, no architecture change and no external knowledge base, so it can be dropped into existing enterprise LLM pipelines as a preprocessing and prompting step. The authors also note a concurrent, independent Anthropic work on context engineering that supports the idea of using XML-based tags to segment context.

Future Directions

• Broader domain validation. The paper's evaluation is limited to question answering in a synthetic setting (NoLiMa+) and a literary setting (NovelQA+); broader testing across diverse tasks and technical domains is needed.
• Semantic category design. TAG's effectiveness depends on choosing appropriate semantic categories, which the authors flag as potentially challenging in specialized domains or low-resource languages.
• Agentic tagging at inference time. The experiments primarily used spaCy-based and LLM-based tagging; comprehensive evaluation of agentic tagging approaches applied on the fly during inference is left as an important direction.
• Fundamental memory limits. TAG operates within existing context windows and cannot address the underlying memory limitations of current LLMs, leaving room for complementary approaches.

Target Audience

Applied NLP and ML engineers building long-context question-answering or RAG pipelines; LLM researchers studying context utilization, attention degradation and the "lost-in-the-middle" problem; and enterprise teams that need to improve document-heavy workflows (technical support, document analysis, narrative comprehension) without deploying retrieval infrastructure or retraining models. Readers should already be comfortable with LLM prompting, benchmark evaluation, and named entity recognition concepts.

Authors’ abstract

Recent investigations into effective context lengths of modern flagship large language models (LLMs) have revealed major limitations in effective question answering (QA) and reasoning over long and complex contexts for even the largest and most impressive cadre of models. While approaches like retrieval-augmented generation (RAG) and chunk-based re-ranking attempt to mitigate this issue, they are sensitive to chunking, embedding and retrieval strategies and models, and furthermore, rely on extensive pre-processing, knowledge acquisition and indexing steps. In this paper, we propose Tagging-Augmented Generation (TAG), a lightweight data augmentation strategy that boosts LLM performance in long-context scenarios, without degrading and altering the integrity and composition of retrieved documents. We validate our hypothesis by augmenting two challenging and directly relevant question-answering benchmarks -- NoLima and NovelQA -- and show that tagging the context or even just adding tag definitions into QA prompts leads to consistent performance gains over the baseline -- up to 17% for 32K token contexts, and 2.9% in complex reasoning question-answering for multi-hop queries requiring knowledge across a wide span of text. Additional details are available at https://sites.google.com/view/tag-emnlp.