小型TransformerでARC-AGI-1を44%達成、コスト67セント
原題: I trained a small transformer in 1.5hrs and it beats many LLMs
なぜ重要か
大規模モデルや膨大なデータに頼らず、低コストで高いサンプル効率を実現する手法は、AI研究の民主化と効率化の観点から業界全体に示唆を与える。
インドの研究者Mithil Vakde氏が、NVIDIA RTX 5090を用いてわずか1.5時間・67セントのコストで小型Transformerをゼロからトレーニングし、ARC-AGI-1公開評価で44%のスコアを達成した。同スコアは多くの大規模言語モデルを上回り、TRM/HRMと同水準。コードはオープンソースとして公開されており、ARC-AGI-2でも7%を記録している。
IIT Bombay出身のエンジニアMithil Vakde氏が、小型Transformerモデルをゼロからトレーニングしてサンプル効率の限界を探る研究の第3弾を公開した。今回のモデルはRTX 5090で1.5時間、コスト67セントでトレーニングされ、ARC-AGI-1の公開評価で44%を達成。以前のモデルの40%から向上しつつ、コストは大幅に削減された。
技術的アプローチとして、入出力ペアをトークン列に変換し、小型Transformerがテスト時にトレーニングセットと評価セットのパズルを自己回帰的に学習する「テスト時トレーニング（TTT）」方式を採用。クロスタスク学習のため各パズルに独立した加算埋め込みを付与し、3D RoPE埋め込みで位置情報を学習している。
前バージョンからの主な改善点は、アーキテクチャのモダン化（SwiGluの採用、RMSnorm化など）、データ多様性の向上、レイヤー数の8層への拡大（以前は4層）。コスト削減には、拡張処理の大幅削減、最適化アルゴリズムのAdamWからNormuonへの変更、flash attentionとflex attentionカーネルの活用が寄与した。また、入力トークンを学習対象から除外し出力トークンのみを損失関数に含める変更を加えたところ、スコアが40%から44%に向上した。理由は不明とされている。
データ面では、ARC-2から重複しない347件の新規タスクを慎重にフィルタリングして追加。この追加データを除いても約40%のスコアが維持されるが、計算量は約2倍必要になるという。Vakde氏はLucas Beyer氏やJeremy Howard氏ら著名研究者からも注目を集めており、コードはGitHubでオープンソースとして公開されている。