ESC
开源 2 分钟阅读

jevlike:训练一个在动态文本选项列表中做选择的小模型

jevlike 是一个独立开源入门模型,接收一段文本和 N 个选项,一次前向计算即为每个选项输出概率,而非逐词生成,输入输出形式与 TypeSafe 未公开设计的商业模型 Jev 相同。项目提供 Doom 与国际象棋演示、合成数据训练流程,并支持冻结的 Hugging Face 预训练编码器。实验显示八选项下单次推理比小型解码器快约 100 倍。

来源:GitHub

Jevlike

训练一个小模型,在动态变化的文本选项列表中做出选择。

Jev 类模型接收一段文本和包含 N 个文本选项的列表,并为每个选项返回一个概率。它通过一次前向计算完成,而不是逐词写出答案。Jev 是 TypeSafe 面向此类任务的商业模型,TypeSafe 尚未公开其设计。本仓库是一个具有相同输入输出形式的独立入门模型。

演示

同一个选项注意力头还能根据图像块为手柄按键打分。这段十秒视频拼接了两个精选的五秒窗口:先是使用 Doom 七个按键的 deadly_corridor 实时战斗,随后是国际象棋控制器用五个按键走位并落子。图示展示了每次决策所用的张量。Doom 窗口来自随附的联合检查点,其在十个记录的回合中平均击杀 0.60、奖励 -97.50。国际象棋窗口来自更强的纯象棋检查点,在 50 局采样对局中对随机走子者取得 4 胜 46 平 0 负,但对 Stockfish level 0 为 0 胜 2 平 48 负。这些窗口是按活跃程度挑选的,并不代表典型对局或能力水平的声明。

安装游戏可选依赖,并使用已发布的联合检查点录制一段全新的 640×480 Doom 轨迹:

uv pip install -e '.[games]'
python examples/doom/play.py examples/checkpoints/joint-imitation.pt --episodes 10 --game-seconds 35.3 --device cpu --capture-resolution 640x480 --output runs/doom.mp4 --trace runs/doom-trace.json

以相同的视觉布局渲染该轨迹。由于作者自有版权的配乐源文件不在本仓库中,生成的影片是无声的。

(cd examples/film && npm install && npx playwright install chromium)
examples/film/make-film.sh runs/doom-trace.json runs/doom-film.mp4 10

本次发布包含 Doom 示例、国际象棋示例、单游戏检查点以及共享的 12 选项检查点。两个游戏都从 jevlike.vision 导入视觉评分器;两个示例中都没有第二份模型副本。

架构

每个选项变为一个查询向量,即表示其文本的一小串数字。该查询为上下文 token 分配注意力权重,这些权重为该选项生成一个上下文向量。一个共享的点积将每个"选项-上下文"对转化为一个分数。softmax(将分数转换为总和为 1 的概率)在所有选项之间运行。

每个选项查询上下文,获得注意力上下文向量,并转化为一个概率。

默认编码器从零开始学习字节嵌入。编码器是将文本转换为向量的部分。可选的 Hugging Face 路径使用冻结的预训练编码器,其现有权重保持固定,而小型评分器在其上学习。

数据格式

每行使用一个 JSON 对象:

{"context":"The customer needs a refund.","options":["refund","sales","technical support"],"label":0}

label 是正确选项的从零开始的索引。每行的选项数量可以不同,最少两个。

快速开始

在仓库根目录运行以下命令。它们会创建本地合成数据、在其上训练、评估保存的模型并为一个新菜单打分。

uv venv
source .venv/bin/activate
uv pip install -e '.[dev]'

jevlike-data synthetic --output data/synthetic
jevlike-train data/synthetic/train.jsonl \
  --validation data/synthetic/validation.jsonl \
  --output runs/synthetic.pt
jevlike-eval runs/synthetic.pt data/synthetic/test.jsonl
jevlike-predict runs/synthetic.pt \
  --context "Choose the exact badge amber badger. Badge: amber badger." \
  --option "azure crane" \
  --option "amber badger" \
  --option "gold heron"

评估会输出 top-1 准确率,即首选正确的比例。Top-3 准确率是正确答案位于三个最高分数之中的比例。期望校准误差将置信度与实际观测准确率进行比较。该命令还会输出一个打乱上下文的对照组,即把每个菜单与错误的上下文配对。有用的模型应当优于该对照。

使用自己的数据

  1. 按上述格式导出训练、验证和测试 JSONL 文件。
  2. 在每一行中保留模型预测时会看到的全部选项。
  3. 将相关记录划分到同一侧。例如,把同一客户或同一目标页面的所有记录放在同一个划分中。这可以防止近似重复样本泄漏进测试集。
  4. 使用你的训练和验证文件运行 jevlike-train。
  5. 在留出的测试文件上运行一次 jevlike-eval。留出指该文件从未用于训练或模型选择。

默认字节编码器将上下文截断为 192 字节,每个选项截断为 32 字节。当你的文本需要更多空间时,可调大 --context-tokens 或 --option-tokens。通过 --device,训练支持 CPU、Mac GPU 的 Apple MPS 以及 NVIDIA GPU 的 CUDA。

使用冻结的预训练编码器

安装可选依赖,并指定任意兼容的 Hugging Face 编码器:

uv pip install -e '.[transformers]'
jevlike-train data/synthetic/train.jsonl \
  --validation data/synthetic/validation.jsonl \
  --output runs/qwen-head.pt \
  --encoder hf \
  --hf-model Qwen/Qwen2.5-0.5B \
  --rank 256 \
  --batch-size 8

检查点保存了训练好的评分器头部和编码器名称,并不复制冻结的编码器权重。因此加载检查点需要能访问同一个 Hugging Face 模型。

--rank 设置小型评分器头部的宽度。更宽的头部拥有更多可训练权重并占用更多内存。

Wikispeedia 示例

scripts/get_wikispeedia.sh 下载公开的 SNAP 存档并构建"下一次点击"JSONL 文件。数据不会存放在本仓库内。

scripts/get_wikispeedia.sh
jevlike-train data/wikispeedia/jsonl/train.jsonl \
  --validation data/wikispeedia/jsonl/validation.jsonl \
  --output runs/wikispeedia.pt

引用请注明 Robert West 和 Jure Leskovec 的 Human Wayfinding in Information Networks(WWW 2012)。源数据条款请查阅 SNAP 数据集页面。

预期效果

在催生本入门项目的实验中,单次前向评分器在合成菜单上达到约 98% 的准确率。在目标不相交的 Wikispeedia"下一次点击"数据上,冻结的 Qwen2.5-0.5B 编码器加评分器达到 26%,而打乱上下文和随机编码器对照约为 8%。一个在 40,000 次点击上从零训练的小模型达到 29%。在八个选项时,单次前向比被迫写出 400 个 token 的小型解码器快约 100 倍。

这些数字描述的是本地实验,而非本次快速开始运行的结果。我们并未证明与 Jev 同等的质量,也未复现 TypeSafe 的私有训练方法。

局限

  • 这是一个研究入门项目,不是 Jev 的复制品。
  • 准确率取决于数据质量、划分质量和编码器。
  • 字节编码器成本低,但对语言语义的理解较弱。
  • 预训练路径可能下载大模型,需要更多内存。
  • 单次前向评分要求在预测前给出完整的选项列表。
  • 速度对比使用的是小型本地解码器,而非大型商业模型。

许可证

代码以 MIT 许可证发布。下载的数据集和预训练模型遵循其各自条款。