The Return of Wake-Sleep
Harvey’s research demonstrates that a wake-sleep agent training loop can boost all-pass performance on legal tasks from 2.9% to 15.7%.
Deep dive
- Methodology: Used 196 tasks from the Legal Agent Benchmark (LAB).
- Performance: Increased pass rates across both familiar and new, unseen legal matters.
- Feedback Loop: Used dual LLM judges to grade performance without the agent seeing the grades, ensuring unbiased evaluation.
- Knowledge Synthesis: Sleep phase distilled raw trajectories into reusable checklists and best practices.
- Optimization: Memory retrieval reduced token costs by 50% while maintaining performance by selecting only relevant lessons per task.
Decoder
- Wake-sleep: A machine learning algorithm where the agent learns online (wake) and then uses its experience to generate and refine knowledge offline (sleep).
- Held-out tasks: Data points withheld from the training phase to test the agent's ability to generalize to new information.
Original article
The Return of Wake-Sleep
Wake-sleep is a classic learning algorithm that trains an AI system by alternating between two phases. In the wake phase, the system acts online and observes real data, such as images, audio, or text. In the sleep phase, it "dreams up" examples offline that might improve online behavior. While originally introduced to train neural networks, later systems used the same concept to build up reusable knowledge (e.g. DreamCoder, LILO and DreamProver).
As production, knowledge-intensive agent use cases proliferate, hybrid online-offline protocols like wake-sleep are increasingly becoming a part of the agent stack. Long-horizon agents now work on individual tasks for hours. As these agents work, they leave a record of actions and deliverables, and with the right verifiers we can tell which actions led to good deliverables and which fell short. That record can be reviewed offline, by LLMs or other agents, and transformed into memories or structured context that the agent can leverage on its next task.
Agent research has started moving in this direction, exploring a range of approaches for persisting notes, workflows, or playbooks across repeated agent interactions (see Reflexion, Agent Workflow Memory, Dynamic Cheatsheet, ACE, etc). Anthropic, with Dreaming in Claude Managed Agents, is a current leader in building frameworks for general-purpose agent context and "sleep-time" context management. Cognition also recently introduced dreaming in its Agent Memory Repo release.
Motivated by these directions, we explored how wake-sleep might apply to long-horizon legal agents. We evaluated our own variant of wake-sleep on tasks from Harvey’s Legal Agent Benchmark (LAB) and found that:
- Wake-sleep meaningfully improves agent performance on complex tasks.
- The memory and context written offline by the sleep phase generalizes to new matters and kinds of task the agent did not see in training.
- Allowing the agent retrieve only the memories relevant to each task, rather than reading all of them, halves the cost per task at about the same average quality.
Wake-sleep for legal agents
So what might a wake-sleep implementation look like for long-horizon agents? In our version of wake-sleep, each learning cycle proceeds through the following stages:
1. Agent rollout. In the wake phase, an agent works through a set of LAB tasks. Each task gives the agent instructions and a folder of documents from a synthetic client matter, and asks for a deliverable such as an issues list, a memo or a marked-up draft. The agent reads and searches across the documents, runs analyses, and writes and edits the final deliverables.
2. Grading. Dual LLM judges grade the deliverables against expert rubrics as in the standard LAB configuration. The agent does not see the grades.
3. Generate candidate lessons. In the sleep phase, a review model reads each graded training run, with the agent’s trajectory, its deliverables, and the verdict on every rubric criterion, and writes candidate lessons about what went wrong. A lesson is either a checklist item, which states something a deliverable must contain, or a practice note, which states how to do part of the work.
4. Merge. The review model then merges the candidate lessons into the agent's memory. It combines lessons that overlap, revises existing lessons and drops some. From the second cycle on, it also judges which existing lessons helped.
5. Commit. The biggest risk in lesson generation is overfitting—e.g. writing down facts specific to one training matter, such as a party’s name or a purchase price, instead of learnings that carry to other matters. To prevent this, we implemented a commit gate that uses Jev to reject lessons that contain client-specific detail, are too narrow, or rest on a single training task. Rejected lessons go to a review queue and are not added to the memory.
6. Store. The memory stores the lessons as text, grouped by kind of work (analysis, drafting and review), with a fourth group for general learnings.
When initialized, the agent's memory contains no lessons or structured context. In each successive wake-sleep cycle, the agent's memory store is updated with the lessons and structured context it learned in the prior cycle.
Experiment and results
We ran a 10-cycle wake-sleep process across 196 tasks from LAB’s Corporate M&A and Capital Markets practice areas (110 for training, 86 for validation), using the steps described in the previous section. We then compared the rubric pass rates between the wake-sleep agent and a vanilla agent with no memory or internalized context across each of the tasks. For these initial experiments, GPT-6 Luna was used as the agent model.
To test generalization, we grouped hold-out tasks into two cohorts based on their distance to the training task distribution:
- Familiar matters: Tasks similar to those found in the training distribution, with similar document types and source instructions to the training tasks.
- New matters: Fundamentally different tasks than the agent encountered in the wake-sleep training cycles.
Results from these runs are shown in Figure 3. Overall, we found wake-sleep to improve performance significantly across all tasks, reaching 15.7% all-pass, up from 2.9% for a vanilla agent. These all-pass gains are attributable to a large increase in rubric-level pass rates across both of our experimental conditions, with the wake-sleep agent completing over 10% more rubric criteria per task on familiar matters, and nearly 5% more rubric criteria per task on new matters relative to the baseline agent.
Unsurprisingly, the biggest per-step gains come from the initial wake-sleep cycles, where the agent picks up its first lessons and best practices for legal knowledge work.
Qualitative Findings
To better understand the impact of wake-sleep, and the evolution of learned memory and context, we examined at agent behavior (online) and memory (offline) across each of the first 10 wake-sleep cycles.
In the first sleep cycle, the sleep process adds 122 lessons to agent memory. The majority of these initial lessons are checklist items that indicate what a high-quality deliverable must contain content-wise. Over the remainder of the sleep cycles, these lessons are refined into both behavioral best practices -- e.g., "When a term changes, explain how each alternative works" -- and tactical instructions -- e.g., "Recompute every figure yourself to verify it". Unhelpful or overly-specific lessons are pruned.
A qualitative example of this progression is shown in Figure 5, which follows one lesson about deadlines and time-related calculations that the agent learned and refined at sleep-time. In the first sleep cycle, the review model writes a rule to state contractual time terms as durations rather than start and end dates, with instructions on how to compute such durations:
State every contractual time term expressly as a duration, such as survival, lease term or notice period, not only as start and end dates. For every future deadline, compute the days remaining from the memo date, show the arithmetic, and put it in the bottom line or in the executive summary.
Over the next nine sleep cycles, the review model added a rule for each new case it met, such as a deadline that falls on a weekend or two documents that set different windows for the same delivery.
These lessons also change online agent behavior by influencing how the agent works, as shown in Figure 6. Using these lessons and best practices, the agent made ~2.5x more tool calls per task to run code, check its work, and draft and edit deliverables in accordance with its developing memory. The agent made about as many calls to read and search documents as before, indicating that these initial gains primarily target post-research work like analysis and writing. An interesting direction of future study is to identify lessons and best practices that improve agent behavior with respect to search and retrieval.
Importantly, this extra analysis and work-checking makes runs slower and more expensive. The wake-sleep agent is trading cost and latency for more thorough, higher-quality work, which raises the question of whether it can maintain that quality while working more efficiently.
Making wake-sleep more efficient
To improve the efficiency of wake-time agent behavior, we evaluated lesson retrieval. To implement this, we again use Jev, this time asking it if each lesson is relevant and helpful the task at hand. Only the lessons Jev judges most likely to help are injected into the agent's context.
We found that retrieving a fixed set of relevant lessons from the sleep bank halved the agent's cost per task while maintaining its rubric criteria pass rate. This suggests that wake-sleep, similar to other agent paradigms, benefits from online context management optimizations.
What we learned about wake-sleep for agents
Production agent deployments will increasingly combine online work with offline processing. Agents will act on tasks in real time, while background processes index documents, summarize past sessions and curate what the agent remembers. Algorithms like wake-sleep, which separate doing the work online from learning from it offline, are a natural next step for agent architectures.
In this work, we explored wake-sleep in a long-horizon legal agent setting, using primarily non-parametric forms of agent optimization (textual representations of context and best practices), but these techniques can be extended to parametric learning and post-training, as in the original wake-sleep formalization.
For verifiable knowledge-intensive work, this loop offers a way for agents to get better over time.
Fresh dev news, AI-digested daily. Never miss a trend at devoured.fyi