Ditch Naive Chunking: Late Chunking RAG in Spring AI

#java #ai #llm #systemdesign

Ditch Naive Chunking: Late Chunking RAG in Spring AI

Naive text chunking breaks your RAG pipeline by slicing context at arbitrary token boundaries before your embedding model ever sees the text. Late chunking fixes this structural flaw by running a long-context transformer over the full document first, then pooling chunk vectors directly from fully contextualized token states.

Why Most Developers Get This Wrong

Splitting documents into isolated 512-token slices with Spring AI's TokenTextSplitter prior to embedding, completely destroying document-level semantics and cross-paragraph references.
Adding heavy LLM contextual-summarization calls or bloated chunk overlaps to salvage lost context, inflating API costs and query latency for zero structural gain.
Blaming retrieval performance on vector database distance metrics when the actual root cause is context degradation at the pre-embedding boundary.

The Right Way

Late chunking encodes full document context into token-level representations before slicing vector spaces into chunk boundaries.

Pass raw, full-length documents into a long-context transformer (e.g., jina-embeddings-v3 with an 8k token window) via extended Spring AI EmbeddingModel pipelines.
Extract raw contextualized token embeddings from the transformer's final layer for the full document in a single forward pass.
Compute final chunk embeddings by mean-pooling token vectors belonging strictly to each target text span's token offset boundaries.
Persist contextualized vectors to PgVectorStore or QdrantVectorStore with zero changes required for downstream similarity retrieval logic.

Show Me The Code

@Service
public class LateChunkingService {
    private final LongContextEmbeddingModel embeddingModel; // Wraps 8k+ context transformer

    public List<VectorChunk> processDocument(Document document) {
        // 1. One forward pass over FULL text generates contextualized token vectors
        TokenEmbeddings fullDocTokens = embeddingModel.encodeDocumentTokens(document.getText());

        // 2. Mean-pool pre-computed token vectors using chunk boundary spans
        return document.getSpans().stream()
            .map(span -> new VectorChunk(
                span.text(),
                fullDocTokens.poolSpan(span.startToken(), span.endToken())
            )).toList();
    }
}

Key Takeaways

Early chunking strips document context before vectorization; late chunking captures global semantics first and slices embeddings later.
You eliminate redundant LLM token overhead while dramatically boosting retrieval precision across multi-page enterprise documents.
Spring AI implementation only requires adjusting the embedding extraction and pooling stage—your vector store schemas remain 100% untouched.