Late Chunking: Embed First, Chunk Later

[!NOTE] TL;DR

Naive chunking embeds each chunk alone, so a chunk that says "its population is 3.8 million" no longer knows that "it" is Berlin.

Late chunking flips the order. Run the whole document through a long-context embedding model once, keep the token embeddings, then mean-pool each chunk's token span. Every chunk vector now carries document context.

It costs about the same as naive chunking at index time and needs no extra storage or query latency. It does need a long-context model that exposes token embeddings and was trained with mean pooling, such as jina-embeddings-v2 or v3.

Anthropic's Contextual Retrieval buys more context with an LLM call per chunk. ColBERT keeps every token vector. Both can beat late chunking, but you pay for them in index cost or storage.

My default: try late chunking first when your documents are long and full of references back to earlier text. It's the cheapest fix for lost context.

Most of my work is turning messy documents into something an agent can reason over: contracts, specs, manuals, long email threads. The retrieval layer under those agents fails in one boring way more often than any other. The right passage exists, but its vector doesn't know what it's about.

This post covers a fix I like because it's almost free: late chunking, from Jina AI's 2024 paper Late Chunking: Contextual Chunk Embeddings Using Long-Context Embedding Models. I'll explain the problem, build it from scratch with transformers, and compare it with the alternatives so you know when to reach for something heavier.

The lost context problem

Here's the standard RAG ingestion pipeline. Split the document into chunks of a few hundred tokens, embed each chunk on its own, and store one vector per chunk.

Now take three sentences from an encyclopedia-style article:

Berlin is the capital of Germany. Its population is about 3.8 million. The city is also one of Germany's sixteen states.

Chunk that by sentence and embed each piece alone. Chunk 2 becomes a vector for "Its population is about 3.8 million." Nothing in that vector says Berlin. A user asks "how many people live in Berlin?" and the one chunk that answers it competes on the words "population" and "million" against every other city in your corpus.

Naive chunking cuts the document before the encoder sees it. Chunk 1 knows it is about Berlin. Chunks 2 and 3 only have a pronoun and "the city".

People usually reach for one of two fixes. Bigger chunks dilute the vector and hurt precision. Overlapping chunks only help if the antecedent happens to fall in the overlap. In real documents the thing "it" refers to is often a section heading three pages up.

The deeper problem is about order. We chunk before the model ever sees the document, so the model can't use context we've already thrown away.

How late chunking works

Late chunking keeps your chunk boundaries and changes when you apply them.

Tokenize the whole document, up to the model's context limit (8,192 tokens for the Jina v2 and v3 families).
Run one forward pass. Every token's output embedding has attended to every other token, so the embedding for "Its" already carries "Berlin".
Keep the token-level embeddings, the last_hidden_state, and skip the model's usual pooling.
Map each chunk's character span to its token span, mean-pool the token embeddings in that span, and normalise.

You still get one vector per chunk, same as before, and the same index, query path and storage. The only change is that each vector was computed with the whole document in view.

Late chunking: "Its" attends to "Berlin" inside the encoder, so the token embeddings h7 to h12 already carry the reference. Pooling just groups them into chunk vectors.

The method only works because of mean pooling. A mean-pooled embedding model is trained so that the average of its token embeddings is a good text embedding. Late chunking assumes the average over a sub-span is also a good embedding of that span, only now it's conditioned on everything around it. The paper tests this on BeIR retrieval tasks and reports consistent gains over naive chunking, with bigger gains on longer documents. That fits the intuition: short documents have little context to lose.

Why not just embed the whole document as one vector?

You could, and you'd get the opposite problem. A single vector for an 8k-token document averages over everything, so a query about one paragraph matches it only weakly. Late chunking gives you both. Each token embedding is contextual, and you still pool locally, so each vector stays specific to its span.

Compare the operations directly. With naive chunking, chunk i's vector is mean(encoder(chunk_i)). With late chunking it's mean(encoder(document)[span_i]). The pooling step is identical. Only the input to the encoder changes.

Building it with transformers

Below is a complete implementation. It uses jinaai/jina-embeddings-v2-small-en because it's small, handles 8k tokens and uses mean pooling. You'll need torch and transformers, plus trust_remote_code=True because the Jina v2 models ship their own model class.

Setup and chunk boundaries

Chunk boundaries live in character space. That way they're independent of the tokenizer, and you can reuse whatever splitter your ingestion pipeline already has (headings, sentences, layout blocks).

import re
 
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
 
MODEL = "jinaai/jina-embeddings-v2-small-en"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModel.from_pretrained(MODEL, trust_remote_code=True).eval()
 
SENTENCE = re.compile(r"\S.*?(?:[.!?](?=\s|$)|$)", flags=re.S)
 
 
def chunk_spans(text: str, max_chars: int = 400) -> list[tuple[int, int]]:
    """Greedily pack sentences into chunks; return (start, end) char spans."""
    spans, start, end = [], None, None
    for m in SENTENCE.finditer(text):
        s, e = m.span()
        if start is not None and e - start > max_chars:
            spans.append((start, end))
            start = None
        if start is None:
            start = s
        end = e
    if start is not None:
        spans.append((start, end))
    return spans

Token embeddings for the whole document

The important flag is return_offsets_mapping=True. It needs a "fast" (Rust) tokenizer and gives you the (char_start, char_end) of each token in the original string. That's how character chunks become token spans.

This function also handles documents longer than the context window, which we'll come back to. For a document that fits, the loop runs once.

@torch.no_grad()
def token_embeddings(text: str, window: int = 8192, overlap: int = 512):
    """Contextual embedding + char offsets for every token of `text`."""
    enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
    ids, offsets = enc["input_ids"], enc["offset_mapping"]
    body = window - 2  # leave room for [CLS] and [SEP]
    step = body - overlap
    pieces, start = [], 0
    while True:
        chunk_ids = tokenizer.build_inputs_with_special_tokens(ids[start:start + body])
        input_ids = torch.tensor([chunk_ids])
        hidden = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids))
        hidden = hidden.last_hidden_state[0, 1:-1]  # drop [CLS]/[SEP]
        # Overlap tokens were only there as left context for this window.
        pieces.append(hidden if start == 0 else hidden[overlap:])
        if start + body >= len(ids):
            break
        start += step
    return torch.tensor(offsets), torch.cat(pieces)

Pooling over spans

A token belongs to a chunk if its character range overlaps the chunk's range. Average those token embeddings, then L2-normalise so that cosine similarity is a plain dot product.

def pool_spans(offsets, embs, spans):
    starts, ends = offsets[:, 0], offsets[:, 1]
    vectors = []
    for s, e in spans:
        mask = (starts < e) & (ends > s)
        vectors.append(F.normalize(embs[mask].mean(dim=0), dim=0))
    return torch.stack(vectors)
 
 
def late_chunk(text: str):
    spans = chunk_spans(text)
    offsets, embs = token_embeddings(text)
    return spans, pool_spans(offsets, embs, spans)

The naive baseline

For a fair comparison, the baseline uses the same model and the same mean pooling. The only difference is that each chunk goes through the encoder alone. The same function embeds queries.

@torch.no_grad()
def embed(texts: list[str]) -> torch.Tensor:
    enc = tokenizer(texts, padding=True, truncation=True, max_length=8192,
                    return_tensors="pt")
    hidden = model(**enc).last_hidden_state
    mask = enc["attention_mask"].unsqueeze(-1).to(hidden.dtype)
    pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1)
    return F.normalize(pooled, dim=-1)
 
 
def naive_chunk(text: str):
    spans = chunk_spans(text)
    return spans, embed([text[s:e] for s, e in spans])

Quick sanity check on the Berlin example. Passing max_chars=1 makes every sentence its own chunk.

doc = ("Berlin is the capital of Germany. Its population is about 3.8 million. "
       "The city is also one of Germany's sixteen states.")
 
 
def compare(doc: str, question: str):
    spans = chunk_spans(doc, max_chars=1)  # one sentence per chunk
    q = embed([question])[0]
    naive = embed([doc[s:e] for s, e in spans])
    offsets, embs = token_embeddings(doc)
    late = pool_spans(offsets, embs, spans)
    for (s, e), n, l in zip(spans, naive, late):
        print(f"naive={float(n @ q):.3f}  late={float(l @ q):.3f}  {doc[s:e]}")
 
 
compare(doc, "How many people live in Berlin?")

I'm not going to print expected numbers. Run it yourself. The thing to look for is the "Its population..." row: its late score should rise relative to its naive score, because that vector now carries Berlin.

A tiny evaluation harness

Anecdotes don't settle retrieval questions. Build a small labelled set from your own corpus: real questions, each mapped to the chunks that answer it. Thirty to fifty questions is enough to see a clear difference. Then measure recall@k for both indexers.

def build_index(docs: dict[str, str], chunker):
    ids, vecs = [], []
    for doc_id, text in docs.items():
        spans, v = chunker(text)
        ids += [(doc_id, s, e) for s, e in spans]
        vecs.append(v)
    return ids, torch.cat(vecs)
 
 
def recall_at_k(index, labelled: list[tuple[str, set]], k: int = 5) -> float:
    """labelled: (question, set of relevant (doc_id, start, end) chunk ids)."""
    ids, vecs = index
    total = 0.0
    for question, relevant in labelled:
        scores = vecs @ embed([question])[0]
        top = scores.topk(min(k, len(ids))).indices.tolist()
        total += len({ids[i] for i in top} & relevant) / len(relevant)
    return total / len(labelled)
 
 
# for name, chunker in [("naive", naive_chunk), ("late", late_chunk)]:
#     print(name, recall_at_k(build_index(docs, chunker), labelled, k=5))

Chunk ids are (doc_id, start, end). Both indexers use the same chunk_spans, so the labels line up across them. If your splitter isn't deterministic, fix that first.

[!TIP] Label questions the way users actually ask them: short, and missing the subject the document assumes. "What's the notice period?" against a contract is a good test. "What is the notice period in the Acme master services agreement?" is not, because it smuggles the context into the query and hides exactly the failure you're trying to measure.

Requirements, and what to do with long documents

Late chunking has three real requirements.

A long-context embedding model. The context you can add is capped by the window. With a 512-token model, late chunking helps a little. With 8k it helps a lot.
Access to token embeddings. Hosted APIs that return only pooled vectors don't give you that. Some providers expose late chunking as an option (Jina's API does). Otherwise you run the model yourself.
Mean pooling in training. jina-embeddings-v2 and v3 are both mean-pooled. Models trained to put everything into the [CLS] token don't promise that averaged token embeddings mean anything. More on that below.

For documents longer than the window, split the token sequence into overlapping macro-windows. That's what the while loop in token_embeddings does. Each window after the first starts overlap tokens early. Those tokens act as left context, and their embeddings are discarded, so every kept token embedding has at least overlap tokens of preceding context.

Macro-windows for long documents. Each window re-reads the tail of the previous one as context, and only the new tokens' embeddings are kept. Stitched together, they form one token-embedding sequence for the whole document.

Context still doesn't cross more than one window. If "it" on page 40 refers to a title on page 1, macro-windows won't save you. Add a cheap header prefix to each window: document title and section path. That's a crude form of what Contextual Retrieval does properly.

Alternatives, and when each one wins

Late chunking is one of several ways to fix lost context. Here are the three I'd compare it with.

Naive chunking. Embed each chunk alone. It's cheap, simple and works with any embedding API. It's fine when chunks are self-contained, such as FAQs, product cards or short tickets.

Contextual Retrieval. Anthropic's Contextual Retrieval (September 2024) asks an LLM, for each chunk, to write a short blurb that places the chunk within the full document. The blurb is prepended to the chunk before it is embedded and before it goes into a BM25 index ("contextual BM25"). Anthropic reports substantial drops in retrieval failures, and the best results came from combining it with BM25 and a reranker. The context is explicit and readable, and an LLM can resolve references across a long document. The cost is one LLM call per chunk at index time. Prompt caching helps, since the document is shared across those calls, but it still dominates your indexing bill.

Late interaction (ColBERT). ColBERT (Khattab and Zaharia, 2020) never pools at all. It stores one vector per token and scores a query with MaxSim: for each query token, find the most similar document token, then sum those scores. It's precise and good at matching specific terms. The catch is that storage scales with tokens rather than chunks, and you need an index that supports multi-vector search. ColBERTv2 added compression to make storage tolerable, but it remains a different operating point.

Where chunking sits in each pipeline. Late chunking is the only one that applies chunk boundaries after the encoder has seen the whole document.

	Naive	Late chunking	Contextual Retrieval	ColBERT
Index-time cost	1 encoder pass per chunk	1 encoder pass per doc window	1 LLM call + 1 encoder pass per chunk	1 encoder pass per chunk
Storage	1 vector per chunk	1 vector per chunk	1 vector per chunk + BM25	1 vector per token
Query latency	ANN lookup	ANN lookup	ANN + BM25 (+ rerank)	multi-vector search + MaxSim
Context reach	none	the model window	whole document (what fits in the LLM prompt)	none across chunks, but fine-grained matching
Model constraints	any	long-context, mean-pooled, token outputs	any embedder + an LLM	ColBERT-style model
Debuggable context	n/a	implicit in the vector	readable text you can inspect	implicit

Here's how I'd choose between them.

Naive wins when chunks already stand on their own, or when you're stuck with a hosted embedding API that returns only pooled vectors.
Late chunking wins when documents are long and full of internal references, and you want better retrieval without changing your index, query path or budget. It's the best quality you get per dollar of change.
Contextual Retrieval wins when references cross the whole document, when you want lexical matching on terms that only appear in context (contextual BM25 is underrated), or when you need to see the context for debugging. Pay for it if quality matters more than indexing cost.
ColBERT wins when you need precise term-level matching, such as part numbers, clause names or rare entities, and can afford the storage and a multi-vector index. It solves a different problem from lost context, and you can stack it with the others.

Nothing stops you combining them. Late-chunked dense vectors plus contextual BM25 plus a reranker is a perfectly sensible stack.

Pitfalls

[!WARNING] Late chunking changes every vector in your index. Old naive vectors and new late-chunked vectors aren't comparable with each other, so plan a full re-index. Don't do a rolling migration that mixes the two.

Boundary chunks. The first chunk of a document has only right context. Chunks at a macro-window seam have less left context than their neighbours. Keep the overlap generous, at least a few hundred tokens, and if you can, align window starts with section boundaries.

Very long documents. Context doesn't flow across macro-windows beyond the overlap. For a 200-page manual, late chunking fixes local references ("this valve", "the above table") but not global ones. Prefix each window with the title and section path, or add Contextual Retrieval for the global part.

Models trained with CLS pooling. If a model was trained to pack meaning into [CLS], averaging its token embeddings gives you something it was never optimised to produce. It may still work. Measure it, don't assume it. Check the model card or its sentence-transformers pooling config before you start.

Tokenizer offsets. return_offsets_mapping needs a fast tokenizer and raises an error with a slow one. Offsets are relative to the exact string you passed in, so if your pipeline normalises whitespace or Unicode after computing chunk spans, the spans and offsets drift apart. Compute both on the same string.

Instruction and task prefixes. Some models expect a passage prefix or a task-specific adapter (jina-embeddings-v3 uses task LoRA adapters). Apply them the same way in both the naive and late paths, and check the model card for how to select the retrieval adapter with the raw transformers API.

Memory. An 8k-token forward pass costs much more activation memory than many 300-token passes, because attention cost grows with sequence length. Batch by total tokens, not by number of documents, and expect to tune window size on smaller GPUs.

Re-indexing cost. Index time per document is roughly the same as naive, but you need to re-run every document, and you need the full document text at index time. If your pipeline only kept chunks, you have to rebuild from source.

Decision checklist

Before adopting late chunking, I go through this list.

Do my chunks contain pronouns, "the above", abbreviations or section-relative references? If not, stop. Naive is fine.
Are my documents longer than a single chunk but mostly under the model's context window? That's the sweet spot.
Can I run the embedding model myself, or does my provider expose late chunking? If neither, look at Contextual Retrieval instead.
Is the model mean-pooled and long-context? If not, switch models first and measure that change on its own.
Do I have a labelled query set and a recall@k baseline? If not, build that before changing anything.
Can I afford a full re-index and keep source text around? Late chunking needs whole documents.
Do references span the entire document? If so, plan to add Contextual Retrieval or a title and section prefix on top.
Do queries hinge on exact rare terms? Add BM25 or consider ColBERT. Late chunking won't help much there.

Closing thoughts

What I like about late chunking is the reframing. Chunking is a storage decision, not a modelling decision. You can let the model read everything first and only then decide how to slice what it understood. It's a few dozen lines of code, the index stays the same shape, and the queries that used to miss because "it" had lost its referent start landing.

Start with the harness, measure naive against late on your own documents, and only pay for an LLM per chunk if the numbers say you need it.