consensus-paper-code

Imported from yixiaosz/consensus-paper-code (AGENTS.md). Install upstream with npx skills add yixiaosz/consensus-paper-code. Copyright stays with the author.

AGENTS.md

本文件供 AI 编程助手阅读。当前仓库是论文 The Consensus Game: Language Model Generation via Equilibrium Search 对应算法的极简参考实现，项目本身没有测试、没有部署流程。下面的所有信息均来自仓库内实际文件，不做推测。

项目概述

• 项目名称：consensus-paper-code（pyproject.toml）
• 版本：0.1.0
• 用途：用 NumPy 实现论文中的 piKL 博弈均衡优化算法。核心函数 equilibrium_ranking 接收生成器概率 P_gen 和判别器概率 P_disc，通过迭代收敛得到精炼后的策略 pi_G* 和 pi_D*。
• 关键文件：
  • simple_equilibrium_2x2.py：最简 piKL 算法示例，固定 2×2 矩阵，无绘图，只打印收敛后的策略。
  • simple_equilibrium_2x4.py：piKL 算法与固定 2×4 矩阵示例，包含 equilibrium_ranking()、plot_history()、plot_start_vs_end() 与 main()。
  • demo_intuition_20260801.py：算法直觉演示脚本。从 simple_equilibrium_2x4.py import equilibrium_ranking，用手写的 2×4 矛盾数据（生成器偏爱 A、判别器认为 A 错误）生成 5 张教学图：双视角一致性对比、诱导联合分布、去噪指标曲线（一致率/贝叶斯残差/KL 漂移）、ER-G/ER-D 答案排序、λ 权衡曲线。
  • ollama_probs.py：独立运行的 Ollama 对接版。用 Ollama /api/generate 的 logprobs 接口从本地模型读出 P_gen/P_disc（候选答案 y 压缩成单 token 字母 A/B/C/D，判别 v 压缩成 Yes/No，每次请求只生成 1 个 token），随后自带一份与 simple_equilibrium_2x4.py 相同的算法与绘图代码完成求解；不 import simple_equilibrium_2x4.py。模型由文件顶部的 MODEL 常量指定，当前为 llama2:7b。
  • MMLU_baseline_test.py：MMLU 分层评测脚本（每个 subject 取前 N 题，默认 N=1 共 57 题，可用命令行参数覆盖），复用 ollama_probs.py 的 build_p_gen/build_p_disc/equilibrium_ranking，输出 G / D / ER-G / ER-D 四种方法的 0/1 得分与平均 accuracy，导出 MMLU_consensus_时间戳.txt（顶部为四方法最终平均，正文含每题预测与迭代 0/1000/2000/3000/4000/4999 的 pi_G/pi_D/Q_G/Q_D 快照）；不生成图片。
  • MMLU_baseline_test_MI_SC.py：在 MMLU_baseline_test.py 基础上扩展为六种方法 G / MI / SC / D / ER-G / ER-D 的 MMLU 评测；MI（Mutual Information）为 P_gen[correct,y] * P_disc[y,correct]，SC（Self-Contrastive）为初始 pi_G1[correct,y]；输出与导出格式同 MMLU_baseline_test.py。
  • pyproject.toml：项目元数据与依赖声明。
  • uv.lock：由 uv 生成的锁定文件，确保依赖可复现。
  • README.md：双语项目说明（英文正文+中文摘要），末尾有纯中文附录，解释"为什么 Ollama 版要查询 6 次而非 2 次"以及"如何理解 demo_intuition_20260801.py 的 5 张图"。
  • .python-version：固定 Python 版本为 3.12。
  • CONSENSUS-GAME-paper.pdf：原始论文 PDF。

技术栈

• 语言：Python 3.12+
• 包管理器：uv（存在 uv.lock，.venv 已由 uv 创建）
• 核心依赖：
  • numpy >= 2.5.1
  • matplotlib >= 3.11.1（保存迭代收敛曲线图；脚本使用无界面的 Agg 后端，只写 PNG 不弹窗）
  • datasets >= 5.0.1（MMLU 评测脚本从 Hugging Face 加载 cais/mmlu）
• 构建后端：pyproject.toml 中未显式声明 [build-system]，uv 会按 PEP 621 的 [project] 元数据管理环境。

代码组织

consensus-paper-code/
├── simple_equilibrium_2x2.py    # 最简 piKL：固定 2×2 矩阵，无图
├── simple_equilibrium_2x4.py    # piKL 算法 + 固定 2×4 矩阵 + 绘图
├── demo_intuition_20260801.py   # 算法直觉演示（import simple_equilibrium_2x4，5 张教学图 + 手写矛盾数据）
├── ollama_probs.py              # Ollama logprobs 对接版（独立运行，算法部分复制自 simple_equilibrium_2x4.py）
├── MMLU_baseline_test.py        # MMLU 四方法基线评测（G/D/ER-G/ER-D）
├── MMLU_baseline_test_MI_SC.py  # MMLU 六方法基线评测（+ MI / SC）
├── pyproject.toml               # 项目配置与依赖
├── uv.lock                      # uv 锁定文件
├── .python-version              # 3.12
├── .venv/                       # uv 创建的虚拟环境
├── README.md                    # 双语说明 + 纯中文附录
└── CONSENSUS-GAME-paper.pdf     # 论文原文

• 算法逻辑目前分散在 simple_equilibrium_2x4.py 与 ollama_probs.py 两份几乎相同的副本中；demo_intuition_20260801.py 与 MMLU 脚本分别 import 前者/后者。
• simple_equilibrium_2x4.py 内有四个顶层对象：
  • equilibrium_ranking(P_gen, P_disc, T=5000, lam_G=0.1, lam_D=0.1, eta_G=0.1, eta_D=0.1)
  • plot_history(history)：把迭代历史按变量各保存为一张 PNG。每条曲线在迭代索引 0, 1000, 2000, 3000, 4000, 4999 处标注对应的 y 值（三位小数），并以 dpi=200 导出。
  • plot_start_vs_end(pi_G, pi_D, pi_G1, pi_D1)：保存开始/结束时 pi_G、pi_D 的 2×2 面板矩阵对比图（时间戳_start_vs_end.png），热力图格内标注三位小数值，矩阵顶部/左侧标注列/行含义（y 用 A、B、C、D 表示，v 用 correct/incorrect 表示），并以 dpi=200 导出。
  • main()：用固定 2×4 矩阵演示算法输出。

运行时架构

• 这是一个本地命令行脚本/研究原型，无网络服务、无 Web 框架、无持久化存储。
• simple_equilibrium_2x4.py 执行流程：
  1. main() 启动时提示用户选择输入数据：选 1 使用 Gen 与 Disc 意见相同的数据，选 2 使用意见不同的数据（选项 y 数量为 4：P_gen 形状 (2, 4)，P_disc 形状 (4, 2)）。
  2. 读取输入矩阵 P_gen（形状 (2, n)，对应 P_LM(y|x,v)）和 P_disc（形状 (n, 2)，对应 P_LM(v|x,y)）。
  3. 分别沿 axis=0 做初始归一化，得到 pi_G1 与 pi_D1。
  4. 在主循环中累计对手历史平均策略，计算 Q 值，再按 piKL 公式做 softmax 更新。
  5. 每次迭代结束后，把 pi_G、pi_D、Q_G、Q_D、avg_G、avg_D 的快照存入 history（键为变量名，值形状为 (T, rows, cols)）。
  6. 返回收敛后的 pi_G*、pi_D*、初始策略及完整 history（共五项）。
  7. main() 把结果按三位小数打印，并调用 plot_history() 为 history 中每个变量保存一张收敛曲线图、调用 plot_start_vs_end() 保存开始/结束矩阵对比图；文件名格式为 YYYYMMDD_HHMMSS_<变量名>.png（结束时的系统时间），曲线图 x 轴为迭代次数 t，y 轴为矩阵每个动作（元素）的取值，且在迭代 0, 1000, 2000, 3000, 4000, 4999 处标注每个动作的取值。所有导出图片使用 dpi=200 以提高清晰度。
• simple_equilibrium_2x2.py 执行流程更简单：硬编码 2×2 矩阵，直接调用 equilibrium_ranking()（无历史记录、无绘图），只打印 pi_G / pi_D / pi_G1 / pi_D1。
• ollama_probs.py 的执行流程：先通过 Ollama（需本地已启动服务并存在 MODEL 常量指定的模型，当前 llama2:7b）发出 2 + n 次请求（2 次 generator、n 次 discriminator），从首个生成 token 的 top_logprobs 中读出候选字母/Yes/No 的概率并归一化，得到 P_gen、P_disc；之后走与 simple_equilibrium_2x4.py 第 3–7 步完全相同的流程。题目与候选答案在文件顶部的 QUESTION/OPTIONS 常量中修改。
• 注释里混用中文解释数学直觉，变量名保留论文英文符号（P_gen、P_disc、pi_G、Q_G 等）。

构建与运行命令

推荐使用 uv（项目已配置虚拟环境）：

# 安装/同步依赖（根据 uv.lock）
uv sync

# 最简 2x2 示例
uv run python simple_equilibrium_2x2.py

# 固定矩阵 demo（启动后输入 1 或 2 选择数据）
uv run python simple_equilibrium_2x4.py

# 运行 Ollama 对接版（需本地 Ollama 服务已启动且已导入 MODEL 指定的模型）
uv run python ollama_probs.py

# MMLU 分层评测（参数为每个 subject 的题数；默认 1 即 57 题，5 即 285 题）
uv run python MMLU_baseline_test.py 5
uv run python MMLU_baseline_test_MI_SC.py 5

也可以直接调用已存在的虚拟环境：

.venv/bin/python simple_equilibrium_2x4.py

simple_equilibrium_2x4.py 示例输出大致如下（选 1 即意见相同的数据；浮点数统一保留三位小数，运行结束还会在当前目录生成 6 张 时间戳_变量名.png 曲线图和 1 张 时间戳_start_vs_end.png 开始/结束矩阵对比图）：

选择输入数据 (1: Gen和Disc意见相同, 2: Gen和Disc意见不同): 1
pi_G1 =
[[0.875 0.200 0.250 0.333]
 [0.125 0.800 0.750 0.667]]
pi_D1 =
[[0.437 0.125]
 [0.125 0.333]
 [0.187 0.292]
 [0.250 0.250]]
pi_G =
[[0.969 0.002 0.005 0.024]
 [0.001 0.596 0.320 0.083]]
pi_D =
[[0.997 0.003]
 [0.024 0.976]
 [0.131 0.869]
 [0.417 0.583]]
saved 20260722_203348_pi_G.png
saved 20260722_203348_pi_D.png
...

测试说明

• 当前没有任何测试文件或测试框架配置（pytest、unittest 均未在 pyproject.toml 的依赖或开发依赖中声明）。
• 验证正确性的方式是运行 simple_equilibrium_2x4.py / simple_equilibrium_2x2.py 并核对输出是否达到“双方几乎确定共识”的预期（pi_G/pi_D 对角线接近 1）。
• 如果要添加测试，建议：
  1. 在 pyproject.toml 的 [dependency-groups] 或 [project.optional-dependencies] 中加入 pytest。
  2. 创建 tests/test_equilibrium.py，用论文中给出的 2×2 示例作为回归用例。

代码风格指南

• 保持中文注释风格：关键数学步骤用中文解释，英文保留算法符号。
• 变量命名沿用论文约定：
  • 大写 P_* 表示原始语言模型概率。
  • 小写 pi_* 表示策略/policy。
  • lam_* / eta_* 分别对应 piKL 正则系数与学习率。
• 使用 np.* 向量化实现，避免手写循环。
• 数值稳定性：每轮更新后先 lG -= lG.max(axis=1, keepdims=True) 再做 exp，防止溢出。

安全与注意事项

• simple_equilibrium_2x2.py / simple_equilibrium_2x4.py / demo_intuition_20260801.py 无外部网络请求、无密钥/令牌/环境变量读取，也不处理用户输入文件。
• ollama_probs.py 会请求本地 Ollama 服务（默认 http://localhost:11434），需要服务已启动且已导入 MODEL 指定的模型（当前 llama2:7b）；除此之外无其他网络访问。MMLU_baseline_test.py 与 MMLU_baseline_test_MI_SC.py 额外会从 Hugging Face 下载 cais/mmlu 数据集。
• 输入 P_gen、P_disc 需要调用方保证是非负概率矩阵；当前实现未对输入做校验。
• T=5000、lam=0.1、eta=0.1 是论文第 3 节推荐值；修改超参数会改变收敛行为。
• 当前实现是纯研究原型，不建议直接用于生产环境做大规模语言模型推理。

常见修改点

• 若要接入真实 LM 的 P_gen/P_disc，替换 simple_equilibrium_2x4.py 或 ollama_probs.py 中的硬编码矩阵/问题即可。
• 若需支持更高维状态空间，保持 P_gen 形状为 (2, n)、P_disc 形状为 (n, 2) 的约定不变。
• 若需可视化收敛过程，equilibrium_ranking 已返回 history（含 6 个变量每次迭代的快照），plot_history() 会据此保存 PNG；改动返回结构时需同步修改 main() 的解包。
• 修改模型时只需改 ollama_probs.py 顶部的 MODEL 常量；MMLU 脚本 import 该常量，会一并生效。
• 若修改了 simple_equilibrium_2x4.py 中的算法，需同步更新 ollama_probs.py 中的同名副本以保持行为一致。