Late Chunking in Long Context Embedding Models

Introduction to Late Chunking in Long Context Embedding Models

Retrieval-augmented generation (RAG) pipelines have become a cornerstone of modern LLM applications, and at the heart of every RAG system lies an embedding model. Traditionally, embedding workflows chunk long documents into smaller pieces, embed each chunk independently, and store the resulting vectors in a vector database. While this approach works, it suffers from a subtle but important flaw: each chunk loses the broader context of the document it came from. Late Chunking is a technique that addresses this problem by reordering the operations — encoding first, chunking second.

What Is Late Chunking?

Late Chunking is a method introduced by Jina AI that leverages long-context embedding models to preserve contextual information across chunks. In a traditional pipeline, the sequence is:

Split the document into chunks
Encode each chunk independently through the embedding model
Pool token embeddings into a single vector per chunk

In Late Chunking, the sequence becomes:

Encode the entire document (or a large window) through the embedding model's transformer encoder in one pass
Split the resulting token-level embeddings into chunks based on the original text boundaries
Apply mean pooling to each chunk's token embeddings to produce chunk-level vectors

Because the transformer encoder uses self-attention across all tokens in the document, every token embedding already contains information from the surrounding context. When you then pool tokens belonging to a particular chunk, the resulting vector reflects not just the chunk's local content but also the document-wide context that flowed into those token representations.

The Key Insight: Self-Attention Is the Context Carrier

Embedding models like jina-embeddings-v2-base or jina-embeddings-v3 support context windows of 8,192 tokens or more. Inside the transformer encoder, each token attends to every other token in the input. This means a token representing the word "it" in chunk 5 has already incorporated information from chunk 1. When you pool those contextually enriched tokens, you get chunk embeddings that carry meaning from the entire document — not just the few sentences in the chunk.

Why Late Chunking Matters

Problem: Context Loss in Traditional Chunking

Consider a document that begins with: "John Smith was born in 1980. He founded Acme Corp in 2005." If you chunk this into two pieces — one about his birth and one about Acme Corp — the second chunk's embedding has no idea who "He" refers to. A query like "When did John Smith found his company?" may fail to retrieve the second chunk because the embedding doesn't connect "He" to "John Smith."

This problem compounds in longer documents with pronouns, references, abbreviations, and domain-specific terminology defined earlier in the text. Traditional chunking creates isolated semantic islands.

Solution: Contextualized Token Embeddings

Late Chunking solves this by letting the encoder see the whole document before any chunking happens. The token for "He" in the second chunk has already attended to "John Smith" in the first chunk. After pooling, the second chunk's vector encodes the connection between the pronoun and its referent.

Empirical Benefits

Experiments by Jina AI and independent practitioners have shown that Late Chunking can improve retrieval performance on benchmarks involving long documents with cross-chunk references. The gains are most pronounced when:

Documents contain pronouns or anaphoric references spanning chunk boundaries
Key definitions or context appear early and are referenced later
Chunks are small relative to the document length
The embedding model supports a long context window (8K+ tokens)

How to Use Late Chunking

Let's walk through a practical implementation. We'll use the transformers library with a Jina embeddings model, but the same principles apply to any long-context embedding model that exposes token-level outputs.

Prerequisites

Install the required packages:

pip install transformers torch

Step 1: Load a Long-Context Embedding Model

You need a model that supports long input sequences and returns per-token embeddings. Jina's models are a natural fit:

from transformers import AutoModel, AutoTokenizer

model_name = "jinaai/jina-embeddings-v2-base-en"
tokenizer = AutoTokenizer.from(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
model.eval()

Step 2: Define Your Document and Chunk Boundaries

First, split your document into chunks using whatever strategy you prefer — fixed-size, sentence-based, or semantic. The key difference from traditional chunking is that you keep track of the character or token boundaries so you can later slice the token embeddings.

document = (
    "John Smith was born in 1980 in Springfield. "
    "He studied computer science at MIT. "
    "After graduation, he founded Acme Corp in 2005. "
    "The company specializes in AI solutions for healthcare. "
    "It grew rapidly and went public in 2015."
)

# Simple sentence-based chunking
import re
sentences = re.split(r'(?<=\.)\s+', document)
chunks = []
start = 0
for sent in sentences:
    end = start + len(sent)
    chunks.append({"text": sent, "start": start, "end": end})
    start = end + 1  # +1 for the space

for i, c in enumerate(chunks):
    print(f"Chunk {i}: {c['text']}")

Step 3: Encode the Entire Document at Once

Tokenize the full document and pass it through the model to get token-level embeddings:

import torch

inputs = tokenizer(document, return_tensors="pt", return_offsets_mapping=True)
offset_mapping = inputs.pop("offset_mapping")  # (1, seq_len, 2)

with torch.no_grad():
    outputs = model(**inputs)

# Token embeddings: shape (1, seq_len, hidden_dim)
token_embeddings = outputs.last_hidden_state.squeeze(0)  # (seq_len, hidden_dim)
print(f"Token embeddings shape: {token_embeddings.shape}")

Step 4: Map Chunks to Token Ranges

Using the offset mapping from the tokenizer, determine which tokens belong to each chunk based on the character boundaries you recorded earlier:

def find_token_range(offsets, char_start, char_end):
    """Find token indices that fall within [char_start, char_end)."""
    token_indices = []
    for i, (start, end) in enumerate(offsets):
        if start == 0 and end == 0:
            continue  # skip special tokens
        if start >= char_start and end <= char_end:
            token_indices.append(i)
    return token_indices

offsets = offset_mapping.squeeze(0).tolist()

chunk_token_ranges = []
for chunk in chunks:
    token_range = find_token_range(offsets, chunk["start"], chunk["end"])
    chunk_token_ranges.append(token_range)
    print(f"Chunk '{chunk['text'][:30]}...' -> tokens {token_range}")

Step 5: Apply Late Pooling per Chunk

Now pool the token embeddings within each chunk's token range. Mean pooling is the most common choice:

import torch.nn.functional as F

def mean_pool(token_embeddings, token_range, attention_mask=None):
    """Mean pool token embeddings for a given token range."""
    if not token_range:
        return None
    chunk_embeddings = token_embeddings[token_range]  # (num_tokens, hidden_dim)
    if attention_mask is not None:
        weights = attention_mask[token_range].unsqueeze(-1).float()
        chunk_embeddings = chunk_embeddings * weights
        return chunk_embeddings.sum(dim=0) / weights.sum()
    return chunk_embeddings.mean(dim=0)

attention_mask = inputs["attention_mask"].squeeze(0)
chunk_vectors = []
for token_range in chunk_token_ranges:
    vec = mean_pool(token_embeddings, token_range, attention_mask)
    if vec is not None:
        vec = F.normalize(vec, p=2, dim=0)
        chunk_vectors.append(vec)

chunk_vectors = torch.stack(chunk_vectors)
print(f"Final chunk vectors shape: {chunk_vectors.shape}")

Step 6: Compare with Traditional Chunking

To see the difference, encode each chunk independently the traditional way:

traditional_vectors = []
for chunk in chunks:
    inputs_chunk = tokenizer(chunk["text"], return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs_chunk = model(**inputs_chunk)
    token_emb = outputs_chunk.last_hidden_state.squeeze(0)
    mask = inputs_chunk["attention_mask"].squeeze(0)
    pooled = mean_pool(token_emb, list(range(token_emb.shape[0])), mask)
    pooled = F.normalize(pooled, p=2, dim=0)
    traditional_vectors.append(pooled)

traditional_vectors = torch.stack(traditional_vectors)

# Compare similarity of chunk 1 ("He studied...") to query
query = "What did John Smith study?"
q_inputs = tokenizer(query, return_tensors="pt")
with torch.no_grad():
    q_outputs = model(**q_inputs)
q_vec = mean_pool(q_outputs.last_hidden_state.squeeze(0),
                  list(range(q_outputs.last_hidden_state.shape[1])),
                  q_inputs["attention_mask"].squeeze(0))
q_vec = F.normalize(q_vec, p=2, dim=0)

late_sim = F.cosine_similarity(chunk_vectors[1], q_vec, dim=0)
trad_sim = F.cosine_similarity(traditional_vectors[1], q_vec, dim=0)

print(f"Late chunking similarity:  {late_sim.item():.4f}")
print(f"Traditional similarity:    {trad_sim.item():.4f}")

In documents with cross-chunk references, the late chunking similarity will typically be higher because the chunk embedding retains the contextual link between "He" and "John Smith."

Using the chunked-json Library

Jina AI provides a convenience library called chunked-json (and the late-chunking package) that wraps this workflow. Here's how to use it:

pip install late-chunking

from late_chunking import chunker

# The chunker handles tokenization, encoding, and late pooling
result = chunker(
    document,
    model_name="jinaai/jina-embeddings-v2-base-en",
    chunk_size=128,
    overlap=0
)

for chunk in result:
    print(f"Text: {chunk['text'][:60]}...")
    print(f"Embedding dim: {len(chunk['embedding'])}")
    print()

This library abstracts away the offset mapping and pooling logic, making it easy to integrate into production pipelines.

Best Practices

Choose the Right Model

Late Chunking only works with models that support long context windows. Models with 512-token limits offer no benefit because you can't fit enough of the document into a single forward pass. Look for models with at least 8,192 token support. Good candidates include:

jinaai/jina-embeddings-v2-base-en (8,192 tokens)
jinaai/jina-embeddings-v3 (8,192 tokens)
Other long-context encoder models

Handle Documents Longer Than the Context Window

If your document exceeds the model's context window, you can apply Late Chunking in a sliding window fashion. Encode a window of the document, produce late-chunked embeddings for chunks fully contained within that window, then slide forward. Ensure windows overlap so chunks near boundaries still receive some cross-window context.

Balance Chunk Size

Smaller chunks give you finer retrieval granularity but may lose internal coherence. Larger chunks are more self-contained but may retrieve overly broad passages. A common sweet spot is 128–256 tokens per chunk. With Late Chunking, you can afford smaller chunks because context is preserved across boundaries.

Preserve Chunk Metadata

Even with contextualized embeddings, store metadata like chunk position, parent document ID, and character offsets. This enables re-ranking, parent-document retrieval, and debugging.

Normalize Embeddings

Always L2-normalize your final chunk vectors before storing them in a vector database. This ensures cosine similarity computations are efficient and consistent.

Benchmark on Your Own Data

Late Chunking is not a universal win. For short, self-contained chunks with no cross-references, traditional chunking may perform identically. Always evaluate on your specific retrieval task using metrics like recall@k and nDCG before committing to the approach in production.

Watch Out for Memory Usage

Encoding an 8,192-token document produces a tensor of shape (8192, 768) for a base model — about 25 MB in float32. This is manageable for batch sizes of 1, but be careful when processing many documents. Consider mixed precision (torch.float16) to halve memory usage.

Conclusion

Late Chunking is a simple yet powerful technique that rethinks the order of operations in embedding pipelines. By encoding the entire document before splitting it into chunks, you allow the transformer's self-attention to propagate context across chunk boundaries, producing embeddings that are richer and more faithful to the document's meaning. For RAG systems dealing with long, reference-heavy documents, this can translate directly into better retrieval quality and fewer missed results. While it requires a long-context embedding model and slightly more memory, the implementation is straightforward and the benefits are measurable. As with any retrieval optimization, the best approach is to implement it, benchmark on your data, and let the metrics guide your decision.