We Trained Our Own Embedding Model for Legal RAG. Here's What It Actually Took.

Short answer: We built a custom embedding model called Tenet 5 because off-the-shelf models could not align legal research queries with the right document chunks without absurdly high top-k. Training our own cut token consumption for RAG-based agents by about a third and improved quality by roughly 10% across five metrics. It took six months, one H100 node, and an evaluation harness we trusted more than the model.

Most teams should not do this. This post is partly about why we did it anyway, and mostly about how to tell which camp you are in.

I've written before about the business case — how this work cut our inference costs by around 50%. This is the engineering companion to that post: what actually broke, what we tried that failed, and what we shipped.

The three problems we were actually solving

Not "our RAG isn't good enough." Three specific, measurable failures:

1. Sub-agents using RAG consumed far too many tokens. In an agentic system this compounds — every sub-agent call carries its own retrieved context.

2. Off-the-shelf embedding models needed ridiculously high top-k to hit the quality bar our users expected.

3. Those models are not trained to align legal queries with the document chunks that answer them. This is the root cause; the first two are symptoms.

We tested many open and proprietary models. They all failed the same way, and the reason is specific enough to be worth spelling out.

Why off-the-shelf embeddings fail on legal queries

A legal research query is a heavy mix of legal metadata and legal context — and those two things are rarely found together in the document text itself.

Here is the part that breaks vector search: two paragraphs can look nearly identical in isolation and mean completely different things in the context of the case. Same words. Different tribunal, different year, different governing treaty. One is authority for your argument, the other is authority against it.

If a human were doing this retrieval, it would take expert knowledge of section-level context, a deep understanding of legal metadata to filter on, and the ability to hold the whole document in mind while comparing specific sections of it.

General-purpose embedding models are trained to capture semantic similarity. In this domain, semantic similarity is close to the wrong signal — it is exactly what makes the two paragraphs collapse onto each other in vector space.

This is the wall I wrote about in every AI team hits a wall — ours had RAG written all over it. This post is what was on the other side of it.

What we tried first (and why it didn't work)

Before training anything, we did what everyone does. For the record, so you can skip it:

Better vector similarity. Different distance metrics, different chunking strategies. No meaningful movement — the problem was the representation, not the comparison.

Metadata extraction and filtering. Reasonable in principle. In practice, legal queries do not decompose cleanly into filters, and the ones that do still need the context to rank what survives the filter.

Stuffing document summaries into the chunk text. This one is worth flagging because it is widely recommended and it actively hurt us. Prepending summaries reduced overall quality by disrupting the latent space of the embeddings — you are pushing every chunk in a document toward the same region of vector space, which is the opposite of what retrieval needs.

Raising top-k. This one worked, which is exactly the trap.

The high-top-k trap

Cranking top-k works because modern LLMs are genuinely good at ignoring irrelevant context. Retrieve 50 chunks instead of 10, and the model will usually find the right one and disregard the rest.

So your quality metrics recover, and the problem looks solved. It isn't. You have moved a retrieval failure into your token bill, where it is invisible until someone reads the invoice.

Two things are worth knowing about that trade:

It is expensive and it compounds. In an agentic system, every sub-agent inherits the bloat. This is why RAG token costs tend to grow super-linearly with agent complexity rather than in step with traffic.

"LLMs ignore noise" has limits. The Lost in the Middle research (Liu et al., TACL) found model performance is highest when relevant information sits at the very start or very end of the context, and degrades noticeably when it is buried in the middle. High top-k does not just cost money — past a point it starts costing accuracy too, and it does so quietly.

High top-k is a legitimate short-term fix. It's a bad long-term architecture.

What we actually built

We used our own data to train an embedding model that aligns document chunks with three things at once: document-level legal metadata, document summaries, and the exact kind of queries lawyers naturally ask.

That last piece matters more than it sounds. Most embedding training uses generic query-document pairs. We had the real queries.

Two techniques did the heavy lifting:

Metadata dropout

We applied a dropout-like method during training so the embeddings stay robust when metadata is missing, malformed, or of a kind the model has not seen. Production legal corpora are not uniform — if your model only performs when all 15+ metadata fields are clean and present, you have built something that works in evaluation and fails on real documents.

Purpose-designed negative sampling

We designed negative sampling specifically to shape a latent space that resembles what an actual lawyer would do. This is the core of it. In contrastive learning, your negatives define the geometry — the model learns what "different" means from the examples of difference you give it. Feed it random negatives and it learns to separate obviously unrelated things, which it could already do. Feed it the two near-identical paragraphs that differ only by tribunal, and it learns the distinction you actually need.

This is well-trodden ground in the literature if you want to go deeper. NVIDIA's NV-Retriever work on positive-aware hard-negative mining is the clearest treatment I know of the central hazard: mine negatives too aggressively and you start sweeping up false negatives — passages that genuinely answer the query but were not labelled as such — and actively teach your model the wrong thing. Their fix uses the positive relevance score as an anchor to filter candidates. Worth reading before you design your own sampler.

The practical setup: six months of iteration on a single H100 node, with Hugging Face and Cursor. No enormous cluster. The constraint was never compute — it was knowing whether each change made things better.

The part that actually mattered: measuring it

If there is one thing to take from this post, it is this. The evaluation harness was harder and more valuable than the model.

Every update to our agent, however small, is gauged against five strict quality metrics. The acceptance bar in legal AI is high and we don't compromise on it — a plausible-sounding wrong citation is worse than no answer.

We tested on ~500 gnarly real-world legal queries through a three-step QA process:

1. LLM-as-a-Judge, via our internal tool Anchor

2. Retrieval benchmarking — measuring retrieval directly, not just end-to-end answers

3. Manual QA by people who know the domain

And the step I'd argue is non-negotiable: we correlate our LLM-as-a-Judge scores against human scores to keep the metrics honest and catch LLM overconfidence.

That is not paranoia. Research on the limits of LLM-as-a-judge without human grounding found judge models agree with human annotators far less than teams assume. For GPT-4o as judge, agreement with expert annotators sat at a Cohen's kappa of about 0.46 with no reference answer, rising to roughly 0.69 when given human-written references (kappa is chance-corrected, so these are lower than raw percentages would suggest).

The sharpest finding is the one that should worry anyone building in a specialist domain: judges agreed with experts at kappa 0.78 on questions the judge could itself answer correctly, and 0.30 on questions it could not. A judge that cannot answer the question cannot reliably grade the answer.

For a general-purpose chatbot, that's a tolerable error bar. For specialist legal research — where every query is the hard tail — an ungrounded LLM judge will confidently tell you a regression is an improvement. Correlating against human scores is what turns the judge from a vibe into an instrument.

The results

• Token consumption for RAG-based agents: down by about a third

• Quality: up ~10% across all five core metrics

The quality number is the one I care about. Cost reductions that trade away quality are easy and worthless — you can get them by lowering top-k and calling it a day. Moving both in the right direction at once is the whole point.

Should you train your own embedding model?

Probably not. Here's the honest test.

Train your own when:

• Your domain has a structural mismatch with semantic similarity — meaning depends on metadata or context that isn't in the chunk text (law, medicine, finance, internal jargon-heavy corpora)

• You're compensating with high top-k and can measure what that's costing you

• You have real user queries in volume — this is the asset, more than the compute

• You have an evaluation harness you trust, ideally before you start

• Retrieval quality is core to your product, not a feature of it

Don't when:

• You haven't yet tried a good off-the-shelf model with sensible chunking — most teams are leaving a lot on the table before they get anywhere near training

• You can't measure retrieval quality separately from answer quality

• Nobody on the team can tell a real improvement from noise across five metrics

• You're doing it because it sounds impressive

Point four is where most attempts die. Six months of iteration is only possible if each iteration produces a trustworthy signal. Without that you aren't training a model, you're doing a very expensive random walk.

If you're earlier in the process, five ways to test your AI readiness without coding or maths is a better starting point than this post.

Frequently asked questions

When should you train a custom embedding model instead of using an off-the-shelf one? When your domain has a structural mismatch with semantic similarity — where meaning depends on metadata or document-level context not present in the chunk text — and you are compensating with high top-k. You also need real user queries in volume and an evaluation harness you trust. If retrieval quality is core to your product rather than a feature of it, the investment can pay for itself in token costs alone.

Why do off-the-shelf embedding models fail on domain-specific retrieval? They are trained to capture general semantic similarity. In specialist domains, two passages can be nearly identical in wording and mean opposite things depending on context — the tribunal, the year, the governing instrument. General models collapse those passages onto each other in vector space, so retrieval cannot separate them at any top-k.

Does adding document summaries to chunks improve RAG retrieval? Not always, and in our case it hurt. Prepending document summaries to chunk text pushed every chunk in a document toward the same region of the latent space, reducing the model's ability to distinguish between them. It's a commonly recommended technique that is worth measuring rather than assuming.

Is raising top-k a good fix for poor retrieval? It works in the short term, because modern LLMs are good at ignoring irrelevant context. But it converts a retrieval failure into a token cost that compounds across sub-agents, and research on long-context behaviour shows accuracy still degrades when relevant information is buried mid-context. It's a reasonable stopgap and a poor architecture.

How reliable is LLM-as-a-Judge for evaluating RAG quality? Less than most teams assume, and least reliable on exactly the hardest questions. Published work measuring GPT-4o as a judge found agreement with expert annotators at a Cohen's kappa of roughly 0.46 without reference answers and 0.69 with human-written ones — and, more starkly, 0.78 on questions the judge could answer itself versus 0.30 on questions it could not. Correlating judge scores against human scores on a sample is the practical fix.

What's next in this series

This is the first of three posts on what's inside our legal AI stack. Coming up: more on the five quality metrics we gauge every change against, and how the agent architecture around this model works.

Keep reading

• How my team reduced our LLM inference cost by about 50% — the business case for this work

• Every AI team hits a wall — ours had RAG written all over it

• The reality of agentic coding: how we are turning everyone into a builder

• Anthropic's Model Context Protocol (MCP): I am not convinced yet