Writing in the Margins: Better Inference Patterns for Long-Context Retrieval

## 1 Introduction

The performance of Large Language Models (LLMs) tends to deteriorate when processing extensive inputs, a limitation linked directly to their fixed context window and attention mechanisms [li2023looglelongcontextlanguagemodels, liu2023lostmiddlelanguagemodels]. In particular, LLMs struggle with tasks involving long contexts, especially when the relevant information is embedded in larger volumes of text [bai2024longbenchbilingualmultitaskbenchmark, shaham2023zeroscrollszeroshotbenchmarklong]. Recent research thus highlights the importance of improving model capabilities to handle more extensive datasets without losing accuracy or requiring exponential increases in computational resources.

There have been various attempts to extend the usable context window of LLMs, such as sparse attention [tworkowski2023focused, chen2024longloraefficientfinetuninglongcontext, mohtashami2023landmarkattentionrandomaccessinfinite], length extrapolation [dai2019transformerxlattentivelanguagemodels, su2023roformerenhancedtransformerrotary, peng2023yarnefficientcontextwindow], and context compression [ge2024incontextautoencodercontextcompression, mu2024learningcompresspromptsgist]. Concurrently, the field has witnessed the rise of sophisticated prompting strategies like Chain of Thought (CoT) and related structured reasoning methods [wei2023chainofthoughtpromptingelicitsreasoning, yao2023treethoughtsdeliberateproblem, Besta_2024]. These approaches have significantly enhanced LLMs’ ability to tackle complex tasks by systematically guiding the reasoning process through predefined structural patterns.

Our work bridges the gap between efficient transformers architecture research and development of new prompting strategies. Specifically, we identify a novel key-value (KV) cache aware reasoning pattern for existing off-the-shelf long context window LLMs in scenarios typical of retrieval-oriented tasks, where the context is substantial and the instructional prompt is comparatively short. We begin by recognizing that long-context prompts are commonly prefilled in the KV cache segment-wise in a process known as chunked prefill. From this insight, we introduce an inference pattern called Writing in the Margins (WiM), which concurrently generates query-based extractive summaries at each step of the prefill that are subsequently reintegrated at the end of the computation. We term these intermediate outputs “margins”, drawing inspiration from the practice of making margin notes for improved comprehension of long contexts in human reading. Using methodologies similar to “scratchpad” techniques, which meticulously record step-by-step calculations, we incorporate margin notes into the final segment predictions. We show that this technique, which adds only minimal additional computation, significantly enhances long context comprehension. The WiM pattern can also provide end-users with real-time insights into computational progress through streamed margin notes, which ultimately help make AI decisions more transparent and explainable. This can enable users to (1) pinpoint the location of essential information and (2) reduce computational load by exiting early if the provided information satisfactorily addresses the query.

In Figure [1](#S0.F1), we provide an illustrative example of WiM inference, which we encourage readers to reference as a practical demonstration to complement the formal algorithm description that will be presented in the following sections.

Our main contributions are as follows:
{tasks}
\taskWe introduce a new inference pattern, Writing in the Margins (WiM), which achieves better performance on long-context window tasks with a relatively minor increase in computational cost.

We demonstrate the application of WiM within an interactive long context retrieval setup, effectively increasing the transparency of the process and reducing the first response latency.

We provide an implementation of this inference pattern using the Hugging Face Transformers library.

## 2 Writing in the Margins

##### Chunked Prefill

Figure: Figure 2: Chunked Prefill. Example of how the attention mask is set across different chunks during prefill iterations (first chunk at the top, second chunk at the bottom). Each new chunk needs to retain causality while attending to all previous chunks. Chunked prefill is mathematically equivalent to prefill without chunking.
Refer to caption: /html/2408.14906/assets/assets/Export_01.png

Typically, the process of inference for generative LLMs consists of two principal phases: the prefill phase and the decoding phase. When an LLM is requested to prefill a substantial prompt—in the range of hundreds of thousands of tokens—it is common practice to prefill the KV cache in chunks [agrawal2023sarathiefficientllminference]. This method is known as chunked prefill and is supported by many inference frameworks, including vLLM [vllmchunkedprefill].

Chunked prefill divides the prompt into fixed-size chunks to populate the KV cache at each layer of the Transformer model [vaswani2023attention]. The rationale for chunked prefill is to reduce overall memory usage, as the quadratic memory complexity of the attention mechanism during prefilling can be prohibitive for larger prompts. By splitting a prompt of length $L$ into $N$ chunks, each of size $K$, where $N=L/K$, the overall memory complexity of prefilling is reduced from $O(L^{2})$ to $O(LK)$. The attention mask must be adjusted to allow each new chunk to attend to all tokens in the previous chunks while maintaining the causal structure only for the new chunk, as illustrated in Figure [2](#S2.F2).

Our work exploits the chunked prefill mechanism to generate intermediate “margins” that can then be appended to the prompt to better guide the model toward performing a specific task.

##### Writing in the Margins

Consider a prompt $P$, composed of a context $C$, and an instruction $I$.

$$ $\text{P}=\text{C}+\text{I}$ (1) $$

Prefilling a decoder-only transformer model $T$ directly with the entire prompt $T(P)$ is computationally inefficient when the prompt is long. Moreover, as shown in [liu2023lostmiddlelanguagemodels], processing the entire prompt in one go can lead to mid-sequence forgetting.

To make this process more efficient, we implement the prefill technique described in the previous paragraph, where the context C is divided into $N$ segments.

$$ $\text{C}=c_{1}+c_{2}+...+c_{N}$ (2) $$

For the first segment, the model $T$ operates on chunk $c_{1}$, resulting in output that includes past key values $\text{pkv}_{1}$. The model continues onto the second segment with the $\text{pkv}_{1}$ cached, i.e., $T(\text{pkv}_{1},c_{2})$, effectively emulating the scenario of processing $T(c_{1}+c_{2})$ in one step.
As the procedure progresses, each sequential chunk, $c_{k}$, is processed with prefilled past key values, noted as $T(\text{pkv}_{[1..k-1]},c_{k})$, mimicking an uninterrupted run of $T$ on C.

The Writing in the Margins (WiM) strategy addresses potential mid-sequence forgetting issues by appending an extractive instruction $I_{A}$ to each chunk, enhancing chunk-specific outputs. It transforms each step into $T(\text{pkv}_{[1..k-1]},c_{k}+I_{A})$, where the instruction $I_{A}$ is embedded alongside each context chunk, then dropped from the KV cache before the next chunk prefilling. The instruction $I_{A}$ is closely related to $I$ - the model is asked to copy over all relevant to $I$ information.

Intermediate outputs from each chunk are referred to as margin notes $M_{i}$, cumulatively forming $N$ notes, described as $M=M_{[1..N]}$. Unhelpful notes, perhaps irrelevant to the instruction, are discarded, enhancing the final contextual construct to $\text{C}+\text{M}+\text{I}$, positioned advantageously towards the end to minimize mid-sequence memory loss. Intuitively, the model is allowed to use relevant intermediate predictions while answering the final query.

To summarize, we modify the chunked prefill algorithm by adding extra decoding steps (green in Table [1](#S2.T1)). Most of these steps can be efficiently batched with the original prefill steps. The query-relevant information extracted from these steps is then added at the end of the context but before the instruction.

| step | Chunked Prefill | WiM | keep |
| --- | --- | --- | --- |
| 1 | $T(\emptyset,c_{1})^{\dagger}$ | $T(\emptyset,c_{1})^{\dagger}$ | $\text{pkv}_{[1]}$ |
| 2 | $T(\text{pkv}_{[1]},c_{2})^{\dagger}$ | $T(\text{pkv}_{[1]},c_{2})^{\dagger}$ | $\text{pkv}_{[1..2]}$ |
|  |  | ${\color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}T(\text{pkv}_{[1]},I_{A})^{\dagger\ddagger}}$ | ${\color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}M_{1}}$ |
| ⋮ | ⋮ | ⋮ | ⋮ |
| N | $T(\text{pkv}_{[1..N-1]},c_{N})^{\dagger}$ | $T(\text{pkv}_{[1..N-1]},c_{N})^{\dagger}$ | $\text{pkv}_{[1..N]}$ |
|  |  | ${\color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}T(\text{pkv}_{[1..N-1]},I_{A})^{\dagger\ddagger}}$ | $M_{N-1}$ |
| N + 1 |  | ${\color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}T(\text{pkv}_{[1..N]},I_{A})^{\dagger\ddagger}}$ | $M_{N}$ |
| N + 2 | $T(\text{pkv}_{[1..N]},I)^{\dagger\ddagger}$ | $T(\text{pkv}_{[1..N]},{\color[rgb]{0,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,.5,.5}M_{[1..N]}}+I)^{\dagger\ddagger}$ |  |

Table: Table 1: Batching Chunked Prefill Steps with WiM margin generation. The inference for generative LLMs consists of two principal phases: the prefill phase (†) and the decoding phase (‡). WiM algorithm adds extra decoding steps that mostly can be batched with chunked prefill steps. We keep margin notes $M_{i}$ produced in extra steps (green) as plain text. We then prefill the model $T$ with all relevant notes $M_{[1..N]}$ before the final instruction $I$.

Algorithms [1](#alg1) and [2](#alg2) present pseudocodes of inference with chunked-prefill and WiM respectively.

##### Generating and classifying margins

It is possible to use the same instance of the model for both generating the margins and classifying them, without affecting the prefilled KV cache. Having generated the margins, it is possible to inference the model with a classification prompt without using the KV cache (past_key_value in Algorithm [2](#alg2)) generated by the prefilling operation. This way, the model will act as it had never been prefilled. Having classified the margins, it is possible to reuse the previously prefilled KV cache to append the classified margins and then generate the final output. The overhead of classifying a single margin, in terms of memory, is just the KV cache size of a single margin and the classification prompt, which is negligible compared to the prefilled long-context prompt.
It is also possible to overlap the generation of margins with their classification using the same model instance and the same request in the batch (see Appendix LABEL:appendix-kv-cache for more details).

Figure: Algorithm 1 Inference with Chunked Prefill

Figure: Algorithm 2 Writing in the Margins

Figure: Figure 3: Design Comparison. Three inference designs for managing long context windows: (Top Left) Long Context LLM (LLM): This design feeds all context, without segmentation, directly to the model. (Top Right) Retrieval-Augmented Generation (RAG): Segments are selected based on a retrieval method (e.g., cosine similarity between vector representations of the query and the segment). All selected segments, along with the task instruction, are then concatenated and fed to a model. (Bottom) Writing in the Margins (WiM): The context is divided and processed segment by segment. At each step, the model is prompted to produce auxiliary information from each segment. This information is then classified and, if deemed positive, it is incorporated into the final step before the task description.
Refer to caption: /html/2408.14906/assets/x2.png

## 3 Experimental Setup

### 3.1 Datasets

Following the RULER [hsieh2024rulerwhatsrealcontext] task categories, we measure the performance of an inference pattern on three types of skills: (I) Multi-Hop Reasoning, (II) Needle Retrieval/ Single-Hop Reasoning, and (III) Aggregation. Table [2](#S3.T2) presents the curated long context datasets used to benchmark all LLMs:

| skill | benchmark | context length | # |
| --- | --- | --- | --- |
| type | name | (tokens) | examples |
| I | MultiHop-RAG ^$\clubsuit$ [tang2024multihopragbenchmarkingretrievalaugmentedgeneration] | 13-32k | 100 |
| I | HotpotQA ^† [yang2018hotpotqadatasetdiverseexplainable] | 16k/ 32k /64k | 100/ 100/ 100 |
| II | SQuAD ^† [rajpurkar2018knowdontknowunanswerable] | 16k/ 32k/ 64k | 100/ 100/ 100 |
| III | CWE ^† [hsieh2024rulerwhatsrealcontext] | 64k | 100 |

Table: Table 2: Datasets We have curated four datasets to evaluate long context window LLMs. Each set consists of $100$ examples, generated either using the RULER code (†) or by subsampling the longest examples from the original benchmark data ($\clubsuit$).

In the following paragraph, we briefly introduce the benchmarks used in each category, describe our curating rationale, and hypothesize the impact of WiM on each skill set.

- I
Multi-Hop QA
The task aims to check the behavior of tracing entities with multi-hop connections based on the HotPotQA and MultiHop-RAG benchmarks [tang2024multihopragbenchmarkingretrievalaugmentedgeneration, yang2018hotpotqadatasetdiverseexplainable]. We used the RULER codebase(^1^11RULER code: [https://github.com/hsiehjackson/RULER](https://github.com/hsiehjackson/RULER)) to generate a subset of $100$ examples based on HotPotQA - a multi-hop queries sourced from Wikipedia articles. Following RULER, we simulated long context retrieval scenarios by generating examples in three length variants: $16k$, $32k$, $64k$. We also selected the $100$ longest examples in the range of $13k$-$33k$ tokens from MultiHop-RAG - a large collection of multi-hop queries based on English news articles.
Hypothesis Transformers are known for not being able to emulate a for loop [zhou2023algorithmstransformerslearnstudy]. In WiM we simulate going through the context twice which can improve the performance by aggregating all interconnected facts in one place at the end of the document.
- II
Needle Retrieval/ Single-Hop Reasoning
In the context of a long context window, the Needle Retrieval and Single-Hop QA task can be jointly seen as a kind of filter benchmark, where the task is to filter irrelevant content and either copy or transform the relevant information. We used the RULER code to generate examples based on SQuAD [rajpurkar2018knowdontknowunanswerable] in three context length variants: $16k$, $32k$ and $64k$, collecting $100$ datapoints in each variant.
Hypothesis The reduction using the WiM inference pattern is in fact a reverse engineering of how filter-type benchmarks are created. The model is asked to filter out injected distractions and copy relevant parts into the margin in each segment-wise prediction step.
- III
Aggregation
This type of task measures a model’s ability to aggregate relevant information that spans long-range context. We use the Common Words Extraction (CWE) benchmark [hsieh2024rulerwhatsrealcontext] where words are sampled from discrete distributions, with the number of common words fixed and increasing with the sequence length number of uncommon words. Following RULER, we generated $100$ examples of average length $64k$ tokens. We scaled original frequencies of common and uncommon words to match the extended length of samples: common words appear $500$ times in the sample, while uncommon words do not appear more than $50$ times.
We modified the task instruction to include the number of word occurrences in text to support aggregation over segments.
Hypothesis The performance of WiM in reduction tasks will be related to the concept of hierarchical reduction. We can think of a summarization task as being a summary of segment summaries, or any aggregate task as being a result over cumulative aggregation over segments.

### 3.2 Long Context Window LLMs

We selected seven off-the-shelf models that officially support context windows up to $128k$ tokens:

- •
Phi-3-small-128k-instruct(^2^22Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone. [abdin2024phi3technicalreporthighly])
- •
Qwen2-7B-Instruct(^3^33Qwen2 Technical Report. [yang2024qwen2technicalreport])
- •
Meta-Llama-3.1-8B-Instruct(^4^44The Llama 3 Herd of Models. [dubey2024llama3herdmodels])
- •
Phi-3-medium-128k-Instruct[2](#footnote2)
- •
Palmyra-4-Chat-128K(^5^55Writer’s proprietary model.)
- •
Meta-Llama-3.1-70B-Instruct[4](#footnote4)
- •
Qwen2-72B-Instruct[3](#footnote3)

In all experiments, we used half precision models with identical sampling parameters — specifically, a temperature setting of $0.0$ and $2k$ maximum new tokens. We used 0 0-shot prompts for all benchmarks. In MultiHop-RAG, HotPotQA and SQuAD experiments, we applied the same model-independent prepossessing step: we used nltk [bird2009natural] to split the context into sentences, then grouped them in segments no longer than $4096$ tokens. This resulted in $4-16$ margin notes per datapoint. In CWE, where the datapoints contain only numbered words, we exchanged nltk for naive words split by space and used $8192$ segment length, which gave on average $8$ margins per sample. We chose to count tokens using GPT-4 tiktoken tokenizer(^6^66 [https://github.com/openai/tiktoken](https://github.com/openai/tiktoken)) since this choice does not favour any of the evaluated models’ tokenizers.

In each case, we measured the relative differences of WiM pattern scores with respect to the following two baselines:

- •
Long Context LLM (LLM) - all context without segmentation is fed to the LLM.
- •
Retrieval Augmented Generation (RAG) - segments are selected based on a retriever (ex. cosine similarity between vector representations of the query and the segment), then all selected segments and the task instruction are concatenated and fed to an LLM.

All three inference patterns, including WiM, are presented in Figure [3](#S2.F3).

In order to make the results more comparable, we replaced the retriever in RAG with the classifier used in WiM. We expect the RAG results to be lower in the real RAG systems (especially for longer segment lengths), as vectorization is a form of lossy compression.

### 3.3 Prompting

For all benchmarks, we respected their original formulation. In all cases, the prompt strategy for the Long Context LLM baseline could be expressed as:

Where system_message and instruction were usually the task instructions split into two parts and appended before and after the main context respectively.

In the RAG approach, we used the original prompt but replaced context with all relevant segments concatenated by a newline sign.

In WiM inference, all constructed prompts shared the common prefix:

This was necessary for the efficient reuse of the KV cache. To ensure that predictions were comparable, we manually identified a promising prompt for the margin generation and final prediction steps for all evaluated models.

#### 3.3.1 Margin generation

For each intermediate context $\texttt{context}_{i}=\Sigma^{i}_{1}c_{i}$ and instruction $I$, we used the following extractive summary prompt $I_{A}$ to generate a margin note $M_{i}$:

In our experiments, the margin generation step was combined with the classification step; the first token generated was a class label. We conditioned the generation of a margin based on the first token; i.e., we continued the generation only if the first token was YES. Additionally, the prompt included an explanation designed to enforce specific formatting and to prevent the model from inserting comments before delivering its judgment.

In Appendix LABEL:appendix-kv-cache, we explore the possibility of decoupling margin generation and classification prompts while using the same instance of the model.

#### 3.3.2 Final WiM prompt with accumulated margins

To distinguish the content of the margins from the original context, and to maintain the document’s logic and structure, we explicitly named the writing-in-the-margins strategy in our last step, while aggregating all relevant margin notes.

We used two variants of the prompt, depending on the number of retrieved margins.

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 10
presentation: 4
questions: 1) One approach the authors could explore would be to use a separate smaller LLM as classifier. Using the base model (which can be very large) adds latency.
soundness: 3
strengths: The paper provides a number of thought provoking outcomes.
1) It showcases how adding a simple strategy of adding notes or summaries in the "margins" after each prefilled chunk can assist in improving LLM reasoning and retrieval capabilities. 
2) The notes written by the LLM can potentially be used to improve explainability of the final decoded output. This is dependent on whether the question asked for the margin generation is useful. In the paper the authors ask the LLM whether the context is relevant to the query (and to provide a summary).
3) The approach is general purpose, it can be applied to any LLM without the need for finetuning which is a big win. 

Overall, strong contribution.
summary: The paper introduces a new inference methodology called "writing in margins" for long context tasks. The method builds upon the chunked prefill strategy (commonly used while dealing with long contexts to avoid the quadratic growth of memory), dividing long input contexts into manageable segments and generates "margins" or intermediate summaries for each chunk. 
The margins are then classified by the same LLM as useful or not-useful and useful margins are kept as part of the context and used during decoding step. 
The approach seems to significantly help LLM (especially smaller LLMs) in better accuracy during decoding.
weaknesses: 1) Latency - while the authors mention that latency is slightly increased, an ablation study for this would be welcome. Since the paper uses 2 steps for each chunk - margin generation and then margin classification, you are effectively doing 2 decoding steps for the model with each chunk. This will add latency, especially if the summaries generated are long.

2) comparison against finetuned models - the paper mentions that this technique the models to perform well on tasks (long context) without the need to finetune the model (similar to rag). It would be good to include a model finetuned for the task and using the standard Long Context LLM decoding approach.

---

confidence: 2
contribution: 1
initial_rating: 3
presentation: 2
questions: * The main results in Table 4 present many metric values repeatedly measured using a fixed dataset and multiple algorithms.  No statistical significance tests are shown.  This severely compromises the integrity of the results. Were these tests conducted—with appropriate corrections for multiple comparisons—but not reported?
soundness: 1
strengths: * **Interesting approach to query-specific representation expansion.**  Bootstrapping decision-making with model information (e.g., writing the margins) is a compelling way for a model to guide itself toward a better response. 
* **Focus on effectiveness and efficiency.** The authors discuss both the effectiveness of their method and how it can improve the efficiency during decoding.  
* **Extensive experimentation.** Notwithstanding concerns (below), testing the approach on multiple settings and across multiple models is a rigorous way to test a model.  The authors could have improved the discussion on how performance varies and what that implies about the proposed method.
summary: The authors present a method for improving the representation of chunked text in a prompt by computing query-specific representations (margin notes) for each chunk. They hypothesize that this expanded and query-specific text allows for more efficient and effective decoding. To test this, the authors apply their method to several baseline models across three tasks: multi-hop reasoning, single-hop retrieval, and aggregation.  Post-hoc analysis involves an ablation study.
weaknesses: * **No formal statement of hypotheses.** This is perhaps implicit, but given the number of experiments, it is essential to be explicit about the precise hypotheses the experiments test.  As best I can tell, one hypothesis is that treatment with margin notes will be better than treatment with other methods (LLM and RAG baselines) across a fixed condition (e.g., length variant, task).  There are some allusions to other hypotheses (e.g., comparisons across columns), but that's less clear.  This is important because of the next point.
* **No formal hypothesis tests.** There are a lot of numbers in Table 4+.  Results in bold seem to be the max within some context.  However, it's not clear if any of these differences are (a) statistically significant and/or (b) if those tests have accounted for multiple comparisons (since these datasets are being reused...a lot).  Without this, it's difficult to understand the robustness of these results.  In order to address this, you can consult the literature on significance testing (Cohen's "Empirical Methods for Artificial Intelligence" is good; tutorials from the RecSys/information retrieval communities are also good) and correcting for multiple comparisons (see those tutorials from the RecSys/information retrieval communities).
* **Writing falls off at the end.** Starting with the ablation experiments (Section 5), the flow and writing of the paper weaken.  Why do these ablation experiments make sense?  What are the implications?  What is the argument of Section 6?  How are all of these things connected to the core hypothesis of the paper?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Larger models seem to profit less from WiM (table 4), and you do not include models larger than 70B. Would models larger than 70B still see improvements with WiM? Can you discuss this in more detail?
2. RAG is best with SQuAD in many cases, and almost always better than WiM. You argue that with multihop Q&A this is no longer the case (as shown in table 4), but isn't this only true for your RAG implementation / approximation, and more sophisticated RAG systems would improve this score?
soundness: 2
strengths: - interesting and original idea
- comparison with several base lines
- improvements over these baselines
summary: The authors propose and investigate the usage of intermediate information (margins) for improving long-context retrieval. They compare different small and medium-size LLMs as well as a RAG like system and find improvements over these baselines in many cases.
weaknesses: - comparison and discussion not complete, as larger models (which show less improvements) and more sophisticated RAG systems are not included

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to the "Weaknesses".
soundness: 3
strengths: * The authors introduce a novel inference pattern called "Writing in the Margins" (WiM) that leverages the chunked prefill mechanism in large language models to generate intermediate "margin notes" that can guide the final prediction. This is a clever way to address the challenges of long-context processing in retrieval-oriented tasks.

* The results show that WiM can significantly boost the performance of off-the-shelf models across a range of long-context benchmarks, including multi-hop reasoning and aggregation. This demonstrates the effectiveness of the proposed approach.
summary: The paper presents a new inference pattern called "Writing in the Margins" (WiM) that addresses the challenges of processing long input contexts in retrieval-oriented tasks. WiM leverages the chunked prefill mechanism in large language models to generate intermediate "margin notes" that summarize relevant information for the given query. These margin notes are then incorporated into the final response, leading to significant performance boosts on benchmarks like HotpotQA and Common Words Extraction compared to vanilla long-context models and retrieval-augmented approaches. The paper also discusses how WiM can enhance the transparency and interactivity of the retrieval process by providing users with real-time insights into the model's reasoning.
weaknesses: * The experimental setup could be expanded to include more baselines, such as state-of-the-art models specifically designed for long-context processing to better assess the relative performance of WiM.

* While the results are strong, the paper could benefit from a deeper analysis of why WiM works well for some tasks (e.g., multi-hop, aggregation) but not as consistently for others (e.g., single-hop QA). Understanding the underlying mechanisms behind these performance differences would strengthen the contributions.

Meta-review: This paper proposes a new method called Writing in the Margins (WiM) to improve inference for long-context retrieval. The strategy presented can be applied to different foundation models and requires minimal additional computation. Although a reviewer gave a very high score, after reading the reviews, discussions, and the paper itself, I think this paper has the following key weaknesses: 1) The experimental section lacks sufficient baseline comparisons, which was pointed out by multiple reviewers who provided different suggestions; 2) The current test samples are too few. Although multiple datasets are used, only 100 cases were selected from each dataset, which is insufficient to validate and demonstrate the method's effectiveness; 3) There is a lack of comparative and analytical experiments. While the authors responded during the rebuttal phase, the issues mentioned above were not addressed. Therefore, I think this paper is not ready for publication by far.

Decision: Reject