Databricks’ Mosaic AI Research team has added a new framework, MemAlign, to MLflow, its managed machine learning and generative AI lifecycle development service.

MemAlign is designed to help enterprises lower the cost and latency of training LLM-based judges, in turn making AI evaluation scalable and trustworthy enough for production deployments.

The new framework, according to the research team, addresses a critical bottleneck most enterprises are facing today: their ability to efficiently evaluate and govern the behavior of agentic systems or the LLMs driving them, even as demand for their rapid deployment continues to rise.

Traditional approaches to training LLM-based judges depend on large, labeled datasets, repeated fine-tuning, or prompt-based heuristics, all of which are expensive to maintain and slow to adapt as models, prompts, and business requirements change.

As a result, AI evaluation often remains manual and periodic, limiting enterprises’ ability to safely iterate and deploy models at scale, the team wrote in a blog post.

MemAlign’s memory-driven alternative to brute-force retraining

In contrast, MemAlign uses a dual memory system that replaces brute-force retraining with memory-driven alignment based on human feedback from human subject matter experts, although fewer in number and frequency than conventional training methods.

Instead of repeatedly fine-tuning models on large datasets, MemAlign separates knowledge into a semantic memory, which captures general evaluation principles, and an episodic memory, which stores task-specific feedback expressed in natural language by subject matter experts, depending on the use case.

This allows LLM judges to rapidly adapt to new domains or evaluation criteria using small amounts of human feedback, while retaining consistency across tasks, the research team wrote.

This reduces the latency and costs required to reach more efficient and stable levels of judgment, making the approach more practical to adapt for enterprises, the team added.

In Databricks-controlled tests, MemAlign was able to show the same efficiency as labeled datasets.

Analysts expect the new framework to benefit enterprises and their development teams.

“For developers, MemAlign helps reduce the brittle prompt engineering trap where fixing one error often breaks three others. It provides a delete or overwrite function for feedback. If a business policy changes, the developer can update or overwrite relevant feedback rather than restarting the alignment process,” said Stephanie Walter, practice leader of AI stack at HyperFRAME Research.

Walter was referring to the framework’s episodic memory, which is stored as a highly scalable vector database, enabling it to handle millions of feedback examples with minimal retrieval latency.

The ability to keep LLM-based judges aligned with changing business requirements, according to Moor Insights and Strategy principal analyst Robert Kramer, is a critical ability as it doesn’t destabilize production systems, which is increasingly important for enterprises as agentic systems scale.

Agent Bricks may soon get MemAlign

Separately, a company spokesperson told InfoWorld that Databricks may soon embed MemAlign to its AI-driven agent development interface, Agent Bricks.

More so because the company feels that the new framework would be more efficient in evaluating and governing agents built on the interface than previously introduced capabilities, such as Agent-as-a-Judge, Tunable Judges, and Judge Builder.

Judge Builder, which was previewed in November last year, is a visual interface to create and tune LLM judges with domain knowledge from subject matter experts and utilizes the Agent-as-a-Judge feature that offers insights into an agent’s trace, making evaluations more accurate.

“While the Judge Builder can incorporate subject matter expert feedback to align its behavior, that alignment step is currently expensive and requires significant amounts of human feedback,” the spokesperson said.

“MemAlign will soon be available in the Judge Builder, so users can build and iterate on their judges much faster and much more cheaply,” the spokesperson added.