Memory-Aligned LLM Judges

AI Evaluation & Safety advanced Building with AI ~32 min read

Overview

LLM judges are useful because they turn expensive human review into scalable automated evaluation. The failure mode is that a generic judge often scores against generic language-quality standards rather than the domain policy, product constraints, or expert taste that actually define quality. A response can be syntactically correct, fluent, and complete while still violating a business rule that only subject-matter experts know.

Memory-aligned judging addresses that gap without immediately fine-tuning model weights. Experts review a small number of examples, write natural-language feedback, and the evaluator distills that feedback into persistent memories. Future judge calls retrieve those memories into the evaluation context, so the judge applies accumulated principles and concrete precedents instead of starting from the same static rubric every time.

The MemAlign approach from Databricks and MLflow is the reference design for this pattern. It separates reusable semantic principles from episodic edge cases, then constructs a working-memory context for each new judgment. The important engineering idea is transferable beyond one product: treat judge alignment as a versioned memory system, not just as a longer prompt or a one-off prompt-optimization run.

Core Concepts

Natural-Language Feedback as a Dense Signal

A scalar label says whether the judge was wrong. Expert feedback says why it was wrong, which constraint was missed, and how similar cases should be handled later. That makes each reviewed example more informative than a bare pass/fail label, especially early in a project when there are only a handful of reviewed traces.

Semantic Memory

Semantic memory stores generalized evaluation rules. Examples include “prefer certified data views over raw operational tables” or “evaluate safety by intent, not just wording.” These records should be concise, non-duplicative, and broad enough to apply across many future examples.

Episodic Memory

Episodic memory stores specific cases where the judge made a mistake or where a rule is hard to generalize. These records provide precedent. When a new example resembles a stored episode, the judge can inspect the prior input, expert correction, and final label before scoring.

Key Concept: Memory Scaling

Memory scaling improves judge quality by accumulating durable feedback over time. It is different from test-time scaling, which spends more compute per request, and different from prompt optimization, which repeatedly searches for a better static prompt. The judge gets better because the evaluation system remembers what experts already taught it.

How It Works

Alignment Stage

The alignment stage starts with traces where the judge output, expert feedback, and corrected target are available. A memory updater extracts general principles for semantic memory and saves salient examples into episodic memory. The update should be auditable: every memory record needs a source trace, a timestamp, and a version so teams can explain why a judge changed.

type FeedbackTrace = {
  input: string;
  judge_output: string;
  expert_feedback: string;
  expert_label: string;
};

type SemanticRule = {
  id: string;
  rule: string;
  source_trace_id: string;
  version: number;
};

type Episode = {
  id: string;
  input: string;
  mistake: string;
  correction: string;
  embedding: number[];
};

async function alignJudge(trace: FeedbackTrace) {
  const rule = await extractGeneralRule(trace.expert_feedback);
  await semanticStore.upsert(rule);

  if (isEdgeCase(trace)) {
    const episode = await buildEpisode(trace);
    await episodicStore.upsert(episode);
  }
}

Judgment Stage

At evaluation time, the system gathers all active semantic rules or a curated subset, retrieves similar episodes, and builds a bounded working-memory context. The judge prompt then contains the rubric, current input, relevant principles, and precedent examples. The model still performs the judgment, but it is now grounded in expert corrections.

async function scoreWithMemory(input: string, response: string) {
  const semanticRules = await semanticStore.activeRules();
  const episodes = await episodicStore.search(input, { topK: 5 });

  return judgeModel.generate({
    rubric: baseRubric,
    principles: semanticRules,
    precedents: episodes,
    input,
    response,
    outputSchema: {
      score: "number",
      pass: "boolean",
      reason: "string",
      applied_memory_ids: "string[]"
    }
  });
}

Memory Maintenance

Alignment memory is not append-only truth. Policies change, experts disagree, and privacy rules may require deletion. A production implementation needs overwrite, deletion, conflict detection, and replay support so teams can compare judge behavior before and after memory changes.

Practical Applications

Domain-Specific Quality Gates

Customer-support, healthcare, finance, and internal-data assistants often have local quality standards that are not captured by public benchmark rubrics. Memory-aligned judges let experts encode those standards incrementally as they review real failures.

Agent Evaluation Loops

Agents generate tool calls, plans, retries, and final answers. A memory-aligned judge can learn recurring failure patterns such as using uncertified data, skipping approval steps, or over-optimizing for answer speed at the expense of safety.

Interactive Evaluation Development

Prompt optimizers and fine-tuning runs can take long enough to discourage tight iteration. A memory-aligned evaluator can incorporate a handful of expert corrections quickly, then expose the impact in the next evaluation batch.

Implementation Patterns

Version Every Judge Artifact

Store the base rubric, model version, semantic memories, episodic memories, retrieval settings, and output schema with each evaluation run. Once memory changes, old and new scores are not strictly comparable unless the memory snapshot is part of the trace.

{
  "judge_id": "support_quality_v4",
  "model": "judge-model-2026-02",
  "rubric_hash": "sha256:...",
  "semantic_memory_version": 17,
  "episodic_memory_index": "episodes_2026_02_03",
  "retrieval_top_k": 5,
  "output_schema_version": 2
}

Separate Principles from Precedents

Do not dump every reviewed example into the judge prompt. Promote only stable, generalizable lessons to semantic memory. Keep edge cases in episodic memory and retrieve them by similarity, metadata, or failure category.

Track Memory Use

Ask the judge to report which memory IDs influenced the score. This creates a debugging path when a memory record is stale, contradictory, or over-applied to unrelated examples.

Performance Considerations

Common Pitfalls

Further Reading