ParamMem: Augmenting Language Agents with Parametric Reflective Memory
| Authors | Tianjun Yao et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2602.23320 |
| Download | |
| Categories | cs.LG, cs.MA |
Abstract
Self-reflection enables language agents to iteratively refine solutions, yet often produces repetitive outputs that limit reasoning performance. Recent studies have attempted to address this limitation through various approaches, among which increasing reflective diversity has shown promise. Our empirical analysis reveals a strong positive correlation between reflective diversity and task success, further motivating the need for diverse reflection signals. We introduce ParamMem, a parametric memory module that encodes cross-sample reflection patterns into model parameters, enabling diverse reflection generation through temperature-controlled sampling. Building on this module, we propose ParamAgent, a reflection-based agent framework that integrates parametric memory with episodic and cross-sample memory. Extensive experiments on code generation, mathematical reasoning, and multi-hop question answering demonstrate consistent improvements over state-of-the-art baselines. Further analysis reveals that ParamMem is sample-efficient, enables weak-to-strong transfer across model scales, and supports self-improvement without reliance on stronger external model, highlighting the potential of ParamMem as an effective component for enhancing language agents.
Engineering Breakdown
Plain English
ParamMem addresses a key limitation in language agents: when they use self-reflection to improve their reasoning, they often produce repetitive reflections that actually hurt performance. The paper shows empirically that more diverse reflections correlate strongly with better task success, then proposes ParamMem—a parametric memory module that encodes different reflection patterns into model parameters and generates diverse reflections through temperature-controlled sampling. They build ParamAgent, an agent framework combining this parametric memory with episodic and cross-sample memory, achieving improvements on reasoning tasks by producing varied, higher-quality reflections across iterations.
Core Technical Contribution
The core innovation is treating reflection diversity as a learnable, parametric property rather than a prompt-engineering problem. ParamMem encodes cross-sample reflection patterns directly into model parameters during training, then leverages temperature-based sampling at inference to generate diverse reflections without retraining. This is fundamentally different from prior approaches that either use fixed reflection prompts or rely on stochastic sampling of the same base model—here, the diversity is baked into the learned parameter space itself. The contribution also includes the ParamAgent framework that orchestrates parametric, episodic, and cross-sample memory sources into a unified reflection strategy.
How It Works
ParamMem works by first collecting reflection traces across multiple problem-solving trajectories during training, then learning a parametric encoding of these reflection patterns into dedicated model parameters (similar to how LoRA or adapter modules encode task-specific knowledge). At inference time, the agent generates reflections by sampling from this parameter space with temperature control—higher temperature produces more diverse reflections, lower temperature produces more consistent ones. The ParamAgent wraps this with episodic memory (remembering past problem-solving episodes) and cross-sample memory (comparing current attempts against solutions to similar problems), feeding these signals back into the parametric reflection module. Each reflection generation step uses the current task state and memory context to condition sampling from the learned reflection parameter space, producing outputs tuned to the specific problem without requiring new fine-tuning.
Production Impact
For production systems, ParamMem solves a concrete problem: language agents that loop on reflection often get stuck in local minima, repeatedly making the same reasoning mistakes. By enabling diverse reflections, you get agents that explore more solution space before settling on an answer, directly improving success rates on complex reasoning tasks. The parametric approach is also compute-efficient—instead of running multiple forward passes through a full model or storing large episodic databases, you store a relatively small set of learned parameters. However, integration requires on-device parameter storage and careful tuning of temperature schedules; you'll need to validate that learned reflection patterns transfer to your specific domain, and initial setup requires collecting diverse reflection traces from successful problem-solving episodes (non-trivial for new domains).
Limitations and When Not to Use This
The paper assumes that reflection patterns learned on training tasks will generalize to new test tasks, which may not hold when task distribution shifts significantly. ParamMem requires offline collection of diverse, successful reflection trajectories—this is expensive and assumes you have access to good problem solvers initially, limiting applicability in domains where generating quality reflections is itself hard. The approach also doesn't address cases where the core reasoning capability is weak; better reflections can't fix fundamentally wrong model knowledge. Finally, the paper doesn't deeply explore the computational cost of maintaining and querying parametric memory at scale, or how this interacts with quantization and other efficiency techniques used in production LLM deployments.
Research Context
This work builds on a growing body of research showing that self-reflection improves language agent reasoning (prior work like ReAct, Chain-of-Thought), but extends it by systematically studying reflection quality and diversity. It fits into the broader trend of moving away from pure prompt engineering toward learned, parametric mechanisms for agent control. The paper likely evaluates on standard reasoning benchmarks (math, planning, or code tasks) where iterative refinement has proven valuable. It opens research directions around learned memory mechanisms for agents, the interplay between parameter-efficient adaptation and agent reasoning, and whether other aspects of agent behavior (planning, memory retrieval) could benefit from similar parametric diversity approaches.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::