AI学習のメモリ問題を解決

Sakana AIと東京大学の研究者たちが「DiffusionBlocks」を提案しました。これは、Transformerベースのネットワークを一度に1つのブロックずつ学習させるものです。学習に必要なメモリは、ブロックの数Bに応じてB分の1に削減されます。様々なアーキテクチャで性能が維持されます。ニューラルネットワーク学習におけるメモリ問題。エンドツーエンドのbackpropagationでは、すべての層の中間活性化を保存する必要があります。メモリ消費量はネットワークの深さに比例して増加します。モデルが深くなるにつれて、これは重大な学習のボトルネックとなります。既存の技術の一つであるactivation checkpointingは、必要に応じて活性化を再計算することで、活性化メモリを削減します。しかし、これはparameters、gradients、またはoptimizer statesのためのメモリを削減しません。Adam optimizerを使用する場合、各層はparameters、gradients、および2つのoptimizer states（momentumとvariance）のためのメモリを必要とします。これは1層あたりparameterサイズの4倍に相当し、activation checkpointingによっても変わりません。ブロックごとの学習は異なるアプローチを提供します。ネットワークをB個のブロックに分割し、それぞれを独立して学習させることで、メモリは約1/Bに削減されます。削減量はブロックの数に比例します。課題は、グローバルに一貫性のあるモデルを生成しつつ、各ブロックに対して原理に基づいた局所的な目的を定義することです。HintonのForward-Forwardアルゴリズムやgreedy layer-wise trainingのような以前のアプローチは、その場しのぎの局所的な目的に依存しています。これらはエンドツーエンドの学習に常に劣り、主に分類タスクに限定されていました。DiffusionBlocksは、理論的なギャップと以前の手法の適用範囲の限界の両方に対処します。https://arxiv.org/pdf/2506.14202。核となるアイデア：EulerステップとしてのResidual Connections。重要な洞察は、文献で確立された関連性に基づいています。Residual networksは、z ℓ = z ℓ − 1 + f θ ℓ ( z ℓ − 1 ) zℓ = zℓ−1 + fθℓ (zℓ−1) を介して各層の入力を更新します。これは、常微分方程式のEuler離散化に対応します。研究チームは、これらの更新が特にpに対応することを示しています。