Retrieval-Augmented Training › A Worked Example: Chunk-by-Chunk Retrieval

To make the chunked retrieval process concrete, walk through how RETRO processes a passage about climate science.

Suppose the model is processing the following 256-token passage (represented schematically as four 64-token chunks):

• Chunk 1: "The Earth's average surface temperature has risen approximately 1.1 degrees Celsius since the late 19th century. This warming is primarily driven by increased atmospheric concentrations of greenhouse gases such as carbon dioxide and methane..."
• Chunk 2: "Carbon dioxide levels have increased from 280 parts per million in pre-industrial times to over 420 parts per million today. This increase is primarily attributable to the burning of fossil fuels, deforestation, and industrial processes..."
• Chunk 3: "The consequences of this warming include rising sea levels, more frequent extreme weather events, and disruptions to agricultural systems. Arctic sea ice extent has declined significantly, and permafrost regions are beginning to thaw..."
• Chunk 4: "International agreements such as the Paris Agreement aim to limit warming to 1.5 degrees Celsius above pre-industrial levels. Achieving this target requires substantial reductions in global greenhouse gas emissions by 2030..."

At the start of processing, the model has seen no context, so Chunk 1 is processed without any retrieved neighbors. This is unavoidable: there is no prefix to query against.

When the model begins processing Chunk 2, it encodes Chunk 1 as a retrieval query. The query embedding is a dense vector representation computed from the tokens in Chunk 1. The nearest-neighbor search finds database chunks that discuss similar material: passages about temperature measurements, greenhouse gas warming, historical climate records. These retrieved chunks, along with their continuations, become the CCA input for Chunk 2. The model can now predict tokens in Chunk 2 with access to related facts from elsewhere in the training corpus, for example data tables of historical temperature records or scientific explanations of the greenhouse effect that may not appear verbatim in the current passage.

When processing Chunk 3, the query is formed from the full prefix up to the end of Chunk 2. Now the model has more context to work with: the passage is clearly about climate science, specifically focusing on CO2 concentrations and their sources. The retrieved neighbors are likely to be even more relevant, potentially including passages about sea level rise data, extreme weather attribution studies, or Arctic observation reports.

This chunk-by-chunk retrieval rhythm means the model gets progressively better retrievals as the sequence unfolds and the topic becomes clearer. The first chunk is the weakest point, but for any passage of reasonable length, the majority of tokens are predicted with high-quality retrieved context.

The worked example also illustrates why the neighbor continuation matters. If the retrieval database contains a chunk beginning with "Arctic sea ice extent reached a record minimum in..." and the query from Chunk 2 matches this chunk, the model receives that chunk and the following text: "...2012, covering only 3.41 million square kilometers, roughly half the average from 1979-2000." This continuation, which extends beyond what was directly retrieved, gives the model access to specific quantitative facts that would be relevant for predicting tokens in Chunk 3 about Arctic conditions. Without the continuation, the model gets the beginning of a fact but not its resolution.