在没有人类示范的情况下解决奥林匹克几何问题

本仓库包含重现 DDAR 和 AlphaGeometry 这两个几何定理证明器所需的代码，它们在 Nature 2024 论文中介绍：

<center>"在没有人类示范的情况下解决奥林匹克几何问题"。</center>

</br> <center> <img alt="fig1" width="800px" src="https://yellow-cdn.veclightyear.com/835a84d5/cc576b12-89ee-4541-b987-85fc79a1a0e6.svg"> </center>

依赖项

对于以下介绍的说明，我们使用 Python 3.10.9，以及 requirements.txt 中列出的具体版本号的依赖项。

我们的代码依赖于 meliad，它不是 pip 的注册包。请参阅下面的说明了解如何手动安装 meliad。

请注意，即使没有 meliad 和 sentencepiece 依赖项，仍然可以运行 DDAR 求解器。

运行说明

可以通过以下方式一次性运行本 README.md 中的所有说明：

bash run.sh

下面，我们将逐步解释这些说明。

安装依赖项，下载权重和词汇表

安装在虚拟环境中完成：

virtualenv -p python3 .
source ./bin/activate
pip install --require-hashes -r requirements.txt

下载权重和词汇表：

bash download.sh
DATA=ag_ckpt_vocab

最后，单独安装 meliad，因为它未在 pip 中注册：

MELIAD_PATH=meliad_lib/meliad
mkdir -p $MELIAD_PATH
git clone https://github.com/google-research/meliad $MELIAD_PATH
export PYTHONPATH=$PYTHONPATH:$MELIAD_PATH

设置通用标志

在运行 Python 脚本之前，让我们先准备一些常用的标志。符号引擎需要定义和推理规则来操作。这些定义和规则在两个文本文件 defs.txt 和 rules.txt 中提供。

DDAR_ARGS=(
  --defs_file=$(pwd)/defs.txt \
  --rules_file=$(pwd)/rules.txt \
);

接下来，我们定义与证明搜索相关的标志。为了重现以下简单示例，我们使用轻量级的证明搜索参数值：

BATCH_SIZE=2
BEAM_SIZE=2
DEPTH=2

SEARCH_ARGS=(
  --beam_size=$BEAM_SIZE
  --search_depth=$DEPTH
)

注意：我们论文中的结果可以通过设置 BATCH_SIZE=32、BEAM_SIZE=512、DEPTH=16 来获得，如方法部分所述。为了保持在国际数学奥林匹克的时间限制内，需要 4 个 V100-GPU 和 250 个 CPU 工作线程，如扩展数据 - 图 1 所示。请注意，我们还去除了其他内存/速度优化，以避免内部依赖并提高代码清晰度。

假设下载的检查点和词汇表放在 DATA 中，安装的 meliad 源代码位于 MELIAD_PATH。我们使用 gin 库来管理模型配置，遵循 meliad 惯例。现在我们定义与语言模型相关的标志：

LM_ARGS=(
  --ckpt_path=$DATA \
  --vocab_path=$DATA/geometry.757.model
  --gin_search_paths=$MELIAD_PATH/transformer/configs,$(pwd) \
  --gin_file=base_htrans.gin \
  --gin_file=size/medium_150M.gin \
  --gin_file=options/positions_t5.gin \
  --gin_file=options/lr_cosine_decay.gin \
  --gin_file=options/seq_1024_nocache.gin \
  --gin_file=geometry_150M_generate.gin \
  --gin_param=DecoderOnlyLanguageModelGenerate.output_token_losses=True \
  --gin_param=TransformerTaskConfig.batch_size=$BATCH_SIZE \
  --gin_param=TransformerTaskConfig.sequence_length=128 \
  --gin_param=Trainer.restore_state_variables=False
);

提示：请注意，即使不定义 SEARCH_ARGS 和 LM_ARGS，您仍然可以运行 DDAR 求解器。在这种情况下，只需在 alphageometry.py 中禁用 lm_inference 模块的导入即可。

运行 DDAR

该脚本通过从文本文件读取问题列表来加载问题，并根据其名称解决列表中的特定问题。我们通过标志 --problems_file 和 --problem_name 传递这两条信息。我们使用 --mode=ddar 来表示我们要使用 DDAR 求解器。

以下我们展示了这个求解器解决 IMO 2000 P1：

python -m alphageometry \
--alsologtostderr \
--problems_file=$(pwd)/imo_ag_30.txt \
--problem_name=translated_imo_2000_p1 \
--mode=ddar \
"${DDAR_ARGS[@]}"

预期输出如下

graph.py:468] translated_imo_2000_p1
graph.py:469] a b = segment a b; g1 = on_tline g1 a a b; g2 = on_tline g2 b b a; m = on_circle m g1 a, on_circle m g2 b; n = on_circle n g1 a, on_circle n g2 b; c = on_pline c m a b, on_circle c g1 a; d = on_pline d m a b, on_circle d g2 b; e = on_line e a c, on_line e b d; p = on_line p a n, on_line p c d; q = on_line q b n, on_line q c d ? cong e p e q
ddar.py:41] Depth 1/1000 time = 1.7772269248962402
ddar.py:41] Depth 2/1000 time = 5.63526177406311
ddar.py:41] Depth 3/1000 time = 6.883412837982178
ddar.py:41] Depth 4/1000 time = 10.275688409805298
ddar.py:41] Depth 5/1000 time = 12.048273086547852
alphageometry.py:190]
==========================
 * 从定理前提：
A B G1 G2 M N C D E P Q : 点
AG_1 ⟂ AB [00]
BA ⟂ G_2B [01]
G_2M = G_2B [02]
G_1M = G_1A [03]

...
[省略日志]
...

036. ∠QEB = ∠(QP-EA) [46] & ∠(BE-QP) = ∠AEP [55] ⇒  ∠EQP = ∠QPE [56]
037. ∠PQE = ∠EPQ [56] ⇒  EP = EQ

==========================

输出首先包括它使用的相关前提列表，然后是逐步构建证明的证明步骤。所有谓词都编了号，以追踪它们是如何从前提推导出来的，并显示证明是完全合理的。

提示：另外传递标志 --out_file=path/to/output/text/file.txt 将把证明写入文本文件。

在 imo_ag_30.txt 中的所有问题上运行将得到其中 14 个问题的解决方案，如我们论文中的表 1 所报告的。

运行 AlphaGeometry：

作为一个简单的例子，我们从examples.txt文件中加载--problem_name=orthocenter。这次，我们传递--mode=alphageometry来使用AlphaGeometry求解器，并传递SEARCH_ARGS和LM_ARGS标志。

python -m alphageometry \
--alsologtostderr \
--problems_file=$(pwd)/examples.txt \
--problem_name=orthocenter \
--mode=alphageometry \
"${DDAR_ARGS[@]}" \
"${SEARCH_ARGS[@]}" \
"${LM_ARGS[@]}"

预期输出如下：

...
[省略日志]
...
training_loop.py:725] 总参数：152072288
training_loop.py:739] 总状态大小：0
training_loop.py:492] 训练循环：为beam_search模式创建任务

graph.py:468] orthocenter
graph.py:469] a b c = triangle a b c; d = on_tline d b a c, on_tline d c a b ? perp a d b c
ddar.py:41] 深度 1/1000 时间 = 0.009987592697143555 分支 = 4
ddar.py:41] 深度 2/1000 时间 = 0.00672602653503418 分支 = 0
alphageometry.py:221] DD+AR未能解决问题。
alphageometry.py:457] 深度 0。有1个节点待展开：
alphageometry.py:460] {S} a : ; b : ; c : ; d : T a b c d 00 T a c b d 01 ? T a d b c {F1} x00
alphageometry.py:465] 从{S} a : ; b : ; c : ; d : T a b c d 00 T a c b d 01 ? T a d b c {F1} x00解码
...
[省略日志]
...
alphageometry.py:470] LM输出（得分=-1.102287）："e : C a c e 02 C b d e 03 ;"
alphageometry.py:471] 翻译："e = on_line e a c, on_line e b d"

alphageometry.py:480] 求解："a b c = triangle a b c; d = on_tline d b a c, on_tline d c a b; e = on_line e a c, on_line e b d ? perp a d b c"
graph.py:468]
graph.py:469] a b c = triangle a b c; d = on_tline d b a c, on_tline d c a b; e = on_line e a c, on_line e b d ? perp a d b c
ddar.py:41] 深度 1/1000 时间 = 0.021120786666870117
ddar.py:41] 深度 2/1000 时间 = 0.033370018005371094
ddar.py:41] 深度 3/1000 时间 = 0.04297471046447754
alphageometry.py:140]
==========================
 * 从定理前提：
A B C D : 点
BD ⟂ AC [00]
CD ⟂ AB [01]

 * 辅助作图：
E : 点
E,B,D共线 [02]
E,C,A共线 [03]

 * 证明步骤：
001. E,B,D共线 [02] & E,C,A共线 [03] & BD ⟂ AC [00] ⇒  ∠BEA = ∠CED [04]
002. E,B,D共线 [02] & E,C,A共线 [03] & BD ⟂ AC [00] ⇒  ∠BEC = ∠AED [05]
003. A,E,C共线 [03] & E,B,D共线 [02] & AC ⟂ BD [00] ⇒  EC ⟂ EB [06]
004. EC ⟂ EB [06] & CD ⟂ AB [01] ⇒  ∠(EC-BA) = ∠(EB-CD) [07]
005. E,C,A共线 [03] & E,B,D共线 [02] & ∠(EC-BA) = ∠(EB-CD) [07] ⇒  ∠BAE = ∠CDE [08]
006. ∠BEA = ∠CED [04] & ∠BAE = ∠CDE [08] (相似三角形)⇒  EB:EC = EA:ED [09]
007. EB:EC = EA:ED [09] & ∠BEC = ∠AED [05] (相似三角形)⇒  ∠BCE = ∠ADE [10]
008. EB:EC = EA:ED [09] & ∠BEC = ∠AED [05] (相似三角形)⇒  ∠EBC = ∠EAD [11]
009. ∠BCE = ∠ADE [10] & E,C,A共线 [03] & E,B,D共线 [02] & ∠EBC = ∠EAD [11] ⇒  AD ⟂ BC
==========================

alphageometry.py:505] 已解决。

注意：点H自动重命名为D，因为LM是在合成问题上训练的，其中点按字母顺序命名，所以在测试时也期望相同的命名。

注意：在这个AlphaGeometry的实现中，我们移除了所有依赖内部基础设施的优化，例如：在多GPU上的并行模型推理、在多CPU上的并行DDAR、 LM和DDAR的并行执行、跨不同问题的共享CPU工作池等。我们还为了代码清晰度，移除了一些内存/速度优化和代码抽象。

从输出可以看出，最初DDAR未能解决问题。 LM提出了两个辅助作图（因为BATCH_SIZE=2）：

e = eqdistance e c a b, eqdistance e b a c，即构造E为圆（中心=C，半径=AB）和圆（中心=B，半径=AC）的交点。这个作图的得分为-1.186。
e = on_line e a c, on_line e b d，即 E是AC和BD的交点。这个作图的得分（-1.102287）高于前一个。

由于第二个作图得分更高，DDAR首先尝试了第二个作图并立即找到了解决方案。因此证明搜索终止，没有第二次迭代。

结果

在尝试复现我们论文中的AlphaGeometry数据之前，请确保通过准备好的测试套件中的所有测试：

bash run_tests.sh

注意：Issues#14报告说，尽管顶部束解码仍然相同，但LM对不同用户给出的分数不同。

然后，传递相应的--problem_file（列）和--mode（行）值，并迭代所有问题以获得以下结果：

<b>解决问题的数量：</b>

	`imo_ag_30.txt`	`jgex_ag_231.txt`
`ddar`	14	198
`alphageometry`	25	228

</center>

源代码描述

该仓库中的文件包括用于运行求解器的Python模块/脚本以及脚本执行所需的资源文件。我们在下面列出了每个文件及其描述。

文件名	描述
`geometry.py`	实现证明状态图中的节点（点、线、圆等）。
`numericals.py`	实现动态几何环境中的数值引擎。
`graph_utils.py`	实现证明状态图的实用工具。
`graph.py`	实现证明状态图。
`problem.py`	实现表示问题前提、结论、DAG节点的类。
`dd.py`	实现DD及其回溯。
`ar.py`	实现AR及其回溯。
`trace_back.py`	实现递归回溯和依赖差异算法。
`ddar.py`	实现DD+AR的组合。
`beam_search.py`	在JAX中实现语言模型的束搜索解码。
`models.py`	实现transformer模型。
`transformer_layer.py`	实现transformer层。
`decoder_stack.py`	实现transformer解码器堆栈。
`lm_inference.py`	实现训练好的语言模型的接口以进行解码。
`alphageometry.py`	主脚本，用于加载问题，调用DD+AR或AlphaGeometry求解器，并打印解决方案。
`pretty.py`	美化求解器输出的解决方案格式。
`*_test.py`	对应模块的测试。
`download.sh`	下载模型检查点和语言模型的脚本。
`run.sh`	执行README中指令的脚本。
`run_tests.sh`	执行测试套件的脚本。

资源文件：

资源文件名	描述
`defs.txt`	不同几何构造动作的定义。
`rules.txt`	DD的推导规则。
`geometry_150M_generate.gin`	meliad中实现的语言模型的Gin配置。
`imo_ag_30.txt`	IMO-AG-30中的问题。
`jgex_ag_231.txt`	JGEX-AG-231中的问题。

引用本工作

@Article{AlphaGeometryTrinh2024,
  author  = {Trinh, Trieu and Wu, Yuhuai and Le, Quoc and He, He and Luong, Thang},
  journal = {Nature},
  title   = {Solving Olympiad Geometry without Human Demonstrations},
  year    = {2024},
  doi     = {10.1038/s41586-023-06747-5}
}

致谢

本研究是谷歌大脑团队（现为谷歌DeepMind）和纽约大学计算机科学系的合作成果。我们感谢Rif A. Saurous、Denny Zhou、Christian Szegedy、Delesley Hutchins、Thomas Kipf、Hieu Pham、Petar Veličković、Debidatta Dwibedi、Kyunghyun Cho、Lerrel Pinto、Alfredo Canziani、Thomas Wies、He He的研究团队、Evan Chen（美国IMO队教练）、Mirek Olsak、Patrik Bak以及《自然》杂志的三位审稿人的帮助和支持。

AlphaGeometry的代码与以下独立库和包进行通信和/或引用：