- 链接:arXiv:2511.20639
- 代码:GitHub
核心愿景:构建一个无需训练(Training-free)、基于潜在空间(Latent Space)进行直接思维传输的多智能体系统。
本文将深入 LatentMAS 的核心代码逻辑,解析其如何通过 methods/ 和 models.py 实现论文中的分布对齐、潜在生成及记忆传输,并着重分析工程实现中针对不同后端(HuggingFace vs vLLM)所做的策略调整与理论偏差。
1. 核心基础设施:分布对齐 (\(W_a\)) 的数值解法#
论文核心痛点之一是 Latent Space 与 Input Embedding Space 分布不一致。论文提出通过线性投影 \(W_a\) 进行对齐。
代码位置: models.py -> ModelWrapper._build_latent_realign_matrix
def _build_latent_realign_matrix(self, model, device, args) -> Tuple[torch.Tensor, torch.Tensor]:
# ... 获取输入输出层权重 ...
# 构建岭回归 (Ridge Regression) 正规方程: (A^T A + lambda*I) X = A^T B
# 其中 A = output_weight, B = input_weight
# 1. 计算 Gram 矩阵 (A^T A)
gram = torch.matmul(output_weight.T, output_weight)
# 2. 添加正则化项 (Ridge term) 对应论文附录 A.1 公式 (8) 中的 lambda
reg = 1e-5 * torch.eye(gram.shape[0], device=gram.device, dtype=gram.dtype)
gram = gram + reg
# 3. 计算 A^T B
rhs = torch.matmul(output_weight.T, input_weight)
# 4. 求解线性方程组得到 Wa
realign_matrix = torch.linalg.solve(gram, rhs)
# ... 计算 target_norm 用于后续缩放 ...
return realign_matrix, target_norm
源码与论文对照:
- 一致性:代码完全复现了论文公式 (3) 及附录 A.1 的推导。它没有简单求逆,而是求解正规方程,且显式引入了
1e-5的正则化项(论文中的 \(\lambda\))以保证数值稳定性。 - 细节补充:代码额外引入了
target_norm,在对齐后对向量模长进行归一化。这是论文未详细展开但对工程稳定性至关重要的细节,防止潜在向量在多次循环后数值爆炸或消失。
2. 内部循环:自回归潜在思维生成#
这是智能体“默念”的过程,跳过 Token 解码,直接在向量层面推理。
代码位置: models.py -> ModelWrapper.generate_latent_batch
# 核心潜在生成循环
for step in range(latent_steps):
# 1. 对齐映射: h_t -> e_{t+1}
latent_vec = self._apply_latent_realignment(last_hidden, source_model)
latent_embed = latent_vec.unsqueeze(1) # 构造 [Batch, 1, Hidden]
# 2. 前向传播 (Forward Pass)
# 关键点:直接传入 inputs_embeds 绕过 Embedding 层
outputs = self.model(
inputs_embeds=latent_embed,
past_key_values=past, # 传入当前的 KV Cache
use_cache=True,
# ...
)
# 3. 状态更新
past = outputs.past_key_values
last_hidden = outputs.hidden_states[-1][:, -1, :] # 获取最后一个 Token 的隐状态
源码与论文对照:
- 一致性:完全符合论文 Section 3.1 的描述。循环中没有
softmax或argmax操作,纯粹是 Hidden State 到 Hidden State 的流转。 - 实现机制:利用 HF Transformers 的
inputs_embeds参数实现了这一逻辑,这是一种标准的“侵入式”推理控制手段。
3. 脑机接口:思维传输的双重实现路径#
这是整个框架中最核心的交互逻辑,但在代码中根据后端不同,出现了两条不同的路径。
路径 A:标准 HuggingFace 模式 (理论的完美复现)#
代码位置: methods/latent_mas.py -> LatentMASMethod.run_batch
# 遍历智能体链 (Planner -> Critic -> ...)
for agent in self.agents:
if agent.role != "judger":
# ...
# 生成潜在思维,并接收更新后的 past_kv
past_kv = self.model.generate_latent_batch(
wrapped_ids,
...
past_key_values=past_kv, # <--- 核心:直接继承上一个智能体的 KV Cache
)
# ...
else:
# Judger 智能体使用累积的 past_kv 进行最终文本解码
generated_batch, _ = self.model.generate_text_batch(
...,
past_key_values=past_for_decoding,
)
解读:
此路径严格对应论文 Section 3.2 的 Theorem 3.3 (Lossless Information Transfer)。智能体之间传递的是 past_key_values 对象(显存中的张量)。
- 计算复杂度:\(O(1)\)。接收方无需处理历史信息,直接在现有 Cache 上追加计算。
- 信息保真度:100% 无损。
路径 B:vLLM 混合模式 (工程上的妥协与变通)#
代码位置: methods/latent_mas.py -> LatentMASMethod.run_batch_vllm
由于 vLLM 的 PagedAttention 机制难以从外部直接注入非标准的 KV Cache,代码实现出现了显著的策略调整。
# 1. 使用 HF 模型计算潜在思维的 Embedding (而非 KV)
past_kv, previous_hidden_embedding = self.model.generate_latent_batch_hidden_state(...)
# 记录生成的向量
embedding_record.append(previous_hidden_embedding)
# ... 在 Judger 阶段 ...
# 2. 拼接历史 Embedding
past_embedding = torch.cat(embedding_record, dim=1).to(self.vllm_device)
# 3. 构造 "软提示" (Soft Prompt):将潜在向量拼接到文本 Embedding 中间
whole_prompt_emb = torch.cat([left_emb, past_embedding[i], right_emb], dim=0)
# 4. 传给 vLLM 进行生成 (vLLM 会对这些 Embedding 进行 Prefill 计算)
outputs = self.model.vllm_engine.generate(
[{"prompt_embeds": embeds} for embeds in whole_prompt_emb],
...
)
⚠️ 关键差异分析:
- 理论偏差:论文声称“传输 KV 以避免冗余重计算 (avoid redundant recomputation)”。但在 vLLM 模式下,传输的是 Embedding。vLLM 接收到这些 Embedding 后,必须在内部进行一次完整的 Prefill(预填充) 计算来生成 KV Cache。
- 工程考量:虽然违背了“无重算”的理论纯度,但利用 vLLM 极度优化的 Attention Kernel,这种“重算”在实际 Wall-clock time 上依然比 HF 的“不重算”要快得多。这是一种用计算量换取利用高性能算子机会的工程策略。
4. 上下文管理与 Prompt 依赖#
4.1 记忆截断 (Truncation)#
代码位置: methods/latent_mas.py -> _truncate_past
论文主要描述了累积式的思维增长,但代码中增加了一个控制机制:
if self.sequential_info_only or self.latent_only:
# 仅保留新增的 latent steps 对应的 KV,丢弃之前的历史
past_kv = self._truncate_past(past_kv, tokens_to_keep)
这表明 LatentMAS 支持一种“遗忘模式”,只传递最新的思维切片。这是论文中未详细讨论的显存优化手段,防止长链条协作时 KV Cache 无限增长。
4.2 显式的 Prompt 引导#
代码位置: prompts.py
虽然 LatentMAS 强调“潜在协作”,但代码显示它高度依赖自然语言 Prompt 来配置智能体状态:
# Judger 的 Prompt
user_prompt = """
You are provided with latent information for reference...
The latent information might contain irrelevant contents. Ignore it if it is not helpful...
"""
这意味着,\(W_a\) 的数学对齐只是必要条件,还需要通过 Prompt Engineering 明确告知模型“你将接收一段非自然语言的潜在输入”,模型才能正确处理这些向量。这补充了论文中关于系统构建的隐含前提。
5. 总结#
LatentMAS 的源码实现通过 ModelWrapper 精巧地封装了底层张量操作。
- 对于算法验证:Standard HF 模式是论文理论的忠实实现,实现了真正的 KV 零拷贝传输。
- 对于高性能部署:vLLM 模式采用了一种“Embedding 拼接 + 重新 Prefill”的变通方案。虽然在机制上与论文强调的“无重算”有出入,但它成功解决了 vLLM 架构兼容性问题。