Joe Abou Nakoul Docs

Late chunking

Properties1
tagsrag, retrieval, embeddings, chunking, long-context

6 min read

Late chunking: pool after attention, not before

Embed the whole document first and slice the token embeddings afterwards, so each chunk’s vector still carries the context around it.

The problem

Standard RAG chunks a document, then embeds each chunk independently. That order throws away everything outside the chunk boundary before the model ever sees it. A chunk reading “The airline attributed the growth to new routes in Asia” contains no vector evidence of which airline — the name was two sentences back, in a different chunk. So a query for “Qatar Airways expansion” scores that chunk low, and the fact you needed sits at rank 40 while a chunk that merely repeats the company name scores high. You see this as a retrieval system that works on documents full of proper nouns and quietly fails on documents full of pronouns, definite articles, and anaphora — which is most real documents.

How it works

Late chunking inverts the two steps. Chunking still happens, but pooling moves to the end.

  1. Tokenize the entire document and run it through a long-context embedding model in one pass. Every token attends to every other token, so the representation of “the airline” now contains information from “Qatar Airways” upstream.
  2. Take the token-level embeddings from the final hidden layer — before the model’s usual mean-pooling step.
  3. Map each chunk back to its token span using the tokenizer’s offset mapping.
  4. Mean-pool over each span separately. One vector per chunk, as before.

The output is shaped identically to naive chunking: same number of chunks, same dimensionality, same index. Nothing downstream changes. What changed is that pooling happened after attention rather than before it, so contextual information crossed the boundaries that chunking would otherwise have sealed.

The constraint is context length. The whole document must fit in the embedding model’s window, which is why this became practical only once 8k-token embedding models existed. Beyond that window you fall back to “long late chunking” — overlapping macro-windows of a few thousand tokens each, pooling within a window, accepting that context does not cross window edges.

Cost is roughly unchanged. You run the same total tokens through the same encoder; you batch differently. Attention over the full document is quadratic where independent chunks were not, so wall-clock indexing time rises on long documents while the vector store stays the same size.

Code

Needs transformers and torch; the model is ~130M params and runs on CPU.

# pip install transformers torch
import torch
from transformers import AutoModel, AutoTokenizer
 
MODEL = "jinaai/jina-embeddings-v2-small-en"  # 8k context window
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL, trust_remote_code=True)
 
doc = ("Qatar Airways reported record traffic in 2025. "
       "The airline attributed the growth to new routes in Asia. "
       "Its cargo division, however, contracted by four percent.")
chunks = [s + "." for s in doc.rstrip(". ").split(". ")]
 
enc = tok(doc, return_tensors="pt", return_offsets_mapping=True)
offsets = enc.pop("offset_mapping")[0]
 
with torch.no_grad():
    # one pass over the WHOLE document: every token attends to every other token
    tokens = model(**enc).last_hidden_state[0]
 
vectors, cursor = [], 0
for c in chunks:
    start = doc.index(c.rstrip("."), cursor)
    end = start + len(c.rstrip("."))
    cursor = end
    span = [i for i, (s, e) in enumerate(offsets) if s >= start and e <= end and e > s]
    # pooling AFTER attention: "the airline" still carries "Qatar Airways"
    vectors.append(tokens[span].mean(0))
 
q = model(**tok("Qatar Airways route expansion", return_tensors="pt")).last_hidden_state[0].mean(0)
for c, v in zip(chunks, vectors):
    print(round(float(torch.cosine_similarity(q, v, dim=0)), 3), c)

Run it, then run it again with the chunks embedded independently. The second sentence is where the two disagree, and that gap is the whole technique.

Where it fits

This is an indexing-time change in a retrieval service, and nothing else in the stack needs to know. If you run a Python ingestion job that chunks documents and writes vectors to a store, late chunking replaces the embed step inside that job — the store schema, the query path, and any Angular front end sitting on top are untouched. That makes it unusually cheap to A/B: build a second index with the same chunk boundaries, run your existing eval set against both, compare recall@k.

The signal that it is time to reach for this: your chunks are prose rather than independent records, your failure cases are queries where the answer chunk never surfaces, and inspecting those chunks shows they refer to an entity named somewhere else in the document. Contracts, policy documents, technical manuals, and transcripts all have this shape. Product catalogues and log records do not.

When not to

If chunks are already self-contained — one row, one ticket, one product — there is no cross-chunk context to preserve and you will measure nothing. If documents exceed the embedding model’s context window, most of the benefit disappears at the window seams. And if your retrieval failures are actually vocabulary mismatch rather than lost referents, a hybrid BM25 + dense setup or a cross-encoder reranker will buy far more recall for far less change; try the reranker first, since it is a query-time knob you can remove.

Check yourself

1. Why does this leave the vector store unchanged in size, when naive chunking and late chunking clearly encode different information?

Answer The number of chunks and the embedding dimensionality are identical — only the values inside each vector differ. Late chunking changes what the pooling operation averages over, not how many vectors come out. This is why it is a drop-in swap rather than a migration.

2. Your documents average 40k tokens against an 8k-context encoder. What breaks, and what do you do?

Answer You cannot run one pass over the document, so context stops crossing window boundaries. Fall back to overlapping macro-windows of a few thousand tokens, pooling chunks within each. Context is preserved inside a window and lost across seams, so place seams at section boundaries where referents are least likely to cross.

3. You ship it and recall@5 moves by half a point. Keep it or revert?

Answer Half a point on an aggregate is usually noise, but the aggregate is the wrong metric here — the technique only acts on queries whose answer chunk contains a cross-boundary referent. Segment the eval set by that property and measure the subset. If the subset does not move either, your corpus is more self-contained than you assumed, and reverting is correct.

Go deeper

  • Late Chunking: Contextual Chunk Embeddings — the original paper (~20 min). Section 3 has the method; the ablation on document length is the part worth your time.
  • Reference implementation — notebooks including the long-document macro-window variant. Read chunked_pooling/ before the README.
  • Contextual Retrieval — a different solution to the same problem: prepend an LLM-generated context blurb to each chunk before embedding. More expensive, no context-window ceiling. Useful to know as the alternative you would benchmark against.

Graph View

Backlinks