Do not be greedy, Think Twice: Sampling and Selection for Document-level Information Extraction

AI summary

Overview

• Research area: Natural Language Processing, specifically document-level Information Extraction (DocIE) and test-time scaling (TTS) for decoder-only large language models.
• Technical level: Intermediate. The paper assumes familiarity with LLM prompting, decoding strategies (greedy vs. sampling), and reward-model training, but the core idea is straightforward.
• Scope (one sentence): The paper argues that greedy decoding under-serves document-level information extraction and proposes a sampling-and-selection framework (ThinkTwice) with one unsupervised and one supervised selector, evaluated on MUC-4, MultiMUC and BETTER Granular.

What This Paper Is About

Document-level information extraction asks a model to read a whole document and produce a structured template containing the entities, relations and events of interest. Because LLM outputs vary between runs, researchers typically use greedy decoding so that comparisons are reproducible — but greedy decoding returns only the single most likely token at each step, not necessarily the best template. This paper shows that sampling many candidate templates and then selecting among them yields substantially better extractions than greedy decoding, especially with reasoning-oriented models, and builds selection mechanisms to exploit that gap.

Key Contributions

1. The ThinkTwice framework. A two-step inference pipeline in which a model π_θ is prompted with annotation guidelines g and a document x_i to produce N candidate templates T_i = {t̂_i,1, …, t̂_i,N}, after which a selector S(g, x_i, T_i) chooses the preferred output. Constrained decoding to a predefined JSON schema is used during sampling.
2. F1 Voting, an unsupervised selector. Instead of the binary match used by Majority Voting, each candidate is scored by its aggregate F1-style similarity to all other candidates, and the highest-scoring candidate is chosen. This addresses the failure mode where Majority Voting selects an empty template for complex IE outputs.
3. A supervised reward-based selector. A preference dataset P is built from the silver candidate set together with the gold annotation, keeping pairs where one template scores higher under the dataset metric than another. The reward model r_θ is trained with the Bradley-Terry objective using an informed margin m = λ·(φ(t^c_i, t_i) − φ(t^r_i, t_i)), with λ set to 3 as suggested by Touvron et al. (2023). At inference the highest-scoring candidate is selected.
4. A rejection-sampling method for silver reasoning traces. Because golden reasoning trajectories do not exist for DocIE datasets, the authors sample N synthetic reasoning–template pairs, rank them with the dataset's scoring function φ against gold templates, keep the top K, build a silver dataset D̂, and iterate until the maximum selection score stops improving; a final aggregation step fine-tunes a newly reinitialized model on the top R pairs.

Main Findings

• Reasoning models beat non-reasoning models in zero-shot DocIE. Comparing Llama and Qwen with and without reasoning traces under greedy decoding, reasoning variants performed consistently better with only a few exceptions. The gap was especially pronounced on BETTER, where the paper reports that constrained decoding is insufficient for non-reasoning models to produce good-quality templates. The authors hypothesize this stems from reasoning models' training on long-form sequential dependencies. All subsequent analyses focus exclusively on reasoning models.
• Sampling plus selection beats greedy decoding. On MultiMUC, F1 Voting improved roughly two F1 points over greedy across both backbones; on BETTER it exceeded three points for Qwen3. Under greedy decoding, MultiMUC/BETTER scores were 12.67/14.78 for Llama R1 and 14.65/16.12 for Qwen3; with ThinkTwice and F1 Voting they rose to 14.41 ± 0.33 / 15.29 ± 0.88 and 16.56 ± 0.37 / 19.45 ± 0.16 respectively.
• Majority Voting is unreliable for IE. On BETTER with Llama R1, Majority Voting scored 10.51 ± 2.30, below the greedy baseline of 14.78. F1 Voting showed lower variance than Majority Voting in all datasets and models. In supervised settings, Majority Voting gains were marginal or even detrimental (Llama R1: 28.41 ± 1.49 vs. greedy 28.52), while F1 Voting remained consistent (36.56 ± 0.22).
• Supervised selection sets a new state of the art on MUC-4. The reward selector reached 41.11 ± 0.48 with Llama R1 and 42.51 ± 0.72 with Qwen3, versus 35.20 for IterX (T5-enc large), 32.30 for GTT (BERT base) and 28.30 for TempGen (BART large). Greedy Qwen3 reached 35.97; greedy Llama R1 reached 28.52.
• A large gap to the oracle remains. Oracle (best-of-N) scores were 31.77 (MultiMUC) and 34.08 (BETTER) for Llama R1, 35.32 and 36.74 for Qwen3, and 56.60 / 58.32 on supervised MUC-4. Even though much smaller open models can beat proprietary systems once sampling and selection are used, the unsupervised and supervised selectors stay well below these upper bounds.
• Reasoning-trace coverage was incomplete. The rejection-sampling procedure obtained gold-quality reasoning traces for only 67.77% (Llama R1) and 71.46% (Qwen3) of the training data.
• Cross-lingual transfer works from English alone. Fine-tuning Qwen3 on English data only (MUC-4) improved the other languages by an average of 4.53 F1 points. Adding the supervised selector provided 6.53 additional points for the zero-shot model and 6.41 for the English fine-tuned model. With both forms of supervision, the method beat systems trained directly on target-language data (Gantt et al., 2024) in four of five languages, with Korean the exception. The English fine-tuned reward configuration averaged 25.95 ± 0.74 across the five non-English languages (MUC-4 English score 42.51 ± 0.72), versus 22.35 for Gantt et al. (2024) and 17.90 for GPT 5.5 zero-shot. Korean remained hardest overall, with GPT 5.5 at 07.59 and the best ThinkTwice configuration at 09.62 ± 0.14.
• Model scale alone is not enough. The paper notes that greedy decoding on modern models can perform below older, less capable LLMs, and that performance improves significantly when sampling is coupled with an effective selection strategy.

Methodology in Plain English

The authors frame DocIE as a template-generation task: each example is a triplet (g, x_i, t_i) of annotation guidelines, document, and gold template. For prompting, the guidelines go into the system prompt as Markdown and the document goes into the user prompt, with the model expected to emit JSON. Llama R1 received a single combined user prompt because system prompts are not recommended for it.

Inference happens twice over the same input: first the LLM generates 64 candidate templates per document, then a selector picks one. Three selectors are compared. Majority Voting picks the most frequent candidate. F1 Voting picks the candidate with the highest summed F1 similarity to the others. The reward selector scores candidates with a trained reward model built on top of Qwen3 Embedding 8B plus a feed-forward projection head that maps the embedding to a single scalar.

For supervised training, the researchers generate reasoning traces by sampling, score each generated template against the gold one with the dataset metric, keep the best ones, and retrain. They repeat this loop until the top score stops improving, then do one final aggregation pass and fine-tune a freshly reinitialized model on the top-ranked pairs. Discarded low-quality candidates are reused with the selected ones to build the preference pairs that train the reward model, which turns the selection problem into something closer to classification.

All experiments use N = 64 samples per instance as a trade-off between exploration and cost, with a maximum reasoning budget of 4000 tokens and bootstrapping over three random seeds. MUC-4 evaluation uses the CEAF-RME metric from Chen et al. (2023). BETTER Granular was used exclusively for zero-shot evaluation because of its longer documents and guidelines; the paper does not report supervised BETTER results, and the size of the training split is not reported.

Why This Matters

The paper reframes output variability in LLM-based extraction as an asset rather than a nuisance, and shows that decoding strategy can matter as much as model choice. For research, it connects DocIE to the broader test-time scaling literature, which has mostly targeted math and reasoning benchmarks, and it supplies a reusable recipe for building reasoning-training data where no golden traces exist. It also questions the common practice of reporting only greedy-decoded numbers as a model's true capability.

Real-world applications implied by the datasets and event types studied:

• News and open-source intelligence monitoring, since MUC-4 covers geopolitical conflict reports with Attack, Arson, Bombing, Kidnapping, Robbery and Forced Work Stoppage event types.
• Humanitarian and disaster response, drawing on BETTER Granular's coverage of disasters, epidemics, displacement and migration.
• Corruption and governance monitoring, using BETTER's corruption and protest templates.
• Multilingual media analysis, since MultiMUC provides parallel annotations in Arabic, Farsi, Korean, Russian and Chinese alongside English, and the method transfers from English-only training.

Industry relevance centers on cost: sampling 64 candidates raises latency and decoding cost relative to greedy decoding, but the paper frames the sample count as a tunable hyperparameter rather than a fixed constant, allowing a trade-off between computation and quality. The supervised selector also lets teams exploit automatically collected supervision instead of expensive human reasoning annotations, and the finding that English-only training transfers to four of five languages reduces the need for per-language annotation.

Future Directions

• Better selectors. The persistent gap to oracle performance on every dataset is the clearest open problem; the authors describe their supervised selector as relatively simple and see room for improvement.
• Reinforcement learning for reasoning traces. The authors suggest GRPO as a way to train reasoning LLMs for DocIE more effectively than rejection sampling, potentially complementing sampling and selection.
• Reducing incomplete trace coverage. With gold-quality traces obtained for only 67.77% and 71.46% of training data, noisy training signals remain a limitation to be addressed.
• Managing inference overhead. Generating multiple candidates increases latency and decoding cost; finding better cost-performance trade-offs is left open.

Target Audience

Researchers and practitioners working on information extraction, template filling, or structured generation with LLMs; engineers building document-processing pipelines that need reliable structured output; and anyone interested in test-time scaling methods outside of mathematical reasoning. Readers who want the implementation details can find the code at the repository the paper links (github.com/MikelZubi/ThinkTwice), and the work is published under a CC BY 4.0 license as arXiv:2601.18395v3.

Authors’ abstract

Document-level Information Extraction (DocIE) aims to produce an output template with the entities, relations, and events of interest occurring in the given document. Standard practices include prompting decoder-only LLMs using greedy decoding to avoid output variability. Rather than treating this variability as a limitation, we show that sampling can produce substantially better solutions than greedy decoding, especially when using reasoning models. We thus propose ThinkTwice, a sampling and selection framework in which the LLM generates multiple candidate templates for a given document, and a selection module chooses the most suitable one. We introduce both an unsupervised method that exploits agreement across generated outputs, and a supervised selection method using reward models trained on labeled DocIE data. To address the scarcity of golden reasoning trajectories for DocIE, we propose a rejection-sampling-based method to generate silver training data that pairs output templates with reasoning traces. Our experiments show the validity of unsupervised and supervised ThinkTwice, consistently outperforming greedy baselines and the supervised state-of-the-art.