ARC-AGI 研究報告：高効率な Transformer モデルの構築

はじめに

NVIDIA GeForce RTX 5090 を使用し、1.5 時間以内で Small Transformer のゼロから（From Scratch）訓練を完了しました。

• 性能: 既存大規模言語モデル（LLM；TRM や HRM など）と同等、あるいはそれ以上のスコアを達成しています。
• 成果: ARC-2 データセットで 7% のスコア（正確には 44% 台の精度）を記録しました。
• コンセプト: 「Faster（高速）」「better（高性能）」「cheaper（低コスト）」かつオープンソースの実現です。

本研究は ARC-AGI シリーズの第 3 報であり、Lucas Beyer や Jeremy Howard, Rohan Anil といったトップリサーチャーからも注目を集めています。

研究の目的

深層学習における**「サンプリング効率（Sample Efficiency）」の解決**に取り組んでいます：

1. Transformer のサンプリング効率的な限界を探求する。
2. コスト削減によるイテレーションの迅速化を達成する。

---

テクニカル詳細

アプローチ概要

前回の実験と比較して、以下のアップグレードを加えています。

• トークン化: 各入力・出力ペアをトークンシーケンスに変換し、Small Transformer を用いてオートリグレッシブに訓練します。（テスト時にはテストラベルを隠蔽した状態でのゼロから訓練を実施）。
• クロスタスク学習: 各パズルに対して学習可能な独立した**加法埋め込み（Additive Embedding）**を割り当てます。位置情報は 3D RoPE で学習されます。
• オーガメンテーション: シーケンスを色と二面角（dihedral angle）の置換により拡張し、推論時に同様の拡張を行い逆変換を適用します（AAIVR: Augment-and-Invert Variants）。

前回のバージョンとの主な変更点

スコア向上とコスト削減のための主要な変更は以下の通りです。

パフォーマンス向上要因

• 現代アーキテクチャ:

  SwiGLU

  （GELU の代替）、

  RMSNorm

  （LayerNorm の代替）の採用。
• データ多様性: より効果的なデータシャッフルの実装。
• スケーリング: レイヤー数を 4 つから 8 つへ拡大。

コスト削減要因

• オーガメンテーションの削減: サンプル効率が大幅に向上。
• 最適化手法の変更:

  AdamW

  から

  Normuon

  （ノルムオン）へ変更。
• アテンション機構: Flex Attention カーネルを備えた Varlen Flash Attention の採用。

最大の仕様変更とデータ拡張

• 教師あり学習への回帰: 入力トークンを訓練対象から除外し、損失関数は出力トークンのみにします（完全な監督学習型へ）。これによりスコアがわずかに向上しました（40% → 44%）。
• データセットの拡張: ARC-2 の非重なりタスクを追加。ただし、データリーク防止のため慎重にフィルタリング済みです。

興味深い振る舞い

• 入力トークンを訓練しないため、アプローチは現在「監督学習型」です。テスト損失が低下したにもかかわらずスコアが良い点、およびばらつきが少ない点が特徴的です。
• Normuon の採用: 移行前は学習率を下げても「迷走（loiter）」する傾向がありましたが、Normuon に切り替えたことで収束問題が解消されました。

---

減法実験 (Ablations)

パフォーマンスへの最大の寄与は、**良い表現（3D RoPE + タスク別埋め込み）**であることが判明しました。

• 3D RoPE またはタスク別埋め込みを削除: スコアが急激に低下し、両方で 25% で飽和します。
• 入力データを訓練: スコアはわずかに低下（〜39%）。
• トレーニングセットの制限 (ARC-1 単体): 約 40% を維持します。
• 3D RoPE を 1D RoPE に変更: スコアが 〜24% に低下。
• タスク別埋め込みの削除: スコアが 〜24% に低下。

CompressARC スタイルの圧縮実験:

• 非監督学習時：パフォーマンスは〜18% に低下。
• 監督学習型時：〜15%。

---

貢献の仕方

コードはオープンソースです。修正してスコアを上げたりコストを削減したり自由にお願いいたします。（トレーニングデータを増やすのは禁止）。

目標

65% を達成してみてください。

改善の可能性

• RoPE の再考:

  PoPE

  （Positional Prompt Embeddings）が同等かそれ以上の性能を示す可能性があります。新しい位置埋め込みを発明することもあり得ます。
• アーキテクチャの近代化: さらに現代的な構造へ変更可能です。
• コスト削減: 手作りの GPU コードでコストをさらに 10 倍削減できる可能性があります。
• オーガメンテーションの除去: データ拡張を除外することでトレーニングコストを下げることが期待されます。

注記: トランスフォーマーだけで 45% を達成できるとは思っておらず、新しいアイデアが必要だと思いました。しかし、低いコストや FLOPs で達成できたことには驚きです。減法実験の結果、拡張子や合成データなしでも驚くほど多くのパフォーマンスが維持されていることがわかりました。

---

アペンドックス：著名リサーチャーによる批判と回答

多くの経験ある研究者からの議論を整理します。

1. テストセットでの訓練は不正／「テストデータを学習」

• 批判: 評価パズルの訓練は Cheating です。
• 回答: 誤りです。「テストデータを学習」はラベルを覚えることを指しますが、ここでは隠蔽したままです。ARC は Meta-learning ベンチマークであり、評価パズルから学ぶのが目的です。Test Time Training (TTT) は Meta-learning 分野で推奨されています。

2. 評価パズルの入力を学習させるのはリーク

• 批判: 入力データを訓練すると情報がリークします。
• 回答: 誤りです。「誘導的推論（Transductive Reasoning）」であり、継続学習の文脈では正当です。本質的に異なるアプローチは単一の前方関数の重みを直接 Meta-learn することです。（最新結果では入力を訓練する必要がなくなりましたが）。

3. テストポリシー違反

• 批判: テストポリシーに反しています。
• 回答: 誤りです。ポリシーは「出題者がテスト内容を知らなければならない」ことを指し、AI システム自体には制限がありません。ARC コミュニティでは常に TTT が許可されています。

4. 訓練コストの除外

• 批判: 訓練コストを含んでいません。
• 回答: 誤りです。全ライフサイクル（初期化からモデル訓練 + 推論実行）のコストを示しています。RTX 5090 を 2 時間使用して約 67 セント。

5. テストタイムトレーニングの実現可能性

• 批判: 同時訓練は非現実的です。
• 回答: LLM がインターネット全体を学習できていないのに「非現実的」と言えないのと同様です。「コスト/タスク」での比較では、同時学習モデルに不利になります。

その他の批判と反証

• ARC-AGI 解決＝AGI 達成ではない: 主張していません。
• 妥当な批判: 合成データの禁止（トップスコアが依存）、オフライン学習の排除（公平性の確保）については同意します。
• 他のアプローチの誤り:
  • TRM/HRM は再帰性を謳いますがアブレーションが少ない点に問題があります。
  • 7M モデルと宣伝されていますが、実際には O(100M+) の埋め込み重みが訓練されており誤解を招きます。
  • LLM のスコア上昇は主にトレーニング後のプロセス（合成データ量）によるものです。

---

LLM から得られた教訓

LLM は現在 ARC-AGI v1/v2 で飽和しています。そこから以下のことが推測できます：

• ARC-AGI は LLM 以前: 当初は事前トレーニングで推論能力が伴いませんでした。
• O1 の勝利: 75% を達成し、十分なデータがあれば何でも学習可能であることを示唆。
• ARC-2 登場時のリセット: 高性能な思考系 LLM もリセットされ、モデルは ARC パズルそのものを学習していることが判明。

結論:

• ARC-2 スコアはトレーニング後に駆動されます（合成データ量に依存）。
• ラボによるベンチマックスングが行われている可能性があります。
• LLM はサンプリング効率において全く優れていません。

---

完全な変更リスト

学習ダイナミクス

• Optimizer:

  AdamW-only

  から

  NorMuon

  + 補助的

  AdamW

  に変更。
• LR スケジュール:

  Warmup+Cosine

  から

  WSD

  （Warmup %, Hold, Linear Decay to Floor）に変更。
• LayerNorm:

  RMSNorm

  に置換。
• FFN:

  Linear->GELU->Linear

  から

  SwiGLU

  様式へ変更。
• Weight Decay: グループ別減衰の明示的実装（Attention Weights, Token Embeddings などに個別設定）。
• トレーニング目的:

  outputs["loss"]

  （非監督風）から

  outputs["output_loss"]

  （監督風）へ変更。

データセット構築

• ソース: ARC-1, ARC-2, ConceptARC、クロスデータセットタスクを使用。
• カラー拡張: グローバル置換から例ごとのアンプルアップル元組へ変更。
• 重複排除: ハッシングによるユニークサンプル多様性の向上。

速度向上（学習ダイナミクス変更なし）

• バッチ構造: Padded [B,S] から

  cu_seqlens

  を持つパケットトークンストリームへ変更。
• アテンションパス: Padded SDMA マスキングから

  Varlen Flash Attention

  サポートへ変更。

その他

• 訓練スプリット: Build-time で

  ("train", "test")

  から

  ("train",)

  へ変更。
• Scheduler Stepping: フラクショナルエポジ進捗に変更。