A Domain-Specific Probabilistic Programming Language for Reasoning about Reasoning (or: a memo on memo)

面向推理的领域专用概率编程语言(或:关于 memo 的备忘录)

https://comosoco.daeh.info/assets/papers/chandra2025memo.pdf

https://osf.io/preprints/psyarxiv/pt863_v1

https://github.com/kach/memo

代码库缩小了 3 倍或更多 ),并且执行和适应数据的速度大大加快(我们已经看到 3,000 倍或更多 的加速)。在惯用的备忘录中,POMDP 求解器是 15 行代码,并且与用 200 行代码编写的手动优化求解器一样快。

需要将我们对心智行为的直觉表达(如 believe、think、see、do、want、imagine 等动词)翻译成概率编程语言的基本操作(如 sample、condition、infer 等)。这种翻译过程中可能会引入一些细微的 bug,导致模型违背我们对心智行为的基本直觉。

摘要

人类对思维本身的思考能力(即“心理理论”)是许多学科研究的基本对象。在最近几十年里,来自不同学科的研究人员逐渐采用了一种丰富的计算范式来建模心理理论,这种范式基于递归的概率推理。然而,实践者常常发现以这种范式进行编程极具挑战性:一方面,因为“思考他人的思考”本身就让程序员感到困惑;另一方面,这类模型运行起来极其缓慢。本文介绍了一种新的、面向心理理论的领域专用概率编程语言——memo。它克服了上述两个挑战:首先,通过为心理理论提供专门设计的语法和语义结构,使建模更加直观;其次,采用了一种独特的推理方法,利用数组编程在现代硬件上实现了良好的扩展性,从而显著提升了性能。使用 memo,研究人员可以用更少的代码编写出运行速度更快的模型。该语言已经得到了多个研究团队的实际应用与采纳。

1 引言

语言的生命力在于它能够描绘其使用者、读者和作者的实际生活、想象生活以及可能的生活。

托妮·莫里森,《诺贝尔文学奖演讲》,1993年

句子“Mary thought she saw a ghost”(玛丽以为她看到了一个鬼)中,主句是真实的(“Mary thought”,即玛丽以为),而嵌入的从句是虚假的(“she saw a ghost”,即她看到了鬼)。这种语法结构可能为语言学习者提供了理解他人错误信念的逻辑工具。

Pyers 和 Senghas [108]

人类不仅具备思考的能力,还能够思考“思考”本身:我们的心智可以对其他心智进行递归推理。这种非凡的能力被称为“心理理论”(theory of mind),它几乎支撑着我们所做的一切:从玩游戏、解释概念、说服陪审团,到发起合作、安慰朋友、识别误解,甚至发动战争。

由于心理理论在人类经验中的核心地位,许多学科都试图通过计算模型来理解它。这些学科包括认知科学 [12, 13, 128]、社会心理学 [2, 3, 15, 35, 77, 80, 138]、语言学 [76, 87, 118, 122, 139]、伦理学 [91, 111, 112] 以及行为经济学 [8, 27, 86, 123]。这类模型经常被用于现实世界系统的工程实践,如机器人技术 [45, 50, 101]、教育 [29, 113, 120]、人机交互 [31, 38]、程序综合 [1, 106, 130] 和计算机图形学 [28, 33] 等领域。

有趣的是,跨学科的心理理论模型往往比它们表面上看起来更为相似。它们共同的核心理念包括贝叶斯推理(因为心智持有不确定条件下的信念)、决策理论(因为心智会理性地追求目标)以及递归性(因为心智会思考其他心智)。来自不同学科的教科书也以类似的方式介绍这些理念如何在建模心智时浮现并运作 [47, 67, 69, 90, 96, 118]。我们推测,这里隐藏着一种完整的编程范式——一种概率编程范式,我们称之为“递归理性”(Recursive Rationality,RR)范式。

RR 模型可以通过使用嵌套推理查询来优雅地实现递归贝叶斯推理 [125]。然而,根据我们多年与实践者的合作得出的结论,即使使用当今最好的概率编程语言(PPLs),在 RR 范式下编程仍然极具挑战性。正如编程领域常见的那样,这些挑战主要体现在两个方面:正确性 (correctness)和效率 (efficiency)。

关于正确性的难题在于抽象层面。心理理论拥有一套丰富的概念词汇(例如表示心理状态的动词,如 believe、think、see、do、want、imagine),但大多数 PPL 只提供少量的基本操作(动词,如 sample、infer、condition)。程序员必须将前者翻译成后者,但由于无法检查语义是否在翻译过程中保持一致,他们常常会遭遇细微的 bug,比如不小心让一个智能体“读取”了另一个智能体的思想,混淆了属于不同心智的变量。

关于效率的问题则源于“思维中的思维”会导致可能世界状态的组合爆炸(不仅在一个心智内部,还包括其他心智所想象的心智状态……),从而导致推理过程极其缓慢。

我们认为这两个问题是阻碍科学进步的重大障碍:在各个学科中,科学家通常需要数周的实现工作和数天的计算时间来验证他们的想法。

在本文中,我们提出了一种新的、面向 RR 范式的领域专用概率编程语言——memo(意为“心智模型”)。具体来说,我们做出了以下贡献:

(1) 我们识别出 RR 实践者常遇到的一类独特 bug(§1.1.1),并设计了一种语言,其语法和语义能够保护用户免受此类 bug 的困扰(§3.1)。

(2) 我们分析了 RR 实践者在实际应用中面临的可扩展性挑战(§1.1.2),并通过观察发现 RR 模型特别适合现代硬件加速(§1.2.2)。因此,我们在推理方法上采取了独特的策略:我们将模型转换为数组程序,进而编译为高效的并行代码(§3.2)。我们的推理算法是端到端可微分的,支持通过梯度下降快速拟合参数。

(3) 通过对几个经典 RR 模型的案例研究,我们展示了 memo 的设计决策使实践者能够用比现有通用 PPL 中专家编写实现少得多的代码写出运行速度显著提升的模型(§4.1)。

(4) 我们使得一些新型、复杂的模型成为可能,这些模型正是 RR 实践者日益关注的方向:高效整合进整个深度学习生态系统的模型(§4.3.1),以及能够反思执行推理所需计算成本的模型(§4.3.2)。

(5) 我们将 memo 作为开源软件发布(https://github.com/kach/memo )。memo 已经被多个研究团队采用,带来了数量级级别的性能提升,并使他们的模型得以进行雄心勃勃的新拓展(§4.2)。

在本引言的其余部分,我们将讨论促使我们设计 memo 的两个关键挑战(§1.1),并预览 memo 设计背后的两个核心洞察(§1.2)。随后,我们将通过示例介绍 memo(§2),探讨其实现机制(§3),并通过一系列案例研究对其进行评估(§4)。

1.1 推动 memo 诞生的两个挑战 1.1.1 正确性(Correctness)

正确实现 RR(递归理性)模型之所以困难,原因有很多。最明显的一点是,涉及多个智能体和递归层次的情境很容易变得混乱和令人困惑(例如可以看看这部来自情景喜剧《老友记》的片段:youtube.com/watch?v=a4CS2tCjAgk)。即使对于研究人员完全理解的简单模型,在代码实现上也可能非常棘手且容易出错。

一类常见的简单错误来源于贝叶斯计算实现中的错误:许多研究者在 R 或 MATLAB 中手动实现推理过程,很容易忘记对概率分布进行归一化等操作。这类问题可以通过使用概率编程语言(PPLs)来自动执行嵌套推理加以解决 [125],因此我们在此不再赘述。

更有趣的一类错误则即使在使用 PPLs 的情况下也依然存在。虽然现有的 PPLs 可以保证概率计算在数学上的正确性(根据该 PPL 的语义),但它们并不能保证这些计算对于我们关于“心智能动性”(agency)的直观理解是有意义的。正如我们之前所讨论的那样,在使用 PPL 编写 RR 模型时,我们需要将我们对心智行为的直觉表达(如 believe、think、see、do、want、imagine 等动词)翻译成概率编程语言的基本操作(如 sample、condition、infer 等)。这种翻译过程中可能会引入一些细微的 bug,导致模型违背我们对心智行为的基本直觉。

让我们来看一个具有代表性的例子。在 RR 模型中一种常见的模式是:一个智能体根据某个效用函数(utility function)做出选择。具体来说,假设某天早上 Alice 根据一家餐厅的食物质量和价格权衡后决定预订哪家餐馆。传统上,这类决策理论情境在 PPL 中通常通过随机均匀采样一个选项,然后根据智能体的效用函数对其施加因子(factor)来进行编码 [47]。例如,在 WebPPL [66] 中,我们可以这样写:

这里的 factor 语句会让推理机制更倾向于高价值的选择路径(体现了“以推理为规划”的理念 [25])。因此,推断出的 r上的分布就可以预测 Alice 的行为:她偏好质量更好、价格更低的餐厅,并在相似选项之间保持中立。

然而,这种 samplefactor 之间的“远程作用”(spooky action at a distance)可能导致微妙的 bug。举个例子,假设 Alice 略微偏好户外餐厅而非室内餐厅(比如质量评分为 11 对比 10),但如果当晚可能有暴风雨,她在户外会面临极高的代价(比如惩罚值为 −100)。我们可能会期望通过引入一个新的随机变量 s 来表示暴风雨的发生,并将其纳入因子中来表达这个额外的需求:

然而,这种看似无害的编码方式实际上是错误的。它无意中将环境中的随机性“控制权”赋予了 Alice,从而错误地预测:Alice 会在一种她能影响天气的幻觉下做出高风险、过于自信的预订决定!在这个例子中,尽管我们的直觉告诉我们 Alice 应该采取保守策略选择室内餐厅,而且选择户外餐厅的预期效用也显著更低(−39 ≪ 10),但 WebPPL 的推理却预测出 Alice 有 58% 的概率会选择户外餐厅。(关于这一问题在强化学习背景下的更多讨论,请参见 Levine [92, §2.3]。)

类比自动微分中的“扰动混淆”(perturbation confusion)[121],我们可以将这类 bug 称为“作恶者混淆”(perpetration confusion):即搞混了哪一个智能体应该对哪一项决策负责。

正如 Siskind 和 Pearlmutter 解决扰动混淆的方法一样,我们提出的解决方案也是类似的:通过给每个选择打上“智能体标签”来区分不同智能体的责任。我们将在第 4.1.2 节中展示一个现实世界中的 perpetration confusion 示例,以及 memo 是如何解决这个问题的。

1.1.2 效率(Efficiency)

RR 模型以运行速度慢而闻名,主要原因在于不同智能体的心智模型中可能出现的世界状态呈组合式爆炸增长(包括它们彼此之间的心智模型!)。这严重限制了研究人员所能探讨的问题范围和规模。

嵌套推理本身也带来了额外的挑战。当嵌套使用近似推理方法时,“内部”推理算法必须产生足够精确的结果,供“外部”推理算法使用。通常这意味着要增加采样数量,而这会极大地拖慢关键的内层循环(参见 §4.2.1 中的一个真实案例)。为每一层递归选择合适的推理算法也给程序员带来了沉重负担,尤其是随着需要调整的超参数数量的增加。

最后,在实际研究中,经验科学家不仅运行模型,还需要将模型的自由参数拟合到实验收集的数据上。这通常是通过网格搜索完成的,即使是极其简单的模型也可能需要数百万次模型评估(参数数量的指数级增长)。研究人员可能还希望为多个不同的实验条件(“处理组”)分别拟合模型、对模型进行多种不同的消融实验,并通过对数据的不同子集独立拟合来交叉验证结果。这些需求会极大放大单次模型评估的成本,使得即便是相对简单的模型也变得非常昂贵。

我们采访的研究人员经常提到,仅推理过程就需要几分钟到几十分钟不等,而模型拟合与交叉验证则往往需要花费多达数天的计算时间

这些问题使得效率成为影响研究生产力和进展的根本性瓶颈:模型拟合的高昂成本使得研究人员难以对模型进行调试、验证直觉以及迭代改进想法。同时,它也迫使研究人员不得不缩小模型规模,或进行比他们期望中更不彻底的超参数调优。

当然,网格搜索和交叉验证属于“极易并行化”(embarrassingly parallelizable)的任务,在现代硬件上可以实现显著加速。但认知科学家通常缺乏必要的专业知识,无法识别甚至利用那些非常容易实现的性能优化机会。

同样地,基于梯度的优化方法无疑会比网格搜索更加高效,但目前只有极少数概率编程语言(PPL)支持在推理过程中自动微分(automatic differentiation)。

除了这些实际问题之外,还有一个重要的科学问题悬而未决。尽管 RR 模型可能非常准确地预测人类的直觉反应,但它们有时也因计算开销过大而受到批评,被认为不太可能是人类大脑产生这些直觉的合理机制 [97, 134]。而实现更快的模型执行方式,为科学家提供了一种解决这一矛盾的途径。

1.2 memo 的设计:两个核心洞察 1.2.1 关于正确性:“尊重我们对心智的元理论”

如前所述,我们在正确性 方面的核心洞察是:提供专门的语法结构,将随机选择与智能体(agent)明确关联。

在 memo 中,我们会这样实现餐馆选择问题:由 Alice 来选择餐馆,而世界(环境)来决定天气 (完整实现请参见补充材料)。这种专用语法使我们能够静态地捕捉并防止诸如“作恶者混淆”(perpetration confusion)之类的错误。

更一般地,通过在语言的语法和语义中显式表达“心智能动性”(agency),memo 能够静态强制执行我们关于心智直觉的多种基本条件,确保模型与智能体行为的基本原则保持一致。例如,memo 可以防止模型中出现无意的“心灵控制”或“读心术”。

我们将在第 3.1 节中详细讨论 memo 的静态语义机制,以及它是如何体现智能体行为的关键原则的。

1.2.2 关于效率:“PPL 遇见 APL”

我们在效率 方面的核心洞察是:大多数现实中的 RR 模型都适合进行精确枚举推理 (exact enumerative inference):从结构上看,它们只需要从基数不大的离散分布中进行有限且静态可知的一系列采样。

幸运的是,近期研究表明,在这类模型中,精确枚举推理往往可以被编译为高效的数组程序(array programs),其中数组表示离散域上的联合分布 [103, 104]。这表明我们有机会大幅加速 RR 模型的运行。

我们可以通过研究人员描述 RR 模型的方式看出这一方向的正确性。许多经典模型在文献中通常有两种互补的描述方式:一种是使用递归树的语言,另一种是使用数组的语言。

例如,著名的值迭代算法(value iteration)有时被描述为带记忆化的图搜索 [16],有时则被描述为对价值表(value table)的迭代操作 [19]。同样,“理性言语行为模型”(Rational Speech Acts model,详见下文)有时被描述为嵌套推理 [52],有时则被描述为对“真值表”(truth table)的操作 [106]。经济学家也常常同时使用“博弈树”(game trees)和“收益矩阵”(payoff matrices)来进行建模。

基于这一点,memo 被设计为能将嵌套的枚举推理编译为数组程序。这种方式为用户带来了诸多直接优势:

  1. 推理是精确的

    ,而且没有超参数;

  2. 数组程序具有简单的循环结构, 易于编译为高效优化的代码

  3. 推理可以 天然地并行化 ,无论是单次评估内部(在可能的世界状态之间),还是在参数拟合过程中(在数据批次之间);

  4. 现有的数组程序编译器可以在 GPU 和其他硬件平台上 加速推理过程

  5. 推理支持 原生自动微分 (automatic differentiation),从而极大地加快了参数拟合速度;

  6. 模型可以无缝接入宿主数组语言的库生态系统,包括深度学习、数据可视化等多个工具包。

我们将在第 3.2 节中详细讨论 memo 是如何将模型编译为数组程序的。

2 memo 的演示

我们将通过 RR(递归理性)范式中的“果蝇”或“Hello, World”来介绍 memo:来自认知语言学中极具影响力的理性言语行为 (Rational Speech Acts,RSA)框架 [43, 52, 54, 64]。RSA 模型描述了人们如何做出超越语言字面语义的语用推理 (pragmatic inferences)。

我们将考虑一个展示这种推理的经典实验,并使用 RSA 来建模其实验结果。(需要说明的是,我们的目标并不是要推广 RSA 而贬低其他理论——RSA 本身也有批评者 [55, 58, 140]。在这里,我们只是将其作为一个典范性例子来引入 memo。)

2.1 理性言语行为框架:我们的示例

考虑以下由哲学家 Lewis [93] 首次提出的经典游戏。想象有三个物体:一个绿色正方形 ☐,一个绿色圆形 ●,以及一个粉色圆形 ○。你的朋友被要求让你从中选择一个,但他只能使用“green(绿色)”、“pink(粉色)”、“square(正方形)”或“round(圆形)”这些词。假设你的朋友说“pink(粉色)”,你会选择哪个?大多数人会选择粉色圆形 ○。

同样地,如果他说“square(正方形)”,大多数人会选择绿色正方形 ☐。到目前为止,这都很直观。但现在,假设你的朋友说“green(绿色)”。这更复杂了,因为绿色正方形 ☐ 和绿色圆形 ● 都符合“green”的字面意义。然而,认知语言学中一个令人惊讶且可复现的发现是:大多数人会在语用层面偏好绿色圆形 ● 而非绿色正方形 ☐ ——例如,在一项研究中,有 64% 的人选择了前者 [109]。这是怎么回事呢?

似乎人们推理如下:如果朋友想让你选绿色正方形 ☐,他本可以明确地说“square”(无歧义),而不是使用模糊的“green”。因此,既然他选择了风险更大、含义模糊的“green”,很可能他是想区分两个圆形 ● 和 ○。

RSA 框架试图以精确的计算方式形式化这一直觉。它将此类情境建模为说话者(S)和听者(L)在不确定性下的相互递归推理。我们假设,最初听者对说话者所指的对象有一个先验信念,其中 r 表示“指称对象”(referent,即说话者想要指代的物体)。给定说话者的言语 u,听者可以通过贝叶斯定理推断出后验概率:

其中 pS(u∣r) 是说话者为了表达对象 r 而选择说出 u 的可能性。

关键在于:听者可能会认为,一个理性、策略性的说话者会选择最有可能让听者正确理解自己意图的词语 u。因此,pS本身可能递归地依赖于 pL。(这就是“递归理性”范式的本质所在。)

2.2 初识 memo:在 memo 中表达“基础情形”的天真听者模型

要在 memo 中表达这个模型,我们首先从实现一个天真听者 (naïve listener)的“基础情形”开始。这个听者认为说话者只是随机选择一个真实的言语表达

对于无歧义的言语表达 (如 “pink” 和 “square”),模型正确地分别选择了 ○ 和 ☐;

而对于有歧义的言语表达 (如 “green” 和 “round”),模型在两个语义上都成立的选项之间没有偏好

这与我们的直觉相悖,也与 Qing 和 Franke 的实验数据不符 [109],说明我们还需要进一步完善模型……

程,如图2所示。

2.4 使用我们的模型:一窥科学家的工作流程

正如我们在第 1.1.2 节中所讨论的那样,一个 RR(递归理性)模型研究者的任务并不仅仅止于表达模型:他们通常还需要将模型拟合到实验数据 上,并对其进行验证 。下面我们来看看在 memo 中是如何完成这一过程的。

2.3 完善模型:加入递归推理

现在,我们来看看如果听者认为说话者也在思考听者 (而听者可能又在思考说话者,以此类推)时会发生什么。我们将通过三个步骤扩展我们的代码来建模这一过

打开网易新闻 查看精彩图片

图 3 展示了通过网格搜索(耗时约 70 毫秒)和梯度下降(耗时约 250 毫秒)将 β 拟合到 Qing 和 Franke 的数据的过程,并展示了在 ℓ=1 和 ℓ=0 时,模型与人类数据的拟合情况。

显然,当 ℓ=1 时,模型能够很好地预测人类的行为;而当 ℓ=0 时,模型的预测结果与人类行为存在显著差异。虽然还需要进行更多的统计分析,但这类结果可以用来论证:人们是在 ℓ=1 的层级上进行推理,而不是 ℓ=0。

2.5 语言的其余部分

至此,我们对 memo 的介绍也接近尾声。图4中给出了该语言的完整抽象语法,其中只包含几个我们尚未介绍的语言结构,下面我们将重点说明它们。

2.5.1 以虚拟语气进行编程(Programming in the subjunctive mood)

智能体常常通过假设和反事实来推理,提出“如果……会怎样?”(what if?)之类的问题,以此指导他们的行为。有大量证据表明,人类在做决策时确实会进行这类模拟 [17, 62, 145]。

为了在 memo 中建模此类推理,我们引入了 imagine 表达式,它允许在一个临时的作用域中想象一个事件 ,然后在那个假设的世界中求值一个表达式。(imagine 类似于 JavaScript 中的 IIFE(立即调用函数表达式),它可以被“解糖”为对一个新生成的模型的调用。)

2.5.2 对“不思考”的思考(Thinking about not thinking)

认知科学中的最新研究探讨了智能体如何考虑“思考本身”的成本 [94, 133]。例如,当某种近似已经足够好时,继续深入思考(在时间、能量、记忆等方面的额外开销)可能就不值得了。

对这种“资源理性”(resource-rationality)的建模一直具有挑战性,因为在传统的概率编程语言中很难反思推理过程的计算成本。

memo 提供了一种机制来反映计算成本:它利用 JAX 的编译时分析提取所需的浮点运算次数(FLOPs)。memo 模型可以使用这些信息,让智能体基于资源考量做出选择。

我们在第 4.3.2 节中展示了这一功能的使用示例。

3 设计与实现

我们的目标是创建一种语言,帮助科学家编写正确且高效的 RR 模型 (递归理性模型)。在本节中,我们将展示 memo 是如何被设计来达成这些目标的。

我们首先介绍 memo 的前端设计 ,解释其静态语义是如何确保模型符合我们对“智能体行为”(agency)的直觉理解,从而保障正确性 (§3.1);

接着,我们描述 memo 模型是如何被编译为数组程序 ,以实现高效推理 (§3.2);

最后,我们讨论支持 RR 研究者的实际考量因素 ,以及它们如何影响了我们的设计决策(§3.3)。

在这里,我们专注于提供具体而直观的解释。有关 memo 静态语义和指称语义的形式化定义,请参见同步进行的研究工作 [5, 匿名引用],该研究建立在 Zhang 和 Amin [143] 对嵌套推理形式化基础上的工作成果之上。

3.1 前端:追踪“心智框架”(frames of mind)

正如我们在第 1.2.1 节中所讨论的,memo 的设计目标是确保模型与我们对现实世界中“智能体行为”(agency)的直觉保持一致。在本节中,我们将通过概述 memo 的静态语义 来说明它是如何实现这一目标的。

在此过程中,我们会强调关于“智能体行为”的四个基本原则(标记为 ① 至 ④),并解释这些原则是如何在 memo 的设计中被具体体现出来的。

“心智框架”(frame of mind):memo 编译器的核心数据结构

memo 编译器的核心数据结构是“心智框架”(frame of mind),它用于跟踪一个智能体的信念(beliefs),以及该智能体对其他智能体信念的信念。

为了引入“心智框架”的概念,我们将继续使用第 2 节中的 RSA 示例,并逐步展示在编译过程中“心智框架”是如何演化的。

如图5所示,一个“心智框架”包含:

  • 一组智能体(包括自身),

  • 这些智能体所做的选择(choices),

  • 指向这些智能体各自心智框架的指针(从而形成嵌套的框架树)。

在一个智能体的框架中,其做出的选择可以是已知的 不确定的 (即随机变量)。

面板 A:初始“根框架”

每个 memo 模型都从一个初始的“根框架”开始,我们称之为观察者框架 (observer’s frame)。

在这个框架中,只有两个变量 r 和 u,它们由模型顶层定义中的维度声明绑定(例如 def L[r: R, u: U]...)。这两个变量对观察者来说是已知的 ,也就是说,它们的值没有不确定性。

面板 B:建立听者对说话者的心理建模

当我们初始化听者对说话者的心理建模(如 listener: thinks[speaker: chooses(...)])时,memo 会为听者创建一个新的心智框架,并为其想象中的说话者再创建一个新的心智框架。

观察者的框架中包含一个指向听者框架的指针,而听者的框架又包含一个指向其所想象的说话者框架的指针。

此时,在说话者的框架中包含了两个变量,这两个变量对说话者而言是已知的 ,因为它们是由说话者自己选择的。

然而,在听者的框架中,这两个变量则是不确定的 ,即被视为随机变量。

这引出了 memo 所体现的第一个关于“智能体行为”的基本原理:

① 不允许“读心术” (No Mind Reading) 智能体总是知道自己做出的选择,但默认情况下并不知道其他智能体的选择,后者对其来说是不确定的。

这一原则体现在语言设计中:不同心智框架之间的变量不能随意共享;只有通过明确的推理机制才能访问他人心智中的信息。

面板 C:听者观察到说话者的言语

面板 C 展示了当听者观察到说话者发出的言语 usaid 等于 u时的情况。这个语句产生了两个效果:

  1. 在听者的框架中, speaker.u_said 的状态从“不确定”变为“已知”。

  2. 在观察者的框架中,听者的 speaker.u_said 被绑定到变量 u (这一信息存储在观察者框架的一个辅助表中)。

自此之后,观察者在建模听者计算过程时,会将 u 替换掉 speaker.u_said。例如,从观察者的视角来看,布尔表达式 listener[speaker.u_said == "green"] 就等价于 u == "green"

请注意,这是模型中唯一一处 将听者的 speaker.u_said 与观察者的 u 关联起来的地方。我们也可以写成 listener: observes [speaker.u_said] is w,其中 w=u。这种设计是有意为之的。

它体现了关于智能体行为的第二个基本原理:

② 智能体可以形成错误信念 (Agents Can Acquire False Beliefs) 观察者知道事实,但听者可能并不知道,并因此持有错误信念。

例如,我们可以模拟一个“中间人攻击”场景:listener: observes [speaker.u_said] is attacker.u_fake。此时听者会认为说话者真的说了 u_fake,而观察者则清楚地知道听者实际上收到了攻击者的消息。

结果是,在观察者的框架中,listener[speaker.u_said]attacker.u_fake 是等价的,但在听者的框架中却不是——听者甚至不知道攻击者的存在。

这体现了第三个关于智能体行为的基本原理:

③ 信念具有指称不透明性 (Referential Opacity of Belief)[56, 110] 同一个对象在不同智能体的信念中可能以不同的方式被表示和理解。

此外,请注意,观察者的框架中并没有包含对说话者的 r_trueu_said 的直接引用 。这也是有意为之的设计。

观察者和听者不必对说话者拥有相同的心理模型

例如,观察者可能认为现实中说话者只是随机胡言乱语,根本没有 r_true 影响其言语;而听者却可能错误地将说话者的言语解释为有目的的行为,并赋予其一个 r_true 的含义(参见 [44, 74])。因此,r_true 应该只出现在听者的框架中,而不是观察者的框架中。

面板 D:听者做出选择

面板 D 展示了当听者做出选择 rpick 时的情形。正如我们在面板 B 中讨论的那样,r_pick 在听者的框架中是一个已知变量 ,而在观察者的框架中则是不确定的

借此机会,我们来看看 memo 所体现的最后一个关于智能体行为的基本原理:

在第 2.2 节中,我们强调了一个重要约束:听者必须知道自己选择 rpick 的概率分布

这一限制实质上体现了以下原则:

④ 不允许“心灵控制” (No Mind Control) 如果一个智能体无法理解自己为何做出某个选择,那就仿佛是别人替他做出了这个决定。

例如,如果 Alice(秘密地)选了一个冰淇淋口味,并期望 Bob 也选同样的口味,那么她对 Bob 的心理建模就剥夺了 Bob 的自由意志。

这类情形会导致像 perpetration confusion(作恶者混淆) (§1.1.1)这样的 bug,更广泛地说,还会引发哲学上的难题,如 Newcomb 悖论 [102, 135]。

memo 通过静态跟踪每个表达式在一个智能体框架中是否“已知”来防止此类问题的发生。

  • 如果一个表达式涉及未知的选择,则它会被标记为“未知”(除了 E[e] 这种期望值表达式,即使 e 是未知的,它本身仍被视为“已知”);

  • memo 强制要求: 智能体在指定自身选择的概率时,必须知道所使用的表达式

这种方式确保了模型不会违反关于“心智自主性”的基本原则,从而避免了许多潜在的逻辑错误和哲学困境。

3.2 后端:将 memo 降低为数组程序

正如我们在第 1.2.2 节中所讨论的那样,RR 模型通常只需要离散枚举推理 (discrete enumerative inference),因此非常适合被编译为数组程序 。在本节的其余部分,我们将通过解释 memo 是如何将我们的 RSA 模型编译为数组程序来完成对示例的说明。

核心思想

memo 的核心思想是:为每一个“心智框架”维护一个数组 ,该数组表示该智能体当前对世界的信念。

  • 这个数组为该框架中跟踪的每一个选择(choice)保留一个维度;

  • 数组中的每一项表示:在已知选择的条件下,未知选择的联合概率。

图6展示了与听者框架相关联的数组是如何在 RSA 模型的三个语句中逐步演化的(为简化起见,设 ℓ = 0)。我们重点讲解听者框架,因为它是最有趣的:观察者和说话者没有不确定的选择,所以他们的数组只是填满了 1。

听者的初始信念数组

听者的信念数组 初始是一个零阶数组 (即无维度的标量),其值为 1:

因为此时听者尚未建模任何不确定的选择,因此只有一个可能的世界状态,听者为其赋予概率 1。

Panel B2 展示了当听话者模拟说话者选择“said”时,同样的过程展开。 首先,用一个新的大小为||的维度来扩展“”,用于“said”,使其成为一个二维数组(以热图形式展示)。然后,memo 对每个“said”和“true”评估可能性表达式“denotes(said, true)”,并沿着“said”轴对结果进行归一化(请注意,这两个操作都可以在现代硬件上轻松并行化)。最后,memo 通过逐点乘以归一化的可能性来更新“”。因此,“”的更新如下:

请注意,通过机械地遵循 memo 的编译规则 ,我们实际上“免费”得到了 Pu 等人 [106] 所描述的 RSA 模型的高效数组实现:即对一个大小为 || × | | 的数组进行行归一化与列归一化的交替操作

我们将在第 4.1.3 节中看到,memo 也以同样的方式自然地还原了经典的表格型值迭代算法(tabular value iteration)。

利用数组 ,我们可以从听者的视角出发,执行多种有趣的计算。例如,假设我们想了解听者对“true 是绿色形状”的信心程度,也就是在听者观察到 usaid 后,在其心智框架中求表达式 Pr[green(true)] 的值。

表达式 Pr[] 在去糖化后等价于 E[],因为伯努利随机变量的期望值就是它的发生概率。

因此我们的目标是计算:

我们可以将这一计算表示为张量收缩 (tensor contraction)来实现:

  1. 首先,在每一个可能的世界状态中计算 green(r_true) 的值,并将结果组装成一个数组 e ;

  2. 然后,将 e 与数组 沿着代表不确定变量的维度进行张量收缩。

在这个例子中,只有一个不确定的维度(即 rtrue),因此我们可以将这个张量收缩写成矩阵乘法的形式:

这样,我们就能够高效地从听者的信念数组中提取出关于世界属性的概率性推理结果。这种基于数组的操作不仅结构清晰,而且可以充分利用现代硬件的并行计算能力,从而实现高效的推理过程。

这表明,对听者来说,Pr[green(r_true)] 的值取决于所观察到的 usaid:

  • 如果他们听到的是 “green”,则认为概率为 1(当然如此!);

  • 如果听到的是 “pink”,则概率为 0;

  • 如果听到的是 “square”,概率也为 1(因为唯一的正方形是绿色的);

  • 如果听到的是 “round”,概率为 1/2(因为两个圆形中只有一个为绿色)。

小结

从前面两节中我们可以总结出几个关键点:

  • 静态层面

    上,memo 维护一棵“心智框架树”(tree of frames),用于追踪智能体、它们的选择,以及它们对其他智能体选择的 知识或不确定性

  • chooses

    创建一个选择:对该选择者来说是 已知的 ,但对建模该选择者的智能体来说是 不确定的

  • observes

    让一个原本不确定的选择对观察者来说变得 已知

这种机制确保了模型在语法和语义层面都能正确表达智能体之间的信念与推理关系,从而避免诸如“读心术”或“心灵控制”等逻辑错误。

智能体必须知道自己做出选择的概率分布。这些规则确保了 memo 模型与我们对“心智能动性”(agency)的直觉保持一致。

在运行时层面 ,memo 为每一个心智框架维护一个数组,该数组表示该智能体在已知选择条件下对不确定选择的联合概率分布。

  • 数组的每个维度对应框架中跟踪的一个选择;

  • chooses

    会引入新的维度;

  • observes

    会对特定维度进行归一化;

  • E[e]

    则通过张量收缩来计算。

这些操作都可以在现代硬件上轻松实现并行化,从而使得推理过程高效运行。

3.3 实现细节

在设计 memo 时,一个关键考量是:与科学家们偏好的工作流保持良好的互操作性

研究人员常常面临的一个痛点是:需要在用于建模的概率编程语言(如 WebPPL)和用于数据分析的脚本语言(通常是 Python、R 或 MATLAB)之间来回传递数据。事实上,我们的许多同事都曾报告说他们手动复制粘贴输入输出内容,或者编写脆弱的 shell 脚本来生成和解析日志文件。

因此,我们选择将 memo 嵌入 Python 中实现 ,并基于 JAX 数组编程库进行构建。这一决定带来了额外的好处:

  • 使 memo 立即兼容现有的 Python 工具链(IDE、Notebook 等);

  • 无缝接入以数组编程为核心的科学计算 Python 生态系统(包括广泛使用的数据分析与可视化包);

  • 并能够利用快速发展的 JAX 生态(包括深度学习、物理仿真和强化学习相关包)。

为了在 Python 中嵌入 memo,我们提供了 @memo 装饰器,它会:

  1. 解析模型定义的源代码;

  2. 将其对应的 Python 抽象语法树(AST)翻译成 memo 的抽象语法树;

  3. 然后将其发送给 memo 编译器处理。

接着,我们将 memo AST 编译为一个调用 JAX API 的 Python 程序。

当调用一个 memo 模型时,JAX 的即时编译器(JIT)会追踪生成的 Python 程序,并将其编译为原生代码,从而实现高效的执行性能。

4 应用与评估

本节首先介绍来自 RR(递归理性)范式的四个经典模型,并将它们作为案例研究来评估 memo 的表现(§4.1)。对于每个模型,我们将 memo 的实现与传统概率编程语言(PPL)中由专家编写实现进行比较。我们展示了 memo 在通常情况下不仅运行速度更快 ,而且代码量更少

接下来,我们简要介绍了目前有四个认知科学研究项目正在使用 memo,并强调 memo 是如何以不同方式影响这些正在进行的研究工作的(§4.2)。

最后,我们展示了两个高级使用示例,构建了一些在其他 PPL 中实现起来特别困难的模型,以展示 memo 的强大表达能力(§4.3)。

我们在表1中总结了所有这些模型的实验结果。根据我们在第 1.1.2 节中的分析,我们测量了计算每个模型所有相关条件所需的总时间——这是研究人员日常工作中最相关的负载指标。我们排除了 JIT 编译器预热所需的时间 ,因为这个固定成本远小于模型拟合和交叉验证所带来的可变成本。

除非另有说明,所有模型都在 Apple M2 处理器上进行了计时测试。我们也报告了每个模型所需的代码行数 (Lines of Code, LOC),以大致衡量实现一个模型所需的开发工作量、潜在 bug 的数量,以及 memo 所提供的领域抽象带来的效率提升。

4.1 案例研究

我们从 RR 范式中的四个经典模型开始研究:

  1. 理性沟通

    (Rational Communication)

  2. 双人博弈

    (Two-player Games)

  3. 顺序规划

    (Sequential Planning)

  4. 信念空间推理

    (Belief-space Reasoning)

我们将这些模型现有的实现(来源于教材、文档等)与使用 memo 重新实现的版本进行了对比。通过这些案例,我们展示了 memo 在建模复杂递归理性行为方面的优势。

4.1.1 标量蕴含(Scalar Implicature)

标量蕴含是一种语言现象:当说话者说出一个较弱的陈述时,听者会将其理解为较强的陈述并不成立 [79]。一个经典例子是:“some(一些)”常常意味着“not all(不是所有)”。

例如,如果一个人说:“我今年有些 PLDI 投稿被接收了”,听者通常会推断出“并不是全部都被接收”。尽管从技术上讲,即使所有投稿都被接收 ,说“some”(即数量 > 0)也并非错误。

理性言语行为框架 (RSA)很好地预测了这一现象 [65]。我们设定:

4.1.2 谢林协调博弈(Schelling Coordination Games)

经济学家谢林(Schelling)[116] 研究了一种情境:两个参与者在无法沟通的情况下被迫达成协调 。在一个如今著名的实验中,他让36名学生设想自己需要在纽约市与一个陌生人会面,但事先并未约定具体的时间和地点。大多数人表示他们会尝试在中午时分前往中央车站 (Grand Central Terminal)。这种“谢林点”(Schelling point)是通过递归推理产生的:人们会思考他人最可能去的地方,从而使得某个显著的地点自发地成为共识。

Stuhlmüller 和 Goodman [125] 用一对相互递归的概率程序(“Alice” 和 “Bob”)对这种推理进行了建模。他们研究了一个简化的情境:一座小镇上有两家酒吧,其中一家略微更受欢迎。他们发现,只需几层递归推理,理性智能体就会几乎一致地选择那家更受欢迎的酒吧。

我们在 memo 中实现了这个模型,并将其与 Evans 等人 [47] 教材中的 WebPPL 实现进行了对比(见图7)。为了测试运行时间,我们将场景扩展为 100 家酒吧 ,并考虑了 100 层递归推理

如表1所示:

  • memo 实现仅需 15 行代码

    ,而 WebPPL 实现则有 25 行;

  • 运行速度方面,memo 仅需 5.8 毫秒 ,而 WebPPL 需要 1.1 秒

传统 PPL 中的“作恶者混淆”风险

这个例子很有趣,因为它突显了传统 PPL 中常见的“作恶者混淆”(perpetration confusion)问题(见 §1.1.1)。

在这个特定案例中,教材中的 WebPPL 实现恰好得出了正确的结果,但只要稍作改动,就可能导致错误。例如,如果我们想预测 Alice 对她能成功与 Bob 会面的信心程度,可能会倾向于写一个新的模型 alice',直接推断 a == b 的概率,而不是推断 a 的概率。然而,这样做的结果会错误地返回 100% 的信心 ,仿佛 Bob 总是会选择与 Alice 相同的酒吧——这显然是不合理的乐观假设。

相比之下,memo 强制我们明确表达我们想要计算的内容 ,从而避免此类错误。在 memo 中,我们需要写出:

即:观察者对 Alice 所赋予的“成功会面”概率的期望值。

为什么这里必须使用观察者的期望呢?这是因为在观察者看来,Alice 自己选择哪家酒吧是不确定的(这会影响她的信心)。很容易忽略这一细节,但如果遗漏了外层的 E[...],memo 就会在编译时静态报错。这正是 memo 帮助程序员发现模型中隐藏假设的一种方式。

观察者视角 vs. Alice 的主观视角

值得注意的是,Alice 对自己成功的信心,与观察者对她是否能成功会面的信心是不同的。

如果我们要在 memo 中建模后者,可以定义一个观察者来模拟 Alice 和 Bob 分别选择 ab,然后计算:

有趣的是,我们发现:观察者对 Alice 成功的信心略高于 Alice 自己的信心

这是因为:

  • Alice 在第 d 层建模的是 Bob 在第 d−1 层的行为;

  • 而观察者在第 d 层同时建模了 Alice 和 Bob 都在第 d 层的行为。

因此,观察者拥有更完整的全局视角,从而得出更高的成功概率估计值。

这一现象再次体现了 memo 在建模复杂心智交互时的清晰性、安全性与表达能力。

4.1.3 通过值迭代进行马尔可夫决策过程(MDP)规划

马尔可夫决策过程 (Markov Decision Process, MDP)是用于建模序贯决策 (sequential decision-making)的基础形式化工具 [19]。

抽象地讲,一个 MDP 描述了:

我们可以在 memo 中实现一个 MDP,方法是将智能体建模为对其未来的自己进行递归推理 。这样编译出的数组程序本质上就是一种策略迭代 (policy iteration)的实现。

更具体地说,在 memo 中我们可以自然地表达如下逻辑:

我们将 memo 实现与 WebPPL 中的实现进行了比较,使用了一个迷宫导航 MDP (这是测试 RR 模型的常见设置)。

我们让一个智能体在一个 9×9 的迷宫中进行决策:

  • 状态空间 S :81 个格子;

  • 动作空间 A :四个基本方向(上下左右);

  • 转移函数 T :表示合法移动;

  • 奖励函数 R :奖励智能体尽快到达目标格子。

为了在这个迷宫中规划路径,我们运行了 1000 轮策略迭代 (policy iteration)。

如表1所示:

  • WebPPL 和 memo 的代码行数相近(75 行 vs. 60 行),

  • 但 memo 的运行速度要快 40 倍 (10.3 秒 vs. 0.27 秒)。

当我们扩展到一个 21×21 的网格 时:

  • WebPPL 需要 79 秒 (我们甚至不得不绕过其内部超时限制);

  • 而 memo 只需 6.9 秒

这正是本文中最慢的模型之一,因此也是展示 memo 中GPU 加速价值 的好机会。

在 NVIDIA GeForce RTX 4070 上,同样的计算只需 0.377 秒 ,而且无需编写任何新代码

图8(a-b) 展示了在一个 21×21 的“memo 主题”迷宫中计算出的价值函数(value function),该迷宫有两个可能的目标位置:西北角(NW)和东南角(SE)。正如预期的那样,价值函数对应于从每个状态到目标的最短路径长度。

借助这个模型,我们可以重现计算社会认知领域最具影响力的模型之一:“逆向规划 ”(inverse planning)[13]。它模拟了人们理解彼此行为的能力。

从贝叶斯视角来看,如果我们假设一个智能体是理性地行动 [61, 83, 84],那么我们就可以根据其观察到的行为 来推断其未知的目标

图8(c-d) 展示了在 memo 中进行的此类推断(仅需额外 10 行代码,推理耗时 9.8 毫秒):

  • 例如,如果智能体位于 “O” 下方,并向西移动了一步,则我们会强烈推测它正前往西北角(图c);

  • 但如果智能体位于下方的 “M” 处,并向西移动,则我们对其意图会感到不确定,因为从那个位置出发,无论是去哪个目标(NW 或 SE),向西都是合理的动作(图d)。

这一示例再次展示了 memo 在建模复杂心智推理任务中的强大表达能力、高效性以及可读性。

4.1.4 在信念空间中进行 POMDP 规划

部分可观马尔可夫决策过程 (Partially-Observable Markov Decision Process, POMDP)将 MDP 形式化扩展为描述那些无法确定自身当前状态的智能体 [7, 85]。

在 POMDP 中,智能体常常表现出“探索-利用”(explore-exploit)行为:它们首先采取一些具有认知价值的动作(epistemic actions),获取有信息量的观测,然后利用这些积累的知识高效地追求奖励。

一种常见的 POMDP 规划方法是信念空间规划 (belief-space planning):我们可以将任意 POMDP 转换为一个扩展的标准 MDP,其中:

  • 智能体对真实状态的当前信念(belief)本身被视为一种“信念状态”(belief state);

  • 信念状态之间的转移遵循贝叶斯信念更新规则 [24]。

在这里,我们考虑了一个经典的“Hello World”级 POMDP 示例:“啼哭婴儿问题 ”(crying baby problem)[90]。

在这个问题中:

  • 婴儿可能处于两个状态之一: S={sated,hungry}

  • 父母每一步可以采取三种动作之一: A={feed,sing,ignore}

  • 尽管无法直接知道婴儿的真实状态,但父母每一步都会观察到婴儿是否在啼哭: o∈Ω={crying,quiet}

我们使用了 Kochenderfer 等人 [90] 提供的状态转移函数 T 和观测函数 O 来建模。

从抽象层面来看,这个 POMDP 的最优策略是:

如果父母确信婴儿饿了,就去喂食;否则就忽略婴儿(绝不唱歌! )

通过 memo,我们可以自然地建模这种基于信念的递归推理过程。智能体可以在其心智模型中模拟自己未来的信念状态,并据此做出最优决策。

这一示例进一步展示了 memo 在处理现实世界中不确定性更强、推理更复杂的任务时的强大建模能力与效率优势。

我们采用了一个现有的“啼哭婴儿”POMDP 实现 [46],并使用一个由专家直接在 Julia 中手动编写并优化的 、现成的 POMDP 求解器对其进行求解(即不是通过概率编程语言中的自动推理实现)[6]。我们将该实现与 memo 中的实现进行了对比。

在这两种实现中,我们都对父母关于婴儿是否饥饿的信念进行了离散化处理:将区间 [0, 1] 均匀划分为 50 个可能的信念状态。

如表1所示:

  • Julia 版本

    需要 43 行代码 来实现模型,另外还需要 199 行代码 实现一个 POMDP 求解器;

  • memo 版本

    则只需 30 行代码 实现模型, 61 行代码 实现 POMDP 求解器。(这个求解器的一个通用版本已经作为工具内置在 memo 中。)

两者的性能相当(25ms vs. 26ms)。

图9 展示了 memo 计算出的信念空间价值函数 ,成功复现了 Kochenderfer 等人 [90] 的结果(见其书中的图 F.10)。

这一案例进一步表明,memo 不仅在代码简洁性和建模效率上优于传统方法,还能保持与专业手工优化代码相当的运行性能,展示了其在复杂不确定性建模任务中的实用性和高效性。

4.2 memo 在实际研究中的应用

接下来,我们简要介绍目前有四个真实世界的研究项目 正在使用 memo。我们感谢这些同事允许我们在本文中描述他们正在进行的、尚处于早期阶段的研究工作。

4.2.1 撒谎行为建模(Lie Production)

Yi 等人 [138] 正在开发一种关于“撒谎”的理论。他们设计了一个简单的实验:两个角色(“说谎者”和“审判者”)各自对世界的了解都是不完整的。“说谎者”需要回答“审判者”的问题,并被激励去撒谎——但又必须小心不要被识破。

这是一个非常复杂的模型,因为其中涉及许多不同层次的心理状态及其对心理状态的推理。例如:

  • 说谎者必须思考:审判者可能知道什么?

  • 审判者认为说谎者知道什么?

  • 审判者会根据说谎者的陈述推断出什么?

Yi 等人最初在 Gen PPL [40] 中实现了他们的模型。我们与他们合作,将该模型重新用 memo 实现。

如表1所示:

  • 模型代码从 220 行 缩短到了 50 行

  • 推理速度提升了约 2,000,000 倍

也许更重要的是,memo 让研究人员更容易扩展他们的模型。他们早就想建立一个“识别谎言”的模型(“call B.S.”),通过对自己所构建的“产生谎言”模型进行推理来实现。他们此前预计这类模型需要数周时间才能实现,但在 memo 的帮助下,他们仅用一个小时就完成了实现,并验证了它能够预测他们希望捕捉的现象。

4.2.2 风险性唾液共享(Risky Saliva Sharing)

Hung 等人 [81] 想要建立一个计算模型,用于解释人们如何通过观察他人的行为来直觉地判断彼此的社会关系。例如:

  • 如果看到两个人共用一杯饮料,我们会根据他们是共用一根吸管、各用一根吸管,还是将饮料倒入两个杯子来判断他们关系的亲密程度。

  • 更令人惊讶的是,即使是年幼的儿童也能从“风险性唾液共享”中做出这种推理 [49, 126]。

在他们的模型中:

  • 两个智能体基于亲密程度和风险感知选择是否共享唾液;

  • 第三个智能体则根据观察到的行为来推断这两个智能体之间的亲密程度。

Hung 等人最初是在 R 编程语言中开发这个模型的 [82]。我们协助他们将其改写为 memo 版本。

如表1所示:

  • 代码行数从 125 行 减少到 38 行

  • 推理速度提升了约 3000 倍

这使得统计分析(包括模型拟合和交叉验证)的时间从原来的 6 小时 缩短到了 6 秒钟 ,极大地提升了研究人员迭代和探索新想法的能力,这是以前无法想象的。

4.2.3 同理心式解释(Empathetic Explanation)

Chandra 等人 [30] 关注情绪在医患互动中的作用。例如,医生如何向病人解释其绝症成因,以最小化病人对其人生选择的后悔感。

研究人员假设:人们期望医生在沟通中考虑情绪因素,并为此收集了实证数据来检验这一假设。他们随后结合以下模型构建了一个“同理心解释”模型:

  • 最新的解释模型 [29];

  • 情绪预测模型 [80];

  • 干预模型 [35]。

如表1所示:

  • 该模型用 memo 编写了 120 行代码

  • 每次运行只需 92 微秒

他们使用现成的 Adam 优化器 [42] 来拟合参数。研究人员报告称,在建模因果性和后悔感的反事实推理过程中,memo 帮助他们发现了细微的问题。

4.2.4 抚养行为建模(Caregiving)

Kleiman-Weiner [88] 致力于研究父母如何照顾孩子。他们设想了一个简单模型:父母可以选择让孩子自己解决问题,或者“接管”任务替孩子完成——但这样做的风险是可能会让孩子觉得自己不被信任。

在最初的 Python 实现中,出于可扩展性的考虑,他们只能使用一个高度受限、可解析求解的子模型 来表示孩子的行为。

而 memo 允许他们摆脱这一限制,转而使用更加灵活的孩子行为模型。该模型利用了 memo 内置的 POMDP 库(见 §4.1.4)。

4.3 扩展功能:新类型模型的实现

最后,我们通过两个演示展示 memo 所使得一些以前难以实现、现在变得可行的新类型模型。

4.3.1 字体设计(Font Design)

七段显示器 (Seven-segment display)是一种简单的电子元件,常用于计算器、微波炉等设备上显示字母和数字。它由七个可独立开关的段组成,共有 种可能的显示配置。

在设计这类显示器的“字体”时,设计师必须确保所使用的字体是无歧义的 。例如:

  • 如果设计师将字母 “g” 映射为某个特定的七段组合(比如:点亮某些段),这个图案可能会被误认为是数字 “9”;

  • 在这种情况下,也许更好的做法是用这一组合作为大写字母 “G” 的表示——尽管这又可能与数字 “6” 产生混淆!

我们可以使用 RSA 来为这种显示器设计“字体”,将 U 设为 128 种显示配置(即所有可能的七段组合),并将 R设为需要显示的一组字符(例如英语中不区分大小写的 36 个字母数字字符)。

核心思想是从说话者 (设计师)的角度出发应用 RSA,其中“听者”是一个基于计算机视觉的光学字符识别模型(ℓ = 0 的天真听者)。

由于 memo 建立在 JAX 之上,我们可以利用其现有的深度学习生态系统,将 RSA 与深度学习无缝结合。在这个示例中,我们使用了 Flax [73] 库训练了一个 ResNet [72] 网络,用于对 EMNIST 数据集 [37] 中的字母数字字符进行分类。

一种表示字符 的天真策略是:选择一个显示配置 ,使得当输入该配置渲染出的图像时,神经网络对字符 的激活值最大。然而,这种策略会导致字体产生歧义,如图 10 所示。

如果我们转而使用神经网络来建模 RSA 中的基础“听者”,那么说话者就能有效地区分那些容易混淆的字符对。

这个模型运行效率很高(仅需 124 毫秒),因为 memo 在内部并行化了图像渲染,并批量调用了神经网络。

这种深度集成在传统概率编程语言(PPL)中很难实现,因为在那些系统中,外部函数必须串行调用(甚至可能需要从头重新实现)。虽然这个例子本身相对简单,但认知科学界越来越关注那些能推理其他智能体感知和认知的 RR 模型——这些模型会将计算机视觉、图形引擎、物理模拟器和语言模型纳入推理循环中 [20, 28, 32, 39, 98, 144]。因此,我们需要能够支持高效实现的新型 PPL,而 memo 正是朝着这一方向迈出的重要一步。

4.3.2 “Takeaway” 游戏

人们有时会根据他人思考的程度 来对彼此做出推断。例如:

  • 如果你向朋友提出一个难题,而他立刻给出了正确答案,你可能会推测他其实早就知道答案。

认知科学领域的最新研究试图建模人们如何将思考过程的时间动态 纳入推理之中 [21, 22, 129]。正如我们在 §2.5.2 中所讨论的那样,memo 的设计正是为了支持这类模型的构建。

在这里,我们将展示一个用于进行此类推理的简化模型。

我们考虑一个名为 “1-2 Takeaway” 的游戏,它是尼姆博弈(Nim)的一个简化变种 [63]。规则如下:

  • 桌上有 n 根火柴;

  • 玩家轮流移除 1 或 2 根火柴;

  • 最后拿走最后一根火柴的玩家获胜。

我们可以很容易地在 memo 中建模这个游戏,并为每个 n 恢复出最优策略。

通过归纳法可以很容易看出(也容易证明),最优策略是:每次留下火柴数为 3 的倍数 。如果当前 n 已经是 3 的倍数,则不存在必胜策略。

知道这个“技巧”的玩家即使面对很大的 n,也能轻松计算出最佳动作;而不知道这个技巧的玩家则必须遍历指数级增长的游戏树来寻找答案。

设想这样一个场景:

你看到某人坐下来快速玩了一轮 Takeaway。他数了数桌上的火柴数量 n,然后立刻做出了一个最优动作。那么,他是否知道那个“技巧”呢?我们如何从他的行为中推断这一点?

我们的直觉是这样的:

  • 对于较小的 n (例如 1 或 2),很难判断——因为即便不知道技巧,也很容易算出正确动作;

  • 如果 n 是 3 的倍数(即处于劣势位置),也很难判断——因为即便知道技巧,玩家也只能随机选择动作;

  • 但如果 n≫1 且不是 3 的倍数,我们就倾向于怀疑他知道这个技巧:因为在这么短的时间内完成完整博弈论计算并得出正确动作的可能性很低。

    • 如果我们观察到他连续做出多个正确动作,这种怀疑就会进一步增强。

我们将通过以下假设来建模这一推理过程:

假设该玩家要么是知道技巧的 (cognizant),要么是不知道技巧的 (oblivious):

  • 如果玩家是 知道技巧的 ,他们会 始终采取最优动作

  • 如果玩家是 不知道技巧的 ,他们会基于真实效用值,以 softmax 的方式选择动作;

    • 其中 softmax 的“温度”(temperature)取决于所需的计算量;

    • 所需计算越多,动作选择越“嘈杂”(noisy)。

观察者(observer)的任务就是根据所观察到的动作来推断:该玩家是知道技巧,还是不知道技巧。

关键的一行代码如下(即观察者如何进行推理):

在这里,cost @ π(t) 表示玩家在时间范围 t 内计算其最优策略 π 所需的浮点运算次数(FLOPs),而 α 是一个比例因子,用于将 FLOPs 转换为“心理努力”的单位。

正如图11所示,这种建模方式确实得出了我们预期的结果:

  • 当 n 很小或 n≡0(mod3) 时,观察者对玩家是否知道技巧 不确定

  • 但随着 n→∞ 且 n≡0(mod3) ,观察者越来越倾向于认为玩家是“知道技巧”的。

这类模型在过去实现起来极具挑战性,但在 memo 中,这个模型仅需 32 行代码 ,运行时间仅为 116 毫秒

小结

memo 具有高度表达能力,能够编码以下各类模型:

  • 各种经典的 RR 模型;

  • 最新的现实世界模型;

  • 以及一些在传统概率编程语言中难以实现的新型复杂模型。

与传统 PPL 相比,memo 的代码通常更短、更容易理解和推理。
在我们考察的所有经典模型中,memo 的运行速度都快于传统 PPL 中专家编写的实现(在 §4.1.4 中甚至与手工优化的求解器相当)。
而在最新的现实世界模型中,memo 的加速效果更加显著,比研究人员之前实现的版本快了多个数量级

虽然我们在本文中没有讨论他们的工作,但已有其他研究人员使用 memo 来研究诸如相互凝视 (mutual gaze)、多模态推理 (multimodal inference)和惩罚行为 (punishment)等现象。一所大学(与作者无关联)甚至开设了一门以 memo 为教学语言的社会认知春季课程。

最后,补充材料中还包含了大量额外的示例,包括:

  • 常见模式:如“赋能”(empowerment)[89] 和“期望信息增益”(expected information gain)[114];

  • 经典经济博弈:如“最后通牒博弈”(Ultimatum)[100]、“贝叶斯说服”(Bayesian Persuasion)[86] 和 “三分之二猜谜游戏”(Guess 2/3 of the Mean)[70];

  • 以及各种逻辑谜题:如“蒙提霍尔问题”(Monty Hall)[119]、“谢丽尔生日问题”(Cheryl’s Birthday)[131]、“共餐密码客”(Dining Cryptographers)[34] 和“纽康悖论”(Newcomb’s Paradox)[135]。

这些案例进一步证明了 memo 在建模心智递归推理方面的强大能力与广泛适用性。

5 局限性与未来工作

memo 的大多数局限性来源于其基于数组的后端设计 。由于联合分布被表示为数组,memo 不支持连续随机变量 (除非离散化),也仅支持静态已知选择序列 (即数组轴)的控制流。

对于典型的现实世界 RR 模型来说,这些限制很少成为问题。尽管如此,我们仍对探索支持这些功能的替代(或混合)推理后端感兴趣。

除了改进 memo 本身,我们也希望利用它来研究一些关于语言和心理理论的长期科学问题。例如:

  • 来自聋童使用手语的研究证据表明:学习用于心理理论的语言可能是理解虚假信念的关键 [41, 99, 107, 108, 117]。

  • 那么,分析 memo 所提供的领域抽象是否能帮助我们理解“专门语言”在心理理论中所起的作用?

此外,目前学界热烈讨论的一个问题是:大语言模型 (LLMs)是否具备类似人类的心理理论能力 [127]。如果我们把 memo 视作一种“思维语言”[51],是否能够提升 LLM 的心智建模能力?就像让 LLM 接触代码、定理证明器和概率编程语言可以提升其他类型的推理能力一样 [105, 136, 142]?

6 相关工作补充

我们在第1节中已经讨论了主要的相关工作。除此之外,memo 还与以下两个方向的工作密切相关。

第一类相关工作 是近期对精确离散推理扩展性 的研究。例如:

  • Dice [78]

  • SPPL [115]

  • Genfer [141]

其中一些系统甚至将连续变量“位爆破”(bit-blast)为离散变量(如 HyBit [57])。
FSPN [124] 和 PERPL [36] 支持在存在递归的情况下进行精确推理。

这些语言通过巧妙地分解问题结构,试图扩展到任意概率程序。而 memo 的一个关键洞察在于:在现实世界的 RR 模型这一特定领域中,我们可以通过在现代并行硬件上高效运行相对简单的算法来实现扩展性 (见 §1.2.2)。

第二类相关工作 认知逻辑 (epistemic logics)[14, 48, 71, 75, 132] 和信念-欲望-意图逻辑 (Belief-Desire-Intention logics)[59, 60, 137]。这些形式化方法长期以来被应用于经济学 [9–11] 和计算机安全 [18, 26] 等领域,用于建模智能体的知识与信念。

memo 受这些逻辑的启发,但建立在用概率建模信念强度的贝叶斯传统之上。

  • 认知逻辑主要用于 证明关于知识本质的一般性陈述

  • 而 memo 则专注于 计算具体情境下人类行为的数值预测

最后,为了透明起见,我们指出:本文的两页扩展摘要正同时提交给一个非归档性质的工作坊进行评审 [4, 匿...