Late Chunking

Late Chunking

The Problem with Naive Chunking

Traditional pipeline: split document into chunks first, embed each chunk independently. Each chunk embedding sees only its own 200-400 tokens — pronouns, anaphora, and shared terminology with earlier sections become invisible.

Document: "Acme Corp reported Q3 revenue of $42M...
          The company attributed growth to...
          Its subscription tier..."

Chunk 2: "The company attributed growth to..."
         ^^^ "the company" refers to Acme Corp — lost in isolation

A query like "Acme Corp revenue growth drivers" may match Chunk 1 but not Chunk 2, even though Chunk 2 is the one answering the question.

The Late Chunking Fix

1. Run the full document through a long-context encoder (8k, 32k, or 128k tokens) in a single forward pass so every token attends to every other token.
2. Record token-level embeddings and chunk token offsets.
3. Mean-pool the token embeddings inside each chunk span to produce chunk vectors.

Every chunk vector now encodes the entire document's context, not just its own span.

Minimal Implementation with jina-embeddings-v3

jina-embeddings-v3 supports 8192 tokens and exposes encode_chunks for this workflow.

from transformers import AutoModel, AutoTokenizer
import torch
import numpy as np

model_id = "jinaai/jina-embeddings-v3"
tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModel.from_pretrained(model_id, trust_remote_code=True).eval()

def chunk_by_sentence(text: str, max_chars: int = 400):
    """Return (chunk_text, (start_char, end_char)) tuples."""
    import re
    sentences = re.split(r"(?<=[.!?])\s+", text)
    chunks, buf, start = [], "", 0
    pos = 0
    for s in sentences:
        if len(buf) + len(s) > max_chars and buf:
            chunks.append((buf.strip(), (start, pos)))
            buf, start = s + " ", pos
        else:
            buf += s + " "
        pos += len(s) + 1
    if buf.strip():
        chunks.append((buf.strip(), (start, pos)))
    return chunks

def char_to_token_spans(text: str, char_spans):
    enc = tok(text, return_offsets_mapping=True, add_special_tokens=True,
              truncation=True, max_length=8192)
    offsets = enc["offset_mapping"]
    token_spans = []
    for c_start, c_end in char_spans:
        t_start = next((i for i, (a, b) in enumerate(offsets) if a >= c_start and b > 0), None)
        t_end   = next((i for i, (a, b) in enumerate(offsets) if b >= c_end and b > 0), None)
        if t_start is not None and t_end is not None:
            token_spans.append((t_start, t_end + 1))
    return enc, token_spans

@torch.inference_mode()
def late_chunk_embed(text: str):
    chunks = chunk_by_sentence(text)
    chunk_texts = [c[0] for c in chunks]
    char_spans  = [c[1] for c in chunks]

    enc, token_spans = char_to_token_spans(text, char_spans)
    input_ids = torch.tensor([enc["input_ids"]])
    attn      = torch.tensor([enc["attention_mask"]])

    out = model(input_ids=input_ids, attention_mask=attn, task_id="retrieval.passage")
    token_embs = out.last_hidden_state[0]  # (T, D)

    vectors = []
    for (ts, te) in token_spans:
        pooled = token_embs[ts:te].mean(dim=0)
        pooled = torch.nn.functional.normalize(pooled, dim=-1)
        vectors.append(pooled.cpu().numpy())
    return chunk_texts, np.vstack(vectors)

Comparison: Naive vs Late Chunking

from sentence_transformers import SentenceTransformer

naive_model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True)

def naive_chunk_embed(text: str):
    chunks = [c[0] for c in chunk_by_sentence(text)]
    vecs = naive_model.encode(chunks, normalize_embeddings=True,
                              prompt_name="retrieval.passage")
    return chunks, vecs

def cos(a, b):
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

query_model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True)
q = query_model.encode("Acme revenue growth drivers", prompt_name="retrieval.query",
                       normalize_embeddings=True)

doc = open("acme_10q.txt").read()
_, naive_vecs = naive_chunk_embed(doc)
_, late_vecs  = late_chunk_embed(doc)

print("naive top-1:", max(cos(q, v) for v in naive_vecs))
print("late  top-1:", max(cos(q, v) for v in late_vecs))

Reported Benchmarks (Jina, 2024)

On BeIR subsets with ~8k-token documents, late chunking improves nDCG@10 by 3-8% over naive chunking using the same encoder. Biggest lift on:

• Scientific papers (SciFact, NFCorpus) — heavy cross-reference.
• Legal documents — anaphoric references to parties defined once.
• Customer support transcripts — pronoun-heavy dialogue.

No lift (or slight regression) on FAQ-like corpora where chunks are already self-contained.

Chunk Size Considerations

When Late Chunking Won't Help

• Your encoder's max context is the same as your chunk size (nothing to gain).
• Documents are already semantically independent (FAQ, product catalog entries).
• You use a cross-encoder reranker on raw text (reranker sees full context anyway).

Storage and Latency

Storage unchanged — you still store one vector per chunk. Latency shifts from many small forwards to one large forward; on GPU this is usually faster due to better hardware utilization. Downside: you cannot parallelize across a cluster by sharding chunks.

Combining with Other Techniques

• Works with hybrid search (BM25 + late-chunked dense).
• Compatible with reranking — use late-chunked vectors for recall stage.
• Can be combined with Matryoshka truncation for storage savings.
• Does NOT combine cleanly with ColBERT (ColBERT already late-interacts at token level).

Anti-Patterns

Production Checklist

• Encoder max context >= longest expected document
• Document truncation policy defined for docs beyond context
• Chunk token offsets validated against tokenizer output (off-by-one-prone)
• nDCG@10 measured on domain eval set vs naive baseline
• Per-chunk vectors stored with document id + token span metadata
• GPU memory profile checked for longest documents
• Re-index path documented (model or chunker changes require full rebuild)