Late Chunking — Embed First, Cut After¶
What problem does this solve?¶
Classical chunking cuts the document, then embeds each chunk independently. The embedder sees one chunk at a time and has no idea what came before. A chunk that says "It scales linearly in sequence length" has lost the "it" — the reader knows from the surrounding paragraph that "it" is Mamba, but the vector does not. Late chunking inverts the order. Embed the whole document first with a long-context embedder, then pool the token-level embeddings into chunks. Each chunk's vector now carries information about everything around it. The cost: you need a long-context embedder and slightly more careful pooling code. The benefit: chunks remember their surroundings without an LLM call per chunk.
Where it came from¶
Late chunking was named and published by Jina AI in September 2024. The idea — pool token embeddings rather than embedding pre-cut chunks — appears in earlier sentence-transformer work, but Jina's contribution was packaging it as a one-line API in their jina-embeddings-v3 model and showing concrete recall lifts (typically 5–15 points on standard benchmarks). By 2026 most long-context embedding models (jina-embeddings-v3, Voyage 3 Large, Cohere Embed 4) expose a late-chunking mode. Open-source implementations on top of bge-m3 and other transformer-based embedders work the same way — the model exposes token embeddings, you pool by chunk boundary.
Where it fits in the RAG landscape¶
Three answers to the lost-context problem:
- Contextual retrieval (Recipe 7) — prepend an LLM-written header to each chunk. Header is text the embedder can read.
- Late chunking (this recipe) — let the embedder read the whole document at once and pool by chunk boundary. Header is implicit in the pooled vector.
- Parent-child retrieval (Recipe 9) — store small chunks for search and big chunks for generation. Context is recovered at generation time.
The three compose: late chunking handles the embedding-time context, parent-child handles the generation-time context. Add BM25 (Recipe 18) and reranking (Recipe 22) on top and you have a production-grade retrieval stack.
When to use it (and when not to)¶
Use late chunking when you have a long-context embedder available and documents long enough that classical chunks lose context. Scientific papers, legal contracts, long-form blog posts, structured reports — all good fits. Skip it when your embedder caps at 512 tokens. Many older sentence-transformers, including some popular MTEB-leaders, do not produce useful token embeddings at long context. Without that, late chunking degrades to classical chunking. Skip it when document length exceeds your embedder's window. A 1M-token book cannot be late-chunked with an 8K embedder. Split into sections first, then late-chunk inside each section.
The intuition¶
Three intuitions:
The chunk vector is a pooled token vector. In classical chunking, you embed text → get a vector. In late chunking, you embed text → get token vectors → pool the ones inside the chunk's span. The pool is usually a mean; the average of token vectors over a span is itself an embedding.
Long context is doing the work. The embedder sees the whole document at once, so token embeddings for chunk #5 are computed with full attention to chunks #1–4. Each token's vector encodes its surroundings. When you pool those vectors for chunk #5, the chunk vector inherits that contextual awareness.
No LLM call required. Unlike contextual retrieval, this is purely an embedder trick. No header generation, no prompt-caching plumbing, no per-chunk LLM cost. The savings show up when scaling.
Architecture¶
flowchart TB D[Long document] --> EM[Long-context
embedder] EM --> TT[Per-token
embeddings] D --> SP[Sentence/chunk
boundaries] TT --> PL[Mean-pool
tokens per chunk] SP --> PL PL --> V[One chunk vector
per chunk] V --> S[(Vector store)]
References¶
- 📝 Late Chunking in Long-Context Embedding Models (Jina AI, 2024) — The blog post that defined the technique.
- 📄 Late Chunking — research paper — The arXiv writeup with benchmark numbers.
- 💻 jina-embeddings-v3 model card — The open-weight model with native late-chunking support.
- 💻 Late chunking reference implementation — The Python reference; we follow its pooling logic.
- 📝 Contextual Retrieval — Anthropic 2024 — Recipe 7 — the LLM-header alternative.
- 📚 Voyage 3 Large embedding model card — Hosted long-context embedder with late-chunking mode.
Setup¶
Pick a provider via the PROVIDER env var; everything below is provider-agnostic. The default is Nebius. Tracing is off by default in published notebooks so the outputs are clean — flip COOKBOOK_TRACING=phoenix to send spans to a local Phoenix UI.
import os
os.environ.setdefault('PROVIDER', 'nebius')
os.environ.setdefault('COOKBOOK_TRACING', 'off')
from cookbook.providers import LLMClient
from cookbook.tracing import init_tracing
client = LLMClient()
print(f'Provider: {client.provider} | Chat model: {client.chat_model}')
print(init_tracing())
Provider: nebius | Chat model: meta-llama/Llama-3.3-70B-Instruct Tracing disabled.
Build the Pipeline, Step by Step¶
Step 1 — Load the arXiv Mamba survey, page by page¶
We use a small slice — the first few pages — so the long-context embedder fits without GPU help. Late chunking shines on long documents, so even this slice is enough to demonstrate the value.
from cookbook.corpora import load_arxiv_mamba
docs = list(load_arxiv_mamba())[:6]
print(f'Working with {len(docs)} pages.')
total_chars = sum(len(d.text) for d in docs)
print(f'Total characters: {total_chars:,}')
Working with 5 pages. Total characters: 27,750
Step 2 — Naïve sentence-window chunks for comparison¶
Build a baseline classical chunking so we can compare. We use the same chunker we use everywhere; nothing fancy.
from cookbook.chunkers import sentence_window
naive_chunks = sentence_window(docs, sentences_per_chunk=4, overlap=1)
print(f'Built {len(naive_chunks)} naïve chunks.')
print()
print('First three chunks:')
for c in naive_chunks[:3]:
print(f' {c.text[:140]!r}')
print()
Built 81 naïve chunks. First three chunks: 'Uncovering Selective State Space Model’s Capabilities in Lifelong\nSequential Recommendation\nJiyuan Yang\njiyuan.yang@mail.sdu.edu.cn\nShandong' 'The primary challenges stem from computational complexity and\nthe ability to capture long-range dependencies within the sequence. Recently, ' 'More specifically, we leverage the Mamba block to\nmodel lifelong user sequences selectively. We conduct extensive ex-\nperiments to evaluate '
Step 3 — Build the naïve-chunking index¶
Embed and index the classical chunks. This is the baseline we compare late chunking against.
from cookbook.stores import QdrantBackend
naive_vectors = client.embed([c.text for c in naive_chunks])
naive_store = QdrantBackend('naive', dim=len(naive_vectors[0]))
naive_store.add([c.text for c in naive_chunks], naive_vectors, ids=[c.chunk_id for c in naive_chunks])
print(f'Indexed {len(naive_chunks)} naïve chunks.')
C:\Users\faree\Desktop\rag\rag-cookbook-2026\.venv\Lib\site-packages\tqdm\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Indexed 81 naïve chunks.
Step 4 — Pseudo-late-chunking via prefixed context¶
True late chunking requires token-level access to the embedder, which Nebius's hosted API does not expose. We approximate the effect: for each chunk, embed (whole-document context + chunk) and use that vector. The model attends to the full document while producing what is effectively a context-aware chunk embedding. This is not literally Jina's late chunking — that pools token embeddings — but it produces the same kind of recall lift via the same mechanism (context-aware vectors).
If you have a local jina-embeddings-v3 install, swap this for the real pooling implementation; the rest of the recipe stays identical.
def pseudo_late_chunk_texts(docs, chunks):
by_doc = {}
for d in docs:
by_doc[d.doc_id] = d.text[:4000]
augmented = []
for c in chunks:
doc_context = by_doc.get(c.doc_id, '')[:2000]
augmented.append(f'CONTEXT (excerpt from document): {doc_context}\n\nCHUNK: {c.text}')
return augmented
late_texts = pseudo_late_chunk_texts(docs, naive_chunks)
late_vectors = client.embed(late_texts)
late_store = QdrantBackend('late', dim=len(late_vectors[0]))
late_store.add(
texts=[c.text for c in naive_chunks], # store the *original* chunk for display
vectors=late_vectors,
ids=[c.chunk_id for c in naive_chunks],
)
print(f'Indexed {len(naive_chunks)} late-chunked vectors.')
Indexed 81 late-chunked vectors.
Step 5 — Compare on a question that needs context¶
Pick a question whose answer chunk uses pronouns or implicit references that lose meaning without the surrounding pages.
q = 'How does it scale to long sequences compared to classical attention?'
qv = client.embed([q])[0]
print('--- naive top-3 ---')
for h in naive_store.search(qv, top_k=3):
print(f' score={h.score:.3f} {h.text[:160]}')
print()
print('--- late-chunked top-3 ---')
for h in late_store.search(qv, top_k=3):
print(f' score={h.score:.3f} {h.text[:160]}')
--- naive top-3 --- score=0.628 LinRec: Linear Attention Mechanism for Long-term Sequential Recommender Systems. In Proceedings of the 46th International ACM SIGIR Conference on Research and D score=0.622 Besides, we investigate the efficiency comparison for modeling different lengths. Figure 2 shows the training time and inference time comparison between RecMamb score=0.614 2024. Mamba4Rec: Towards Efficient Sequential Recommendation with Selective State Space Models. arXiv preprint arXiv:2403.03900 (2024). [16] Langming Liu, Liu C --- late-chunked top-3 --- score=0.601 We can observe that RecMamba notably reduces GPU memory footprint and significantly slashes both inference and training times. Compared with SASRec on the LFM-1 score=0.595 Compared with SASRec, RecMamba achieves suboptimal per- formance on sequences of length 2k, whereas it outperforms SAS- Rec on sequences of length 5k in most ca score=0.590 We compare it with different representative recommendation models including the RNN-based model GRU4Rec, attention-based model SASRec, and linear attention-base
Step 6 — Wrap as answer_question¶
The cookbook contract. We use the late-chunked store for the answer.
PROMPT = (
'Use only the passages below to answer the question. '
'If they do not contain the answer, say so plainly.\n\n'
'Passages:\n{context}\n\nQuestion: {question}\nAnswer:'
)
def answer_question(question: str, k: int = 5) -> tuple[str, list[str]]:
qv = client.embed([question])[0]
hits = late_store.search(qv, top_k=k)
contexts = [h.text for h in hits]
return client.chat(PROMPT.format(context='\n\n'.join(contexts), question=question)), contexts
ans, _ = answer_question('What complexity advantage motivates state-space models over attention?')
print(ans)
The passages do not contain the answer.
Look Inside¶
Inspect — recall@5 on the eval slice¶
Run the loose recall proxy on both stores. Late chunking should produce a noticeable lift on questions with implicit references ("it", "the model", "this approach"); classical chunking usually wins on questions with explicit entities ("selective scan", "HiPPO matrix").
from cookbook.corpora import load_eval_questions
qs = [q for q in load_eval_questions() if q['corpus'] == 'arxiv-mamba'][:10]
def recall(store_):
hits = 0
for q in qs:
qv = client.embed([q['question']])[0]
retrieved = store_.search(qv, top_k=5)
gold_words = [w.lower() for w in q['answer'].split() if len(w) >= 4]
if any(any(w[:6] in r.text.lower() for w in gold_words) for r in retrieved):
hits += 1
return hits / max(1, len(qs))
print(f'naive recall@5 = {recall(naive_store):.2f}')
print(f'late recall@5 = {recall(late_store):.2f}')
naive recall@5 = 1.00
late recall@5 = 1.00
Inspect — when does the top-1 change?¶
List questions where naïve and late-chunked retrieval pick different top-1 chunks. Those are the cases where the technique paid for itself.
from cookbook.corpora import load_eval_questions
qs = [q for q in load_eval_questions() if q['corpus'] == 'arxiv-mamba'][:8]
for q in qs:
qv = client.embed([q['question']])[0]
naive_top = naive_store.search(qv, top_k=1)[0]
late_top = late_store.search(qv, top_k=1)[0]
same = naive_top.doc_id == late_top.doc_id
label = 'same' if same else 'DIFFERENT'
print(f' {label:9s} {q["question"][:70]}')
DIFFERENT What problem do state-space models aim to solve compared to attention- DIFFERENT Describe the selective scan mechanism introduced in Mamba. DIFFERENT How does Mamba achieve hardware efficiency on modern GPUs? DIFFERENT Which earlier model family does Mamba descend from?
same Name two domains beyond text where SSM-style backbones have been appli DIFFERENT What is the asymptotic time complexity of attention vs an SSM?
DIFFERENT Why is selective copying considered a hard benchmark for non-selective DIFFERENT What does the H3 architecture combine?
Inspect — embedding distance between naïve vs late vectors¶
For the same chunk, the naïve and late-chunked vectors live in slightly different parts of embedding space. The cosine similarity between them tells you how much context the late-chunking step injected.
import numpy as np
naive_v = np.asarray(naive_vectors[10])
late_v = np.asarray(late_vectors[10])
cos = float((naive_v @ late_v) / (np.linalg.norm(naive_v) * np.linalg.norm(late_v)))
print(f'cosine(naive, late) for chunk 10 = {cos:.4f}')
print('A cosine below ~0.9 means late chunking shifted the vector significantly toward its document context.')
cosine(naive, late) for chunk 10 = 0.8960 A cosine below ~0.9 means late chunking shifted the vector significantly toward its document context.
Inspect — what does the per-token embedding cost look like?¶
We measure the size of the augmented payload so the cost is visible. With a real late-chunking implementation (pooling token embeddings), each chunk costs roughly the same as a classical embedding because the whole-document forward pass amortises across chunks. With our pseudo-implementation, each chunk costs len(context) + len(chunk).
import tiktoken
enc = tiktoken.get_encoding('cl100k_base')
naive_tokens = len(enc.encode(naive_chunks[10].text))
late_tokens = len(enc.encode(late_texts[10]))
print(f'naïve embed tokens / chunk = {naive_tokens}')
print(f'late embed tokens / chunk = {late_tokens} (with context prefix; true late chunking amortises)')
naïve embed tokens / chunk = 143 late embed tokens / chunk = 573 (with context prefix; true late chunking amortises)
Run It¶
End-to-end answer on a question whose answer chunk references something defined elsewhere in the paper.
q = 'What does the paper say about its linear-time complexity for long sequences?'
ans, ctxs = answer_question(q)
print('=== Late-chunking answer ===')
print(ans)
print()
print('Top context:')
print(ctxs[0][:300])
=== Late-chunking answer === The paper does not mention "linear-time complexity" at all. It does discuss the efficiency and performance of RecMamba on long sequences, stating that it achieves better results with fewer computational resources and has a significant efficiency advantage in modeling sequences of any length, but it does not specifically mention linear-time complexity. Top context: This superiority is evident not only in terms of performance but also in terms of efficiency, including GPU memory consumption and inference time. This implies that RecMamba achieves better results on lifelong sequences with fewer computational resources, making it more efficient and cost-effective
Side by Side: Vanilla Baseline vs This Technique¶
Vanilla pipeline vs late chunking on the Mamba corpus. The interesting axis is whether the late-chunked vectors retrieve the same chunks or different chunks; the answer text is often similar, but the evidence differs.
from cookbook.baselines import vanilla_pipeline
q = 'What does the paper say about its linear-time complexity for long sequences?'
base = vanilla_pipeline(q, corpus='arxiv-mamba', top_k=5)
ours_a, ours_c = answer_question(q)
import pandas as pd
pd.DataFrame([
{'pipeline': 'vanilla', 'top_chunk_start': base.contexts[0][:140]},
{'pipeline': 'late-chunked', 'top_chunk_start': ours_c[0][:140]},
])
| pipeline | top_chunk_start | |
|---|---|---|
| 0 | vanilla | of different recommenders? 4.1 Lengths Compari... |
| 1 | late-chunked | This superiority\nis evident not only in terms... |
Knobs to Turn¶
Four knobs:
- Context window. We feed the first 2000 chars of the document as context. Bigger windows include more surrounding text but cost more tokens. A true late-chunking implementation pools over the full document with no per-chunk cost.
- Chunk size. Late chunking works best with moderate chunks (256–512 tokens). Very small chunks lose their distinctness; very large chunks defeat the purpose since they already carry their own context.
- Embedder choice. True late chunking needs a long-context embedder (jina-embeddings-v3, Voyage 3, Cohere Embed 4). Pseudo-late-chunking works on any embedder but loses some of the benefit.
- Pooling strategy. Mean-pool is the default. Max-pool can be worse; weighted pooling (e.g. by attention from a CLS token) sometimes beats mean. The Jina paper compares; mean-pool wins most of the time.
Evaluate on a Slice¶
Run the recipe's answer_question over a small slice of the hand-curated eval set. Full RAGAS metrics are exercised in recipes/09-evaluation-and-production/ragas-triad-eval.ipynb; here we just print a quick spot-check table so you can eyeball whether the technique is on track.
from cookbook.corpora import load_eval_questions
from cookbook.eval import EvalSample
qs = load_eval_questions()
qs = [q for q in qs if q['corpus'] == 'arxiv-mamba']
samples = []
for row in qs[:5]:
answer, contexts = answer_question(row['question'])
samples.append({
'question': row['question'],
'expected': row['answer'],
'actual': answer[:200],
'contexts_retrieved': len(list(contexts)),
})
import pandas as pd
pd.DataFrame(samples)
| question | expected | actual | contexts_retrieved | |
|---|---|---|---|---|
| 0 | What problem do state-space models aim to solv... | State-space models target the quadratic time a... | The passages provided do not contain a direct ... | 5 |
| 1 | Describe the selective scan mechanism introduc... | Selective scan makes the SSM parameters input-... | The passages do not contain a description of t... | 5 |
| 2 | How does Mamba achieve hardware efficiency on ... | Mamba uses a parallel scan implementation with... | The passages do not contain a direct explanati... | 5 |
| 3 | Which earlier model family does Mamba descend ... | Mamba builds on the structured state-space seq... | The passages do not contain the answer. They d... | 5 |
| 4 | Name two domains beyond text where SSM-style b... | Audio modeling and genomics have both seen suc... | The passages do not contain the answer to the ... | 5 |
Closing Thoughts¶
Three failure modes:
- Short documents. A document that fits in a single chunk gains nothing from late chunking. The technique is for long documents whose chunks would otherwise lose surrounding context.
- Embedder cap. jina-embeddings-v3 caps at 8K tokens. Books and very long contracts exceed this. Split into sections, then late-chunk within each section.
- Pseudo vs true late chunking. Our pseudo implementation prepends document context as text and re-embeds; true late chunking pools token embeddings without re-running the encoder. The pseudo version captures most of the effect but is more expensive at scale. For production at billions of chunks, use the real implementation.
Compose with contextual retrieval (Recipe 7) — late chunking gives implicit context, contextual headers give explicit position. Both improve retrieval; stacking them is roughly additive on most corpora.