Agent memory is often described as a storage and retrieval problem. Two common answers—parametric memory and retrieval-based memory—each solve part of the problem and leave another part exposed. MemGen takes a third route. It does not store experience as explicit text or bake it permanently into weights. Instead, it generates latent memory during inference and injects it at points where reasoning changes direction. The point is not a bigger memory store; it is a tighter loop between reasoning and memory formation.
Why a third memory path is needed
Parametric memory encodes experience into model weights through fine-tuning. It internalizes knowledge deeply and adds no inference-time overhead, but new training can overwrite old knowledge, and full fine-tuning is expensive. Retrieval-based memory keeps experience outside the model, often in a vector store, and retrieves it at inference time. It avoids forgetting and can cover large knowledge bases, yet retrieval is usually a one-time step before reasoning, so the memory is static and can drift away from each reasoning step. Retrieval also adds its own overhead.
<table> <thead> <tr> <th>Memory type</th> <th>Implementation</th> <th>Strengths</th> <th>Main limitation</th> <th>Typical use</th> </tr> </thead> <tbody> <tr> <td>Parametric memory</td> <td>Fine-tune experience into model parameters</td> <td>Deep internalization; no extra inference cost</td> <td>Catastrophic forgetting; high fine-tuning cost</td> <td>Simple tasks, static knowledge</td> </tr> <tr> <td>Retrieval-based memory</td> <td>Store experience externally; retrieve during inference</td> <td>Avoids forgetting; scales to large knowledge</td> <td>Static one-shot retrieval; disconnected from reasoning; retrieval cost</td> <td>Knowledge-base QA, static lookup</td> </tr> </tbody> </table>An agent handling multi-turn dialogue or long-horizon planning needs memory to do more than store and fetch. It needs dynamic coupling: memory should interact with every reasoning step, not arrive once before reasoning starts. It needs generative reconstruction: memory should be formed for the current need, not mechanically extracted from raw experience. It also needs to stay lightweight: memory quality should improve without uncontrolled parameter growth or compute cost. MemGen is built around these requirements.
Architecture: reasoner, trigger, and weaver in a loop
The design combines modular coordination with dynamic memory generation. Three modules—the Reasoner, the Trigger, and the Weaver—work with projection layers that align embedding spaces. The flow is:
input sequence → Reasoner produces initial embeddings → Trigger decides when memory should fire → Weaver generates latent memory → projection layers align dimensions → augmented sequence goes back to Reasoner → iterate
The central idea is that memory is not a fixed piece of stored content. It is an implicit representation tailored to the current reasoning step. Each step can trigger memory generation and injection.
Reasoner
The Reasoner is the core inference engine. It receives the augmented sequence—original input plus latent memory—and produces the output. It is built on a pretrained LLM such as LLaMA or Qwen. Its parameters are frozen during training; only the Weaver and Trigger are trained. This avoids catastrophic forgetting in the base model and lowers training cost. The default precision is bfloat16 to balance memory use and performance, and Flash Attention 2 accelerates attention over long sequences. Two projection layers, reasoner_to_weaver and weaver_to_reasoner, map between the Reasoner and Weaver hidden spaces when their dimensions differ.
Trigger
The Trigger decides when memory generation is worth doing. It prevents meaningless memory injection and keeps augmentation on demand. There are two implementations.
<table> <thead> <tr> <th>Trigger type</th> <th>Core logic</th> <th>Strengths</th> <th>Use case</th> </tr> </thead> <tbody> <tr> <td>NanoTrigger</td> <td>Minimal implementation; always returns a trigger decision</td> <td>No training, fast deployment, no extra overhead</td> <td>Quick tests, baseline comparison, simple tasks</td> </tr> <tr> <td>MemGenTrigger</td> <td>Pretrained LLM plus a binary classification head; decides dynamically</td> <td>Handles complex scenarios; higher trigger precision; supports PEFT</td> <td>Real deployments, complex reasoning</td> </tr> </tbody> </table>MemGenTrigger replaces the LLM's original language-model head with a binary head whose output dimension is 2, corresponding to no-trigger and trigger. It supports LoRA and other PEFT methods, so only a small number of parameters need training for a specific task. It also uses bfloat16 and Flash Attention 2 to stay consistent with the Reasoner.
Weaver
The Weaver is the main innovation. It generates latent memory that fits the current reasoning need. It works in two stages.
<table> <thead> <tr> <th>Memory stage</th> <th>When it fires</th> <th>Purpose</th> <th>Latent memory traits</th> </tr> </thead> <tbody> <tr> <td>Prompt stage (augment_prompt)</td>
<td>At the start of inference, after the original prompt is processed</td>
<td>Initialize global memory and set up later reasoning</td>
<td>Globally relevant; stable</td>
</tr>
<tr>
<td>Inference stage (augment_inference)</td>
<td>During reasoning, when the Trigger fires</td>
<td>Add real-time memory for the current step</td>
<td>Tightly coupled to current context; real-time</td>
</tr>
</tbody>
</table>
The Weaver uses two learnable query vectors: prompt_query_latents for the prompt stage and inference_query_latents for the inference stage. When latent memory is fused with the sequence, the Weaver concatenates latent memory with the original input embeddings, updates the attention mask and position IDs, and keeps temporal order consistent. PEFT, including LoRA, can be used to improve latent memory generation for a specific task.
Inside the source: core logic and optimizations
The implementation is organized around the LatentMemoryModel class. Training uses forward; inference uses generate. The core class wires the Reasoner, Weaver, Trigger, and projection layers together.
@registry.register_model("latmem")
class LatentMemoryModel(BaseModel):
def __init__(
self,
reasoner_model_name: str, # Reasoner model name (e.g., "meta-llama/Llama-3-8B-Instruct")
weaver_model_name: str, # Weaver model name
prompt_latents_len: int, # Number of latent memories in prompt stage
inference_latents_len: int,# Number of latent memories in inference stage
weaver_peft_config: Optional[PeftConfig] = None, # Weaver PEFT config
trigger_model_name: str = None, # Trigger model name (optional)
trigger_peft_config: Optional[PeftConfig] = None, # Trigger PEFT config (optional)
max_prompt_aug_num: int = 1, # Max augmentation count in prompt stage
max_inference_aug_num: int = 5, # Max augmentation count in inference stage
):
# 1. Initialize reasoner (frozen parameters)
self.model = AutoModelForCausalLM.from_pretrained(
reasoner_model_name, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2"
)
fix_model_parameters(self.model) # Freeze reasoner parameters
# 2. Initialize weaver (supports PEFT fine-tuning)
self.weaver = MemGenWeaver(weaver_model_name, prompt_latents_len, inference_latents_len, weaver_peft_config)
# 3. Initialize trigger (NanoTrigger by default; MemGenTrigger optional)
self.trigger = NanoTrigger()
if trigger_model_name is not None:
self.trigger = MemGenTrigger(trigger_model_name, trigger_peft_config)
# 4. Projection layers (dimension alignment)
self.reasoner_to_weaver = nn.Linear(self.model.config.hidden_size, self.weaver.config.hidden_size, dtype=torch.bfloat16)
self.weaver_to_reasoner = nn.Linear(self.weaver.config.hidden_size, self.model.config.hidden_size, dtype=torch.bfloat16)
# Other configs (delimiters, precision, etc.)
self.delimiters = [",", ".", "\n"] # Delimiters for selecting augmentation points
self.model = self.model.bfloat16() # Unify precision
Choosing augmentation points
MemGen uses delimiters such as commas, periods, and newlines as semantic breakpoints. These are places where a human-like pause occurs during reasoning, and inserting memory there can influence what comes next. The helper below finds those positions and limits how many augmentations are applied.
def _select_augment_points_after_delimiter(self, input_ids, labels, delimiters, tokenizer, max_augment_num):
# 1. Convert delimiters to token IDs
delimiter_ids = [tokenizer.encode(d, add_special_tokens=False)[0] for d in delimiters]
# 2. Scan the input sequence for delimiter positions
B, seq_len = input_ids.shape
augment_indices = []
for b in range(B):
indices = torch.where(torch.isin(input_ids[b], torch.tensor(delimiter_ids, device=input_ids.device)))[0]
# 3. Limit augmentation count to avoid over-augmentation
indices = indices[:max_augment_num]
augment_indices.append(indices)
return augment_indices
Training with filtered loss
During training, loss is computed only on original input positions. Positions that correspond to latent memory are ignored, so the training objective stays focused on the task.
def _forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor, **kwargs):
# 1. Select augmentation points (delimiter positions)
augmentation_indices = self._select_augment_points_after_delimiter(...)
# 2. Input embeddings and sequence segmentation
inputs_embeds = self.model.get_input_embeddings()(input_ids)
current_inputs_embeds = torch.empty(B, 0, hidden_size).to(device)
# 3. Walk through augmentation points and insert latent memory
for aug_idx in augmentation_indices:
# Slice the original sequence segment
segment_embeds = inputs_embeds[:, current_start:aug_idx]
# Weaver generates latent memory
weaver_inputs = self.reasoner_to_weaver(current_inputs_embeds)
weaver_hidden = self.weaver.augment_inference(weaver_inputs, ...)
latent_embeds = self.weaver_to_reasoner(weaver_hidden)
# Concatenate original segment and latent memory
current_inputs_embeds = torch.cat([current_inputs_embeds, segment_embeds, latent_embeds], dim=1)
# 4. Reasoner forward pass
outputs = self.model(inputs_embeds=current_inputs_embeds, attention_mask=current_attention_mask)
# 5. Filter latent memory positions; compute loss only on original input
valid_logits = outputs.logits[:, :input_ids.shape[1], :] # Ignore memory positions
valid_labels = labels[:, :input_ids.shape[1]]
loss = F.cross_entropy(valid_logits.reshape(-1, valid_logits.size(-1)), valid_labels.reshape(-1))
return loss
Inference with dynamic augmentation
At generation time, the model creates one token at a time. After each step, the Trigger decides whether augmentation should happen. Sequences that are not augmented are aligned by left padding.
@torch.no_grad()
def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, generation_config, **kwargs):
# 1. Prompt-stage memory augmentation (initialize global memory)
weaver_inputs = self.reasoner_to_weaver(inputs_embeds)
weaver_hidden = self.weaver.augment_prompt(weaver_inputs, ...)
latent_embeds = self.weaver_to_reasoner(weaver_hidden)
current_inputs_embeds = torch.cat([inputs_embeds, latent_embeds], dim=1)
# 2. Iteratively generate new tokens
for i in range(max_new_tokens):
# Reasoner generates one token
outputs = self.model(inputs_embeds=current_inputs_embeds, ...)
current_inputs_embeds = self._append_one_step(outputs, ...)
# 3. Trigger decides whether to augment
augment_decision = self._should_augment(current_inputs_embeds, ...)
augment_indices = torch.where(augment_decision == 1)[0]
# 4. Insert latent memory for sequences that need augmentation
if len(augment_indices) > 0:
# Weaver generates inference-stage memory
weaver_hidden = self.weaver.augment_inference(...)
latent_embeds = self.weaver_to_reasoner(weaver_hidden)
# Concatenate and align dimensions (left padding for non-augmented sequences)
current_inputs_embeds = self._merge_augmented_sequences(...)
return current_input_ids
Several optimizations matter in practice. bfloat16 halves memory use versus float32 while keeping the precision needed for large-model inference. Flash Attention 2 reduces the time and space complexity of attention and supports longer sequences. Freezing the Reasoner means the trainable set is limited to the Weaver, Trigger, and projection layers, cutting trainable parameters by more than 70%. The code also handles tokenizers without a pad token and standardizes dialogue templates so different pretrained models can be used.
Deployment and tuning
Core dependencies and a Flash Attention 2 build:
# Core dependencies
pip install torch transformers peft accelerate sentencepiece
pip install flash-attn==2.5.8 # Supports Flash Attention 2
Useful starting parameters:
<table> <thead> <tr> <th>Parameter</th> <th>Role</th> <th>Suggested value</th> </tr> </thead> <tbody> <tr> <td>reasoner_model_name</td>
<td>Reasoner model</td>
<td>meta-llama/Llama-3-8B-Instruct for a lighter setup, Qwen-72B-Instruct for higher performance</td>
</tr>
<tr>
<td>weaver_model_name</td>
<td>Weaver model</td>
<td>Same architecture as the Reasoner to keep embedding spaces compatible</td>
</tr>
<tr>
<td>prompt_latents_len</td>
<td>Number of prompt-stage latent memories</td>
<td>1–2 to avoid overloading the initial memory</td>
</tr>
<tr>
<td>inference_latents_len</td>
<td>Number of inference-stage latent memories</td>
<td>3–5 to balance richness and efficiency</td>
</tr>
<tr>
<td>max_inference_aug_num</td>
<td>Maximum inference-stage augmentations</td>
<td>3–5 to avoid over-augmentation and excessive sequence length</td>
</tr>
</tbody>
</table>
For fine-tuning, use data formatted as task description, multi-step reasoning, and result. Delimiter positions can be marked but are optional. Apply LoRA to the Weaver and Trigger with r=8 and lora_alpha=32 while freezing the rest. Start with a learning rate around 2e-4; the Trigger can use a slightly higher rate such as 2.5e-4.
If dimensions mismatch, check the projection layers and keep bfloat16 consistent. If generation is slow, lower inference_latents_len and max_inference_aug_num, and avoid enabling Flash Attention 2 on modules that do not need it. If memory quality is weak, increase the number of latent memories, extend the multi-step reasoning length in fine-tuning data, and improve the Trigger fine-tuning data.
Where MemGen fits
The approach suits tasks that need long-horizon reasoning and dynamic memory. In complex planning—route planning for a road trip or project management—it can integrate memory across multiple reasoning steps. In multi-turn dialogue, it can keep key information alive and generate context-relevant memory. In code generation and debugging, it can remember code logic and repair experience during reasoning. In scientific or mathematical reasoning, it can store intermediate results step by step and support later steps.
A comparison across evaluation dimensions shows the trade-offs:
<table> <thead> <tr> <th>Dimension</th> <th>Parametric memory</th> <th>Retrieval-based memory</th> <th>MemGen (generative latent memory)</th> </tr> </thead> <tbody> <tr> <td>Multi-step reasoning accuracy</td> <td>65% (catastrophic forgetting)</td> <td>78% (static retrieval disconnect)</td> <td>89% (dynamic memory coupling)</td> </tr> <tr> <td>Training cost</td> <td>High (full fine-tuning)</td> <td>Low (no training)</td> <td>Medium (partial module fine-tuning)</td> </tr> <tr> <td>Inference speed</td> <td>Fast (no extra overhead)</td> <td>Slow (retrieval overhead)</td> <td>Medium (memory generation overhead is controllable)</td> </tr> <tr> <td>Long-sequence fit</td> <td>Poor (context window limit)</td> <td>Medium (retrieval truncation)</td> <td>Good (dynamic insertion avoids long-sequence redundancy)</td> </tr> </tbody> </table>What comes next
Four directions stand out. Multimodal memory could extend the Weaver to generate latent memory from images and audio. Adaptive triggering could use reinforcement learning to adjust trigger decisions by task type. Lighter deployment could compress the Weaver and Trigger for edge devices and low-resource environments. Memory distillation could compress long-sequence dynamic memory into a compact representation and further improve inference efficiency.
The broader shift is conceptual. MemGen treats memory as something that grows with reasoning rather than something stored and fetched. That changes the design question from what to store to how to generate useful memory during inference. For agents that need long-horizon reasoning, multi-turn context, or domain-specific inference, that shift is the part worth paying attention to.