【Agent】生成式隐式记忆 MemGen 源码解读

【Agent】生成式隐式记忆 MemGen 源码解读

目录

• 【Agent】生成式隐式记忆 MemGen 源码解读
  • 0x00 概要
  • 0x01 背景
  • 0x02 源码解析
    • 2.1 模型
      • 2.1.1 核心特色
      • 2.1.2 网络结构
      • 2.1.3 代码
      • 2.1.4 插入阶段
        • forward
        • generate
          • 核心作用
          • 核心特色
          • 推理生成流程图
    • 2.2 Trigger
      • 2.2.1. 核心作用
      • 2.2.2. 核心特色
      • 2.2.3 网络架构
      • 2.2.4 代码
    • 2.3 MemGenWeaver
      • 2.3.1 核心作用
      • 2.3.2 核心特色
      • 2.3.3 网络架构
      • 2.3.4 代码
  • 0xFF 参考

0x00 概要

MemGen旨在构建一个动态、生成式的记忆框架，其核心由两个协同工作的轻量级模块构成：一个基于强化学习（RL）训练的记忆触发器（Memory Trigger）和一个记忆编织器（Memory Weaver）。

论文：MemGen: Weaving Generative Latent Memory for Self-Evolving Agents
链接：https://arxiv.org/abs/2509.24704
代码：https://github.com/KANABOON1/MemGen

0x01 背景

MemGen 提出动态生成式记忆框架，由记忆触发器与记忆编织器两个轻量模块协同构成，旨在突破现有智能体记忆范式的局限。

当前主流的记忆实现路径为：

• 参数化记忆通过微调将经验编码进模型参数，虽能深度内化知识却易引发灾难性遗忘；
• 基于检索的记忆将经验外化存储，虽规避了遗忘问题，但静态的一次性检索机制无法体现记忆与推理动态交互的认知特性。

这一现状引出两大核心问题：如何实现记忆与推理在每一步思考中的无缝耦合，以及如何让记忆从提取式升级为满足当前需求的生成式重构，而动态生成式隐式记忆正是应对这些挑战的第三种探索路径。

0x02 源码解析

MemGen项目旨在创建一个动态且自生成的记忆框架，该框架由两个协同工作的轻量级模块组成：一个基于强化学习训练的记忆触发器和一个记忆编织器。这一框架的核心思想是解决大型语言模型（LLM）智能体能力涌现时对“自进化”机制的探索需求，其中记忆扮演关键角色。

2.1 模型

LatentMemoryModel 是 MemGen 框架的核心实现，旨在构建动态生成式隐式记忆系统，解决传统记忆范式的局限性。通过整合推理器（Reasoner）、记忆编织器（Weaver）和记忆触发器（Trigger），实现记忆与推理过程的无缝耦合，让智能体在任务执行中动态生成、使用记忆，而非依赖静态检索或参数化存储。

2.1.1 核心特色

模型的核心特色如下：

• 模块化协同设计：由推理器（核心推理）、编织器（生成潜在记忆）、触发器（控制记忆触发）三大模块构成，模块间通过投影层实现嵌入空间映射，结构清晰且解耦。
• 动态记忆增强：在推理过程中自动识别分隔符位置作为记忆增强点，动态插入编织器生成的潜在记忆，突破静态记忆注入的局限，贴合人类认知中记忆与推理的动态交互特性。
• 精度与效率优化：默认使用 bfloat16 精度，推理器采用 Flash Attention 2 提升计算效率；冻结推理器参数，仅训练编织器和触发器，实现参数高效学习。
• 灵活配置与兼容性：支持自定义触发器模型、PEFT 微调配置、记忆增强次数等参数；自动处理 Tokenizer 缺失 pad token 的问题，标准化对话模板，提升跨场景兼容性。
• 损失计算精准过滤：通过潜在记忆掩码排除记忆嵌入对应的位置，仅对原始输入位置计算损失，确保训练目标聚焦于核心任务性能，避免记忆生成过程干扰主任务学习。

2.1.2 网络结构

关键说明（核心设计亮点）

1. 三大模块协同逻辑：
   • 推理器（Reasoner）：核心推理组件，权重冻结以保留基础能力，仅通过潜在记忆调整解码路径。
   • 触发器（MemGenTrigger）：动态判断记忆插入时机，输出二分类触发概率，决定是否调用编织器。
   • 编织器（MemGenWeaver）：生成针对性潜在记忆，分提示词 / 推理两阶段设计，支持 PEFT 高效微调。
2. 核心流程闭环：输入 → 推理器生成原始嵌入 → 触发器 + 增强点选择模块确定插入位置 → 编织器生成潜在记忆 → 投影层适配维度 → 重组增强序列 → 推理器完成最终推理 → 过滤无效位置输出。
3. 关键技术细节：
   • 跨模块投影：通过 reasoner_to_weaver 和 weaver_to_reasoner 解决推理器与编织器嵌入维度不匹配问题。
   • 动态记忆增强：按分隔符拆分序列，逐段插入记忆，避免长序列冗余，贴合人类 “思考 - 记忆” 交互模式。
   • 精度与效率：全流程采用 bfloat16 精度，推理器 / 编织器启用 Flash Attention 2，平衡性能与速度。
4. 训练与推理适配：
   • 训练时：通过 labels 和 valid_logits 计算损失，仅优化编织器、触发器及投影层参数。
   • 推理时：无需 labels，自动完成 “触发判断 - 记忆生成 - 推理增强” 全流程，实现动态自进化。

具体网络结构如下

2.1.3 代码

LatentMemoryModel 的代码如下：

@registry.register_model("latmem")
class LatentMemoryModel(BaseModel):  # 定义了一个名为 LatentMemoryModel 的类，继承自 BaseModeldef __init__(self, reasoner_model_name: str,  # 推理模型名称weaver_model_name: str,  # 记忆编织器模型名称prompt_latents_len: int,  # 提示长度inference_latents_len: int,  # 推理长度weaver_peft_config: Optional[PeftConfig] = None,  # 记忆编织器配置，可选trigger_model_name: str = None,  # 触发模型名称，可选trigger_peft_config: Optional[PeftConfig] = None,  # 触发器配置，可选max_prompt_aug_num: int = 1,  # 最大提示增强数量max_inference_aug_num: int = 5,  # 最大推理增强数量):   super().__init__()  # 调用父类构造函数# 构建推理模型self.model = AutoModelForCausalLM.from_pretrained(  # 从预训练模型加载推理模型reasoner_model_name, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")self.tokenizer = AutoTokenizer.from_pretrained(reasoner_model_name)  # 加载入分词器self.config = self.model.config  # 获取模型配置# 构建记忆编织器self.weaver = MemGenWeaver(  # 初始化记忆编织器weaver_model_name, prompt_latents_len, inference_latents_len, weaver_peft_config)# 构建触发器self.trigger = NanoTrigger()  # 默认触发器，始终返回 trueif trigger_model_name is not None:self.trigger = MemGenTrigger(  # 如果指定了触发模型，则加载相应的触发器trigger_model_name, trigger_peft_config)logging.info(f"Use Trigger: {trigger_model_name}")  # 记录日志# 投影层，用于在推理模型和记忆编织器之间映射嵌入# 将推理模型输入嵌入映射到记忆编织器输入嵌入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)self.delimiters: List[str] = [",", ".", "\n"]  # 用于检测增强点的分隔符self.max_prompt_aug_num = max_prompt_aug_num  # 提示后提示中插入潜在数量self.max_inference_aug_num = max_inference_aug_num  # 指定分隔符后插入潜在数量# 后处理self._postprocess_models()  # 后处理模型self.warnings_issued = {}  # 存储发出的警告self.model_tags = None  # 存储模型标签log_trainable_params(self)  # 记录可训练参数def add_model_tags(self, tags: Union[list[str], str]) -> None:  # 添加模型标签r"""向模型添加自定义标签，这些标签将被推送到 Hugging Face Hub。不会覆盖模型中现有的标签。参数：tags (`Union[list[str], str]`)：要添加到模型的标签例子：```pythonfrom transformers import AutoModelmodel = AutoModel.from_pretrained("google-bert/bert-base-cased")model.add_model_tags(["custom", "custom-bert"])# 将模型推送到您的命名空间，名称为 "my-custom-bert"。model.push_to_hub("my-custom-bert")"""if isinstance(tags, str):tags = [tags]if self.model_tags is None:self.model_tags = []for tag in tags:if tag not in self.model_tags:self.model_tags.append(tag)def _postprocess_models(self):"""后处理记忆模型的组件：推理模型、记忆编织器、触发器和分词器。步骤：1. 冻结推理模型的所有参数（不更新梯度）。2. 将所有模型转换为 bfloat16 以提高内存和计算效率。3. 确保分词器有一个有效的填充符：- 如果缺少填充符，使用 EOS 符作为填充符。- 设置 `padding_side` 为 "left" 以兼容生成任务。4. 标准化分词器的模板为 `CONVERSATION_TEMPLATE`。"""# 默认冻结推理模型的所有参数fix_model_parameters(self.model)# 将所有子模型转换为 bfloat16self.model = self.model.bfloat16()self.weaver = self.weaver.bfloat16()self.trigger = self.trigger.bfloat16()# 确保分词器有一个填充符if self.tokenizer.pad_token is None:self.tokenizer.pad_token = self.tokenizer.eos_tokenself.tokenizer.pad_token_id = self.tokenizer.eos_token_idself.tokenizer.padding_side = "left"logging.info(f"Tokenizer has no pad token. Using EOS token ({self.tokenizer.eos_token}) as pad token.")# 标准化分词器的模板self.tokenizer.chat_template = CONVERSATION_TEMPLATE

2.1.4 插入阶段

LatentMemoryModel 的两个关键函数 forward 和 generate 区别如下：

• forward 函数
  • 训练时候计算损失，由训练循环自动调用。
• generate 函数
  • 推理时候生成文本，由代码显式调用。

forward

forward 函数的主体如下：

    def _forward(self, input_ids: torch.Tensor,attention_mask: torch.Tensor,labels: torch.Tensor,   **kwargs) -> torch.Tensor:# 预处理输入assert input_ids.shape == attention_mask.shape == labels.shapetokenizer = self.tokenizerreasoner = self.modelweaver = self.weaverdelimiters = self.delimitersmax_augment_num = self.max_inference_aug_num  # 限制推理增强点的数量以避免过度增强device = self.deviceembeds_dtype = reasoner.get_input_embeddings().weight.dtypeB, _ = input_ids.shapehidden_size = reasoner.config.hidden_size# 选择增强索引augmentation_indices = self._select_augment_points_after_delimiter(input_ids, labels, delimiters, tokenizer, max_augment_num)# 输入嵌入inputs_embeds = reasoner.get_input_embeddings()(input_ids)# 初始化开始索引和空张量以累积处理的段current_start_idx = 0current_inputs_embeds = torch.empty(B, 0, hidden_size).to(device, dtype=embeds_dtype)current_attention_mask = torch.empty(B, 0).to(device, dtype=attention_mask.dtype)current_latents_mask = torch.empty(B, 0).to(device, dtype=torch.bool)# 遍历所选增强点for aug_idx in augmentation_indices:# 切片原始嵌入和注意力掩码segment_inputs_embeds = inputs_embeds[:, current_start:aug_idx]segment_attention_mask = attention_mask[:, current_start:aug_idx]segment_latents_mask = torch.zeros(B, segment_inputs_embeds.size(1).to(device, dtype=torch.bool)# 连接当前段到累积嵌入和掩码current_inputs_embeds = torch.cat([current_inputs_embeds, segment_inputs_embeds], dim=1)current_mask = torch.cat([current_mask, segment_attention_mask], dim=1)current_position_ids = generate_position_ids(current_mask)current_latents = torch.cat([current_latents, segment_latents], dim=1)# 将推理模型嵌入映射到记忆编织器嵌入weaver_inputs_embeds = self.reasoner_to_weaver(current_inputs_embeds)# 确定此点是否为提示（增强）的结束is_prompt_end_aug = (labels[:, aug_idx] != -100).all() and (labels[:, aug_idx-1] == -100).all().item()# 根据类型，使用记忆编织器增强提示或推理if is_prompt_end_aug:weaver_hidden_states, attn_mask, pos_ids = weaver.augment_prompt(weaver_inputs, current_attention_mask, current_position_ids)else:weaver_hidden_states, attn_mask, pos_ids = weaver.augment_inference(weaver_inputs, current_attention_mask, current_position_ids) # 将记忆编织器隐藏状态映射回推理模型嵌入latent_inputs_embeds = self.weaver_to_reasoner(weaver_hidden_states)# 更新累积嵌入和掩码与新增强段current_inputs_embeds = torch.cat

generate

核心作用

该 generate 方法是 MemGen 模型的推理核心，实现了动态记忆增强与序列生成的无缝融合。通过迭代生成新 token，每步自适应判断是否插入编织器生成的潜在记忆，让推理器在生成过程中实时利用动态记忆调整解码路径，最终输出增强后的序列（可选返回记忆增强位置掩码）。

核心特色

• 双阶段记忆增强：先执行提示词阶段记忆增强（初始化全局记忆），再在迭代生成中动态触发推理阶段增强（补充实时记忆），适配不同生成阶段的记忆需求。
• 自适应触发机制：通过 _should_augment 结合触发器决策，仅对需要记忆支持的序列执行增强，避免无意义的计算开销。
• 维度对齐优化：非增强序列采用左填充（_left_pad）方式对齐增强序列维度，确保批次内所有序列格式统一，不影响批量生成效率。
• 高效推理设计：
  • 禁用梯度计算（@torch.no_grad()），节省内存并加速推理；
  • 启用推理器缓存（use_cache=True），减少重复计算；
  • 仅在必要时输出隐藏状态，降低计算成本。
• 灵活配置与可解释性：支持控制最大生成 token 数、采样策略等参数；可选返回 augmentation_pos 掩码，标记记忆插入位置，提升模型可解释性。
• 鲁棒性保障：提前终止机制（所有序列生成 EOS 或达最大增强次数时终止），避免无效迭代；重构生成配置固定关键参数，确保生成稳定性。

推理生成流程图

潜在记忆插入的完整流程：

• 初始化阶段：对输入提示进行增强，插入初始潜在记忆。
• 生成循环：逐个生成token。
• 条件检查：在每个步骤检查是否满足插入条件。
• 决策判断：使用trigger模型决定是否插入潜在记忆。
• 潜在记忆生成：通过weaver模型生成潜在记忆表示。
• 嵌入连接：将潜在记忆嵌入连接到当前输入序列。
• 继续生成：使用增强后的序列继续生成下一个token。

具体流程如下图所示：

代码如下：

@torch.no_grad()  # 禁用梯度计算，适用于推理阶段，提升效率并节省内存
def generate(self, input_ids: torch.Tensor,  # 输入token ID序列，形状[batch_size, prompt_len]attention_mask: torch.Tensor,  # 注意力掩码，形状与input_ids一致generation_config: GenerationConfig = None,  # 生成配置（如最大新token数、采样策略等）return_augmentation_mask: bool = False,  # 是否返回记忆增强位置掩码**kwargs
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """执行MemGen模型的推理生成流程：动态融合潜在记忆与推理器，生成增强后的输出序列。核心逻辑：1. 初始化提示词阶段的记忆增强2. 迭代生成新token，每步判断是否触发推理阶段记忆增强3. 对需增强的序列插入编织器生成的潜在记忆，非增强序列左填充对齐维度4. 生成完成后返回结果（可选返回增强位置掩码）"""tokenizer = self.tokenizerreasoner = self.modelweaver = self.weavertrigger = self.triggerdelimiters = self.delimitersmax_augment_num = self.max_inference_aug_num  # 单序列最大推理阶段增强次数invalid_token_id = -100  # 无效位置标记（用于增强位置掩码）# 预处理输入：转移到模型所在设备input_ids = input_ids.to(self.device)attention_mask = attention_mask.to(self.device)# 提取生成配置关键参数max_new_tokens = generation_config.max_new_tokens  # 最大生成新token数do_sample = generation_config.do_sample  # 是否启用采样生成temperature = generation_config.temperature  # 采样温度（控制随机性）pad_token_id = tokenizer.pad_token_id  # pad token IDeos_token_id = tokenizer.eos_token_id  # 结束token IDprompt_len = input_ids.size(1)  # 提示词长度# 重构生成配置（固定必要参数，确保生成稳定性）generation_config = GenerationConfig(do_sample=do_sample,temperature=temperature,pad_token_id=pad_token_id,eos_token_id=eos_token_id,use_cache=True  # 启用缓存加速生成)# 将输入token ID转换为嵌入向量inputs_embeds = reasoner.get_input_embeddings()(input_ids)B, _, hidden_size = inputs_embeds.shape  # B=batch_size，hidden_size=推理器隐藏层维度device = inputs_embeds.device  # 模型所在设备（CPU/GPU）# 初始化生成过程中的关键张量current_inputs_embeds = inputs_embeds  # 当前输入嵌入（含原始提示词+潜在记忆）current_attention_mask = attention_mask  # 当前注意力掩码current_position_ids = generate_position_ids(current_attention_mask)  # 当前位置IDcurrent_input_ids = input_ids  # 当前已生成的token ID序列# 提示词阶段记忆增强：生成并插入提示词专用潜在记忆weaver_inputs_embeds = self.reasoner_to_weaver(current_inputs_embeds)  # 映射到编织器嵌入空间weaver_hidden_states, attn_mask, pos_ids = weaver.augment_prompt(weaver_inputs_embeds, current_attention_mask, current_position_ids)latent_inputs_embeds = self.weaver_to_reasoner(weaver_hidden_states)  # 映射回推理器嵌入空间# 拼接提示词与增强记忆current_inputs_embeds = torch.cat([current_inputs_embeds, latent_inputs_embeds], dim=1)current_attention_mask = torch.cat([current_attention_mask, attn_mask], dim=1)current_position_ids = torch.cat([current_position_ids, pos_ids], dim=1)# 生成循环初始化sentence_augment_count = torch.zeros(B, dtype=torch.int, device=device)  # 各序列已增强次数augmentation_pos = torch.full((B, max_new_tokens), fill_value=invalid_token_id, device=device)  # 增强位置掩码inserted_embeds: List[List[torch.Tensor]] = [[] for _ in range(B)]  # 记录插入的潜在记忆（用于后处理）for i in range(max_new_tokens):# 若所有序列均已生成EOS token，提前终止if (current_input_ids[:, -1] == eos_token_id).all():break   # 若所有序列均已达到最大增强次数，一次性生成剩余tokenif (sentence_augment_count >= max_augment_num).all():# 调整剩余生成长度generation_config.max_new_tokens = max_new_tokens - i# 推理器生成剩余tokengenerated = reasoner.generate(inputs_embeds=current_inputs_embeds,attention_mask=current_attention_mask,generation_config=generation_config,)current_input_ids = torch.cat([current_input_ids, generated], dim=1)break# 推理器前向传播，获取当前步输出outputs = reasoner(inputs_embeds=current_inputs_embeds,attention_mask=current_attention_mask,position_ids=current_position_ids,output_hidden_states=False,  # 推理阶段无需输出隐藏状态，提升效率)# 生成并追加一个新token，更新关键张量current_inputs_embeds, current_attention_mask, current_position_ids, current_input_ids = self._append_one_step(outputs, current_inputs_embeds, current_attention_mask, current_position_ids, current_input_ids, do_sample, temperature)# 若为最后一步生成，终止循环if i == max_new_tokens - 1:  break # 判断当前批次中哪些序列需要进行推理阶段记忆增强augment_decision = self._should_augment(current_input_ids, current_attention_mask, sentence_augment_count=sentence_augment_count, do_sample=do_sample, temperature=temperature  )augmentation_pos[:, i + 1] = augment_decision  # 记录增强位置（1=增强，0=不增强，-100=无效）augment_indices = torch.where(augment_decision == 1)[0]  # 需增强的序列索引# 对需增强的序列执行记忆增强，非增强序列左填充对齐维度if len(augment_indices) > 0:# 递增需增强序列的增强次数计数sentence_augment_count[augment_indices] += 1# 提取需增强序列的嵌入、掩码和位置IDcandidate_inputs_embeds = current_inputs_embeds[augment_indices]candidate_attention_mask = current_attention_mask[augment_indices]candidate_position_ids = current_position_ids[augment_indices]# 编织器生成推理阶段潜在记忆weaver_inputs_embeds = self.reasoner_to_weaver(candidate_inputs_embeds)weaver_hidden_states, attn_mask, _ = weaver.augment_inference(weaver_inputs_embeds, candidate_attention_mask, candidate_position_ids)latent_inputs_embeds = self.weaver_to_reasoner(weaver_hidden_states)  # 映射回推理器空间# 拼接原始嵌入与潜在记忆candidate_inputs_embeds = torch.cat([candidate_inputs_embeds, latent_inputs_embeds], dim=1)candidate_attention_mask = torch.cat([candidate_attention_mask, attn_mask], dim=1)# 构建合并张量（适配所有序列，包括增强和非增强）new_len = candidate_inputs_embeds.size(1)  # 增强后序列长度merged_inputs_embeds = torch.zeros((B, new_len, hidden_size), device=device, dtype=current_inputs_embeds.dtype)merged_attention_mask = torch.zeros((B, new_len), device=device, dtype=current_attention_mask.dtype)# 填充增强序列merged_inputs_embeds[augment_indices] = candidate_inputs_embedsmerged_attention_mask[augment_indices] = candidate_attention_mask# 填充非增强序列（左填充对齐长度）non_augment_indices = torch.where(augment_decision != 1)[0]if len(non_augment_indices) > 0:non_aug_inputs_embeds = current_inputs_embeds[non_augment_indices]non_aug_attention_mask = current_attention_mask[non_augment_indices]non_aug_inputs_embeds, non_aug_attention_mask, _ = self._left_pad(non_aug_inputs_embeds, non_aug_attention_mask, None, weaver.inference_latents_num)merged_inputs_embeds[non_augment_indices] = non_aug_inputs_embedsmerged_attention_mask[non_augment_indices] = non_aug_attention_mask# 更新当前关键张量current_inputs_embeds = merged_inputs_embedscurrent_attention_mask = merged_attention_maskcurrent_position_ids = generate_position_ids(current_attention_mask)  # 重新生成位置ID# 记录插入的潜在记忆（用于后处理或可解释性分析）for idx, embed in zip(augment_indices, latent_inputs_embeds):inserted_embeds[idx].append(embed.clone().detach().cpu())# 后处理：调整增强位置掩码长度与生成结果一致new_generated_len = current_input_ids.size(1) - prompt_lenaugmentation_pos = augmentation_pos[:, :new_generated_len]# 根据配置返回结果：仅生成序列 或 序列+增强位置掩码if not return_augmentation_mask:return current_input_idselse:return current_input_ids, augmentation_pos

2.2 Trigger

2.2.1. 核心作用

该模块定义了 MemGen 框架中记忆触发器的核心接口与两种具体实现，核心作用是动态决策记忆增强的时机—— 即在推理过程中判断何时插入编织器生成的潜在记忆，实现记忆与推理的动态耦合，突破传统静态记忆注入的局限。

2.2.2. 核心特色

• 抽象接口统一规范：Trigger抽象基类定义了触发器的核心接口，确保后续扩展新触发器时遵循统一标准，提升代码可扩展性。
• 双实现适配不同场景：
  • NanoTrigger：极简实现，始终触发记忆增强，无需训练，适用于快速测试、基线对比或无需动态控制的简单场景。
  • MemGenTrigger：基于预训练 LLM 的智能触发器，通过二分类头适配决策任务，支持 PEFT 参数高效微调，能根据输入序列动态判断是否触发，适配复杂真实场景。
• 高效适配与灵活扩展：
  • 采用 bfloat16 精度和 Flash Attention 2 优化计算效率；
  • 支持 PEFT 微调，在不冻结基础模型的前提下实现参数高效学习；
  • 替换 LLM 原始输出头为二分类头，精准适配 "是否插入记忆" 的决策需求。
• 模块解耦设计：触发器决策独立于编织器模块，仅基于输入序列和数据分布做出判断，保证了模块间的低耦合和高内聚。

2.2.3 网络架构

网络架构图如下。

说明如下：

1. 模型支持PEFT参数高效微调（如LoRA），适配于Transformer Blocks层
2. 整体精度采用bfloat16，平衡计算效率与数值稳定性
3. 注意力计算通过Flash Attention 2优化，提升长序列处理速度

2.2.4 代码

class Trigger(torch.nn.Module, ABC):"""记忆触发器的抽象基类（Trigger）。定义了触发器的核心接口，用于决定在推理过程中何时触发记忆增强（插入潜在记忆）。所有具体触发器实现都需继承此类并实现forward方法。"""def __init__(self):super().__init__()  # 调用父类Module的初始化方法@abstractmethoddef forward(self, **kwargs) -> bool:"""抽象前向传播方法：接收输入数据，返回是否触发记忆增强的决策。子类必须实现此方法，定义具体的触发逻辑。Args:**kwargs: 可变关键字参数，包含输入序列、注意力掩码等模型所需数据Returns:bool: 触发决策（True表示触发记忆增强，False表示不触发）"""...class NanoTrigger(torch.nn.Module):"""极简触发器（NanoTrigger）：始终触发记忆增强的基础实现。无需复杂逻辑，固定返回触发决策，适用于基础测试或无需动态控制的场景。"""def __init__(self):super().__init__()  # 注册一个缓冲区张量，用于获取模型所在设备（无实际计算意义）self.register_buffer("_device", torch.tensor(0.0))@propertydef device(self):"""获取模型所在设备（CPU/GPU）"""return self._device.devicedef forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> bool:# 该"极简触发器"始终预测需要插入记忆# 输出logits张量，其中插入决策（索引=1）的概率被设为1.0# 适用于批次中的每个token位置batch_size, seq_len = input_ids.shape# 初始化logits张量：形状为[batch_size, seq_len, 2]，2表示"不插入"（0）和"插入"（1）两类logits = torch.zeros(batch_size, seq_len, 2, device=input_ids.device)logits[..., 1] = 1.0  # 将所有位置的"插入"决策概率设为1.0return logitsclass MemGenTrigger(torch.nn.Module):"""MemGen框架的专用触发器模块（MemGenTrigger）。- 输入：接收推理器模型当前解码序列的`inputs_embeds`（或input_ids）- 输出：生成形状为[batch_size, seq_len, 2]的logits张量，表示每个位置"不插入"（0）和"插入"（1）记忆的概率，用于动态决策记忆增强时机。"""def __init__(self, pretrained_model_name_or_path: str,  # 预训练模型名称或路径（用于初始化触发器LLM）peft_config: Optional[PeftConfig] = None  # PEFT配置（可选，用于参数高效微调）):super().__init__()# 构建基础LLM模型（作为触发器的核心推理组件）self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, torch_dtype=torch.bfloat16,  # 使用bfloat16精度提升效率attn_implementation="flash_attention_2"  # 启用Flash Attention 2优化注意力计算)self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)  # 对应的Tokenizer# 对基础模型进行后处理（设置可训练、替换输出头）self.model = self._postprocess(self.model)# 若提供PEFT配置，应用参数高效微调if peft_config is not None:self.model = get_peft_model(self.model, peft_config)self.config = self.model.config  # 保存模型配置@propertydef device(self):"""获取模型所在设备（CPU/GPU）"""return self.model.devicedef _postprocess(self, model: PreTrainedModel):"""对基础模型进行后处理，适配触发器的二分类任务需求。Args:model: 原始预训练LLM模型Returns:处理后的模型（可训练、替换为二分类输出头）"""# 设置所有模型参数为可训练for parameter in model.parameters():parameter.requires_grad = True# 将原始语言模型的输出头（lm_head）替换为二分类头hidden_size = model.config.hidden_size  # 模型隐藏层维度classification_head = nn.Linear(hidden_size, 2)  # 输出维度为2（不插入/插入）model.lm_head = classification_head# 确保新的二分类头参数可训练for param in model.lm_head.parameters():param.requires_grad = Truereturn modeldef forward(self, input_ids: Optional[torch.LongTensor] = None,  # 生成序列的token ID，形状[batch_size, seq_len]attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码，避免关注填充token**kwargs: Unpack[TransformersKwargs],  # 传递给底层模型的额外参数) -> torch.Tensor:"""序列生成的触发决策机制。触发器基于已生成的`input_ids`做出决策，受数据分布影响，但独立于编织器模块。Args:input_ids (Optional[torch.LongTensor]): 生成序列的token ID张量attention_mask (Optional[torch.Tensor]): 注意力掩码，默认None**kwargs: 传递给底层模型的额外关键字参数Returns:torch.Tensor: Logits张量，形状为`(batch_size, seq_len, num_classes)`num_classes=2，分别对应"不插入"（索引0）和"插入"（索引1）的概率"""   # 调用基础模型前向传播，返回二分类logitsreturn self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs).logits

2.3 MemGenWeaver

2.3.1 核心作用

MemGenWeaver 是 MemGen 框架的核心组件之一，负责生成动态潜在记忆并将其与推理器的输入序列融合，从而实现记忆与推理过程的无缝交织。它通过可学习的潜在记忆查询向量，在提示词阶段和推理阶段分别生成针对性的记忆表示，引导推理器调整解码路径，提升智能体的动态决策能力。

2.3.2 核心特色

• 双阶段记忆生成：区分提示词阶段（augment_prompt）和推理阶段（augment_inference），使用各自独立的可学习潜在记忆查询向量，适配不同阶段的记忆需求，增强记忆生成的针对性。
• 灵活的潜在记忆融合：通过_augment方法统一实现潜在记忆与输入序列的融合，包括嵌入拼接、注意力掩码扩展和位置 ID 计算，确保记忆与原始输入在语义空间和时序上的一致性。
• 高效的模型设计：
  • 基于预训练 LLM 构建，支持 PEFT 参数高效微调，在保留基础能力的同时降低训练成本；
  • 采用 bfloat16 精度和 Flash Attention 2 优化，提升计算效率和内存利用率。
• 动态记忆编织机制：生成的潜在记忆并非静态检索结果，而是基于当前输入序列动态生成的隐藏状态，能够捕捉实时上下文信息，实现 “生成式记忆” 的核心特性。
• 模块化与可扩展性：与推理器、触发器解耦，通过标准化接口交互；潜在记忆的数量可通过参数灵活配置，适配不同任务对记忆容量的需求。

2.3.3 网络架构

网络架构图如下。

说明如下：

1. 核心组件：

   • 可学习潜在记忆向量：分阶段设计（P=提示词阶段数量，I=推理阶段数量），支持动态生成记忆
   • 预训练LLM：作为记忆生成核心，默认启用bfloat16精度和Flash Attention 2优化
   • 序列融合层：确保输入与记忆在语义、掩码、时序上的一致性

2. 核心流程：

   • 输入 → 选择对应阶段的潜在记忆 → 融合序列 → LLM生成隐藏状态 → 提取潜在记忆输出
   • 支持PEFT参数高效微调（如LoRA），适配于Transformer Blocks层

3. 输出用途：

   • 生成的潜在记忆将通过投影层映射到推理器的嵌入空间，与原始输入融合以引导解码

2.3.4 代码

两个关键变量如下：

• prompt_query_latents。
  • 作用：增强模型在处理prompt时候的表现。 在模型处理完原始提示之后会被注入到序列中，为模型提供额外的上下文信息。
  • 使用场景：在 augment_prompt 方法中使用，在生成阶段的开始阶段使用一次。
• inference_query_latents。
  • 作用：在生成过程中动态增强模型的推理能力。可以在生成过程中的多个点被注入，以提供实时上下文增强。
  • 使用场景：在 augment_inference 方法中使用，在生成阶段中多次被使用。通常在遇到特定分隔符（逗号，句号等）后触发插入。

这两个变量都通过_augment 方法获得（获取学习到的潜在向量，并将其附加到输入嵌入中）。其流程如下：

• 将潜在变量附加到当前输入嵌入序列的末尾。
• 更新注意力掩码和位置ID，以考虑新增的潜在向量。
• 将增强后的序列通过Weaver模型处理。
• 提供于潜在向量位置对应的隐状态作为增强表示。

判断是否插入是通过函数 _should_augment 完成的。

• 检查当前生成的文本是否是特殊字符（逗号等）
• 使用触发模型（trigger model）进一步判断是否应该增强。
• 考虑最大增强次数限制。

class MemGenWeaver(torch.nn.Module):"""MemGen模型的编织器模块（MemGenWeaver）。- 输入：接收接收来自推理器模型当前当前解码序列的`inputs_embeds`（输入嵌入入）- 输出：生成长度为K的隐藏状态序列，
这些状态将与原始`inputs_embeds`拼接，以改变推理器的解码路径"""def __init__(self, pretrained_model_name_or_path: str,  # 预训练模型的名称或路径prompt_latents_num: int,    # 提示词阶段生成的潜在记忆数量inference_latents_num: int, # 推理阶段生成的潜在记忆数量peft_config: Optional[PeftConfig] = None  # PEFT配置（可选）):super().__init__()# 基础模型初始化self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path,torch_dtype=torch.bfloat16,  # 使用bfloat16精度以提高效率attn_implementation="flash_attention_2"  # 启用Flash Attentionention 2优化)self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)  # 对应的分词器# 若提供PEFT配置，则应用参数高效微调if peft_config is not None:self.model = get_peft_model(self.model, peft_config)self.config = self.model.config  # 保存模型配置# 提示词阶段的潜在记忆查询向量（可学习参数）self.prompt_query_latents = nn.Parameter(torch.randn(prompt_latents_num, self.config.hidden_size),  # 形状：[prompt_latents_num, hidden_size]requires_grad=True  # 允许反向传播更新)# 推理阶段的潜在记忆查询向量（可学习参数）self.inference_query_latents = nn.Parameter(torch.randn(inference_latents_num, self.config.hidden_size),  # 形状：[inference_latents_num, hidden_size]requires_grad=True  # 允许反向传播更新)@propertydef prompt_latents_num(self) -> int:"""返回提示词阶段的潜在记忆数量"""return self.prompt_query_latents.size(0)@propertydef inference_latents_num(self) -> int:"""返回推理阶段的潜在记忆数量"""return self.inference_query_latents.size(0)@propertydef device(self):"""返回模型所在的设备（CPU/GPU）"""return self.model.devicedef _augment(self, latents: torch.Tensor,                # 潜在记忆查询向量，形状：[latents_num, hidden_size]inputs_embeds: torch.Tensor,          # 输入嵌入，形状：[batch_size, seq_len, hidden_size]attention_mask: torch.Tensor,         # 注意力掩码，形状：[batch_size, seq_len]position_ids: torch.Tensor            # 位置ID，形状：[batch_size, seq_len]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:"""通用的潜在记忆增强方法：将潜在记忆与输入序列融合，生成增强后的隐藏状态。参数：latents: 潜在记忆查询向量inputs_embeds: 输入序列的嵌入表示attention_mask: 输入序列的注意力掩码position_ids: 输入序列的位置ID返回：三元组 (latents_hidden_states, latents_mask, latents_position_ids)- latents_hidden_states: 生成的潜在记忆隐藏状态，形状：[batch_size, latents_num, hidden_size]- latents_mask: 潜在记忆的注意力掩码，形状：[batch_size, latents_num]- latents_position_ids: 潜在记忆的位置ID，形状：[batch_size, latents_num]"""batch_size = attention_mask.shape[0]  # 获取批次大小latents_num = latents.size(0)         # 获取潜在记忆数量# 扩展潜在记忆维度以匹配批次大小：[1, latents_num, hidden_size] → [batch_size, latents_num, hidden_size]latents = latents.unsqueeze(0).repeat(batch_size, 1, 1)# 将潜在记忆嵌入与输入嵌入拼接：[batch_size, seq_len + latents_num, hidden_size]inputs_embeds = torch.cat([inputs_embeds, latents], dim=1)# 构建潜在记忆的注意力掩码（全为1，表示有效）并与输入掩码拼接latents_mask = torch.ones(latents.shape[:-1], dtype=attention_mask.dtype, device=attention_mask.device)attention_mask = torch.cat([attention_mask, latents_mask], dim=1)  # 形状：[batch_size, seq_len + latents_num]# 生成潜在记忆的位置ID（在输入序列最后位置的基础上递增）last_position_ids = position_ids.max(dim=1)[0]  # 获取输入序列的最大位置IDlatents_relative_positions = torch.arange(latents_num, device=attention_mask.device)  # 潜在记忆的相对位置# 计算绝对位置：输入序列最大位置 + 相对位置 + 1（避免重叠）latents_position_ids = last_position_ids.unsqueeze(1) + latents_relative_positions + 1# 拼接位置ID：[batch_size, seq_len + latents_num]position_ids = torch.cat([position_ids.long(), latents_position_ids.long()], dim=1) # 验证拼接后的维度是否一致assert inputs_embeds.shape[:2] == attention_mask.shape == position_ids.shape# 模型前向传播，获取隐藏状态outputs = self.model(inputs_embeds=inputs_embeds,attention_mask=attention_mask,position_ids=position_ids,  output_hidden_states=True,  # 输出所有层的隐藏状态)# 取最后一层的隐藏状态，并提取潜在记忆部分（序列末尾的latents_num个位置）hidden_states = outputs.hidden_states[-1]latents_hidden_states = hidden_states[:, -latents_num:, :]return latents_hidden_states, latents_mask, latents_position_idsdef augment_prompt(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:"""提示词阶段的潜在记忆增强：使用提示词专用的潜在记忆查询向量。参数与返回值同_augment方法"""return self._augment(latents=self.prompt_query_latents,inputs_embeds=inputs_embeds,attention_mask=attention_mask,position_ids=position_ids)def augment_inference(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:"""推理阶段的潜在记忆增强：使用推理专用的潜在记忆查询向量。参数与返回值同_augment方法"""return self._augment(latents=self.inference_query_latents,inputs_embeds=inputs_embeds,attention_mask=attention_mask,position_ids=position_ids)

0xFF 参考

最新成果！Agent记忆的第三种可能：生成式隐式记忆