ARC Prize 2024 竞赛第二名解决方案：Omni-ARC 方法

非常欢迎对本文或论文提出改进建议 👍！

本文几乎是论文的复刻版，我建议直接访问论文，因为那里的格式更好。

TLDR (太长不看版)

我的解决方案是 MindsAI 团队方法（测试时微调）的实现。除了训练模型生成测试输出（原始 ARC 任务）外，它还训练模型执行其他任务，例如通过生成新输入来学习输入分布。

背景

• 商业背景
• 数据背景

我的方法动机

在 ARC 挑战 中，我们需要根据少量高维的输入输出图像对来学习转换规则。图像大小可达 30x30 像素，每个像素有 10 种不同的颜色。虽然这些图像不如现实世界的图像复杂，但它们仍然是高维数据。

如何从少量高维示例中学习？

要解决每个 ARC 问题，我们必须找到数据的正确表示。当人类解决任务时，最大的挑战是找到看待问题的正确视角。一旦我们有了数据的正确视角，ARC 问题就变得微不足道了。

数据的正确表示可以降低数据的维度，使得从极少量的示例中学习转换成为可能。

如何学习 ARC 问题的良好表示？

如果我们训练模型执行需要良好数据表示的任务，那么模型内部很可能会开发出所需的表示。

我的见解是，我们可以用多种不同的方式使用 ARC 问题来学习这种表示，而不仅仅是原始提出的任务（即给定少量输入输出对，为图像生成输出）。

Omni-ARC：训练单个模型执行多个 ARC 相关任务

• examples + input -> output。ARC 数据集的原始任务。
• inputs -> input。生成新输入需要理解网格的分布。这也适用于输出，输出也应遵循某种分布。
• examples -> code。这是 Ryan Greenblat 使用 GPT-4o 采用的方法。
• code + input -> output。这等同于第一个任务，但不是将示例作为输入，而是给出问题的代码定义。
• code -> inputs。任务的每个输入都遵循某种分布，给定分布的描述，模型应该能够生成该分布的样本。
• inputs -> code。我们也可以做相反的任务，给定一些输入，编写代码来生成该分布。
• examples + input + output -> is the output correct?。可以训练模型来验证提议的输出是否正确。
• examples + input + output options -> select the correct output。我们可以训练模型在多个选项中选择正确的输出。

列出的所有任务都要求模型学习 ARC 图像的一些有用表示。Omni-ARC 方法背后的想法是训练单个模型来执行所有任务，期望跨所有任务的共享表示比训练模型执行单个任务具有更好的泛化能力。

Omni-ARC，一个执行所有 ARC 相关任务的单一模型（而且它有一个非常酷的 logo）

方法概述

解决方案简述：

1. 获取 Qwen2.5-0.5B-Instruct 并在公开可用的 ARC 数据集上对其进行微调。模型被微调以：
   1. 为测试样本生成输出
   2. 学习输入分布并生成新输入。
2. 使用私有测试数据进行测试时微调，仅针对生成测试输出的任务。
3. 使用数据增强进行推理，并通过投票选择预测结果。
4. 与 2020 年公开解决方案进行集成。

训练

数据

我使用了以下公开可用的数据集进行训练：

对于所有数据集，我训练模型执行两个任务：

• examples + input -> output。ARC 数据集的原始任务。
• inputs -> input。生成新输入需要理解网格的分布。这也适用于输出，输出也应遵循某种分布。

新生成输入的示例

数据增强

对每个问题，对所有输入和输出应用相同的数据增强。数据增强是以下增强的组合：

• 旋转
• 翻转
• 颜色变化
• 训练和测试示例之间的交换

问题增强

除了数据增强外，我还通过仅对输入或输出应用转换来进行问题增强。这种转换通过将原始 ARC 转换与随机选择的新转换组合来创建新的 ARC 问题。

这些新转换需要是可逆的，否则新生成的问题可能无法解决。我使用了以下额外的转换：

• 旋转和/或翻转
• 填充图像
• 放大 (Upscale)
• 镜像

论文中的问题增强示例

问题表示

我使用了一种非常简单的 ARC 网格文本表示作为 LLM 的输入。网格被包含在 Markdown 代码片段中，形状定义在第一行，每一行都有编号。

```grid shape: 3x3
1 100
2 010
3 001
```

训练超参数

模型使用 LoRA 进行微调。在进行全模型微调时未发现显著改进，并且在测试时微调中，微调已训练的 LoRA 适配器似乎比创建全新的适配器更有益。

• 模型：Qwen2.5-0.5B-Instruct
• LoRA rank: 128
• 学习率：5e-5，带有预热的线性调度
• Batch size: 16
• 训练步数：2e5
• 最大序列长度：8196
• 在 2xA6000 GPU 上训练

我使用 huggingface 的 trl 和 accelerate 库进行训练。

测试时微调 (Test-time fine-tuning)

仅在 ARC 任务上微调模型不足以在私有测试集上表现良好。通过应用测试时微调，我们可以将一个模型解决的问题数量从 11 个提高到 33 个。

这是我对测试时微调的解释：

• 对于每个有 n 个训练样本的测试问题，我使用 n-1 个训练样本微调模型，并使用剩余的样本作为测试样本。测试样本的选择是在训练过程中随机进行的。
• 我像之前的训练一样使用了数据增强。
• 我为每个测试问题微调了一个模型，因此每次提交都会生成 100 个微调模型。
• 我在测试时微调中使用 batch size 1，以便尽可能快地学习新问题。
• 模型在每个问题上微调了约 300 步。
• 令人惊讶的是，测试时微调的最佳学习率是 8e-5，高于训练时使用的学习率 (5e-5)。如果有更多的计算资源，使用更慢的学习率、更大的 batch size 和更长的时间进行训练，很可能会获得更好的结果。

由于提交时间有限，测试时微调仅应用于预测测试输出的规范 ARC 任务。但它也可以应用于生成新输入的任务，或验证输出正确性的任务。

这种为每个任务训练单个模型且 batch size 为 1 的非寻常配置是由于计算资源和提交时间的限制而产生的。这是能够最快学习新测试问题的配置。

推理

推理时也应用了数据增强，并从预测中还原数据增强以获得原始输出。对每个问题进行 96 次预测，并使用投票来选择最有希望的预测。就像 MindsAI 的 AIRV（增强、推理、还原增强、投票）一样。

推理使用 temperature 0 进行。

使用 VLLM 生成预测。每个微调模型用于为其问题生成预测。

集成 (Ensemble)

我将我的模型预测与2020 年解决方案进行了集成。由于 2020 年解决方案只需要 CPU，我设法在后台运行它，同时使用 GPU 进行我的模型的微调和推理。我只需要注意 RAM 使用情况，因为两个作业必须共享相同的内存。

集成策略非常简单，只需从每个解决方案中获取第一次尝试的结果。

结果

这种方法在 ARC 排行榜上得分为 40。

相同的方法（没有测试时微调）可以解决 32% 的评估任务，当使用 32 次预测进行投票时，top_2 准确率达到 22%。由于硬件资源有限，我通常不在评估数据集上评估带有测试时微调的模型。Kaggle 每周提供 30 小时的 GPU，但我们可以每天进行 3 次提交，这相当于 36 小时的计算量。因此，使用提交来查看测试时微调的性能要便宜得多，在那里我们每周有 7 倍以上的可用计算量。

提交细节

提示词不够，需要测试时微调

显然，这次比赛表明 LLM 需要测试时微调才能执行新任务。少样本提示词 (Few-shot prompting) 不足以让模型学习新任务。

可以训练模型来验证任务的正确性

在挑战的最后几周，我尝试继续 Omni-ARC 方法并训练模型以：

1. 验证输出是否正确
2. 在两个选项之间选择正确的输出

我们的想法是，如果我们用更准确的机制替换投票选择机制，我们可以提高 LB 分数。使用训练好的模型，我使用接近 1 的采样温度为原始 ARC 数据集生成了错误的预测。

如上表所示，我能够在评估数据集上取得有希望的结果。这些数字是针对 32 次预测的。

然而，我无法使用这种方法提高 LB 分数。我的假设是，测试时微调模型的预测分布与冻结模型的分布不同。因此，测试时微调模型的投票准确率可能比表中显示的冻结模型高得多。

这些验证器模型可以从测试时微调中受益，但由于提交时间有限，我无法测试这一假设。

更多信息请参阅 Iteration 47 和 Iteration 45。

使用代码解决任务对我没用

我还尝试通过训练模型执行以下额外任务来扩展 Omni-ARC 方法：

• examples -> code。这是 Ryan Greenblat 使用 GPT-4o 采用的方法。
• code + input -> output。这等同于第一个任务，但不是将示例作为输入，而是给出问题的代码定义。

为此，我构建了一个小型领域特定语言 (DSL)，并用 Python 代码重现了 285 个 ARC 训练任务。这是一个耗费了大约 3 周的劳动过程。

不幸的是，模型泛化能力不佳。它只能解决 5% 的评估任务，而且这些任务与训练任务非常相似。在私有测试集上，一次幸运的提交解决了 1 个任务。

我相信这种方法有很大的潜力，但由于挑战即将结束且其他团队在排行榜上不断进步，我不得不转向其他方法。

更多信息请参阅 Iteration 26, Iteration 40 和 Iteration 42。

数据集的质量是相关的

在挑战的最后几周，我尝试将 BARC 数据集 添加到训练数据中。令人惊讶的是，尽管声称有大量的不同任务 (400k)，但我没有在评估数据集或排行榜上看到任何显著改进。更多信息请参阅 Iteration 48. 更多外部数据。

这令人惊讶，因为原始 ARC 数据集在增加训练任务数量时显示出明显的趋势：

我猜 GPT4 自动生成的任务相对于原始 ARC 任务没有太多的新颖性。

正确的模型大小

Qwen2.5-0.5B 是我的方法和可用提交计算资源的正确模型大小。

在第一步中，我尝试了较小的模型，如 SmolLM2-135M 和 NanoLM-0.3B，但它们没有达到与 Qwen2.5-0.5B 相同的准确率。更多信息请参阅 Iteration 46. 重访小型 LLM。

在我的最后一次尝试中，我还尝试了较大的模型，如 Qwen2.5-1.5B 和 Qwen2.5-7B。这些模型表现出更高的数据效率，它们在相同的训练步数下达到更小的训练损失。这些模型的问题在于它们在提交时的微调和推理速度较慢。此外，由于 VRAM 要求，我们必须减少训练样本的长度。如果提供更好的硬件和更多的提交时间，LB 分数很可能会通过这些更大的模型得到提高。