Share
Embed
2025年度人工知能学会全国大会チュートリアル講演「深層基盤モデルの数理」
2025年度人工知能学会全国大会 (JSAI2025) にて行ったチュートリアル講演の資料です.
Taiji Suzuki
May 30, 2025
Transcript
161 slides
-
鈴木大慈 2 所属 ➢ 東京大学大学院情報理工学系研究科数理情報学専攻・教授 ➢ 東大次世代知能科学研究センター研究部門研究者(研究知能部門) ➢ 理研 革新知能統合研究センター
深層学習理論チーム チームディレクター 専門 ➢ 機械学習の理論:数理統計学,統計的学習理論,確率的最適化 解釈可能性: 説明可能性,データの可視化,メンテナ ンスの容易化 各種テクニックの解析: アーキテクチャの解析,損失関数の設計, 最適化技法の解析 深層学習の原理解明: 「表現理論」「汎化誤差理論」「最適化 理論」 学習の本質解明: “良い”学習手法の特徴付け,統一理論, 深層学習を優越する方法論の提唱 応用 基礎 鈴木大慈 情報理工学系研究科 確率論 幾何学 関数解析 最適化理論 数学 数理統計 スパース推定 関連する機械学習理論 特徴抽出 カーネル法 深層学習の理論 主な研究内容 ➢ 深層学習を含む様々な学習機構に関する理論 ➢ 学習理論を通じた各種学習手法の汎化解析や学習アルゴリズムの収 束理論 ➢ 確率的最適化による大規模複雑な機械学習問題の効率的解法 著書/授賞 ➢ICLR2021 outstanding paper award. ➢『確率的最適化(機械学習プロフェッショナルシリーズ)』講談社,2015 年8月8日. ➢金森敬文,鈴木大慈,竹内一郎,佐藤一誠:『機械学習のための連続最適化 (機械学習プロフェッショナルシリーズ)』講談社,2016年12月7日. ➢文部科学大臣表彰・若手科学者賞「深層学習の原理解明に向けた統計的学習 理論の研究」.文部科学省,2020年4月7日. ➢第11回日本統計学会研究業績賞 (2017年度).2017年9月5日. ➢Satoshi Hayakawa and Taiji Suzuki:日本神経回路学会論文賞.日本神経回 路学会,2021年9月23日. ➢日本応用数理学会,ベストオーサー賞(論文部門).2019年9月4日. 主な活動場所 • 国内:IBIS, 統計連合大会 • 国外:NeurIPS, ICML, ICLR, ACML, ... (ACML steering committee member) -
知能の起源 4 有櫛動物 平板動物 刺胞動物 左右相称動物 海綿動物 [Najle et al:
Stepwise emergence of the neuronal gene expression program in early animal evolution. Cell, 186, 4676–4693, 2023] [Hayakawa, E., Guzman, C., Horiguchi, O. et al. Mass spectrometry of short peptides reveals common features of metazoan peptidergic neurons. Nat Ecol Evol 6, 1438–1448, 2022] -
5 約7億~8億年前:ペプチド分泌細胞 ➢ 光,pH,個体群密度,食物源由来のグリシンなどに反応 [Moroz LL. Convergent evolution of neural
systems in ctenophores. J Exp Biol. 2015 Feb 15;218(Pt 4):598--611] [Najle et al: Stepwise emergence of the neuronal gene expression program in early animal evolution. Cell, 186, 4676–4693, 2023] • 外部刺激に反応する機能の発現:生存確率の増大 (受動から能動へ) • 入力「センサー」→ 出力「運動機能」 「遺伝子レベルの自然淘汰」から「個体レベルの実時間環境適応」 → いずれ「脳」へ発達 光受容体 繊毛 [G.Jekely et al : An option space for early neural evolution, Pil. Trans. R. Soc. B: 370:2015.0181 -
汎化能力 6 餌の認識・追尾 捕食可能かどうかの判断 ➢ 毒の有無/安全性 捕食者の認識・逃走判断 餌の探索 予測精度の向上 =
生存確率の増大 常に変化する環境への適応 ⇒ 丸暗記ではない汎化の必要性 Image credit: @BobNichollsArt -
情報圧縮による恩恵 8 𝑥1 𝑥2 𝑥3 𝑔 𝑥1 2 𝑥2 2
𝑥3 2 3層 ( 𝑥 を中間層で生成):O(poly(𝑑𝑥 ))ノードで十分 2層 (𝑥を直接使用して近似):Ω(exp(𝑑𝑥 ))ノードが必要 まず二乗和𝑥1 2 + ⋯ + 𝑥𝑑𝑥 2 を作ってから𝑔を作用. (Eldan&Shamir, 2016) gはBessel関数を元に構成 全方向をケアする必要がある (座標軸方向だけではダメ) 3層 2層 (中間層で座標軸方向 だけを見ればよい) 𝑥の表現として圧縮された表現 𝑥 を用 いるかどうかで差が生まれる. -
言語の形成 9 • シンボル化された概念を音声で伝達 • 情報の共有:未経験事象も予測可能に ⇒ 生存確率の飛躍的向上 • 餌の場所
• 道具の作り方 • 摂取可能な食糧 • 共同での主張 概念形成 ↓ 特徴学習 -
モデル訓練の計算量 11 [Sastry et al.: Computing Power and the Governance
of Artificial Intelligence. arXiv:2402.08797] 訓練時計算量:6ヵ月で2倍 モデルサイズとパフォーマンス [Real et al.: Regularized evolution for image classifier architecture search. 2019] モデルサイズの指数的増大 [Luccioni. The mounting human and environmental costs of Generative AI. Apr. 2023.] -
12 Alex-net 2 × GTX 580 1.581 TFLOPS for FLOAT
1.5GB memory xAI 200,000 × H100 ・・・ 800 TFLOPS for FP16 80GB memory 2012 2024 [参考] 産総研 ABCI 3.0: 6,128 × H200 ・・・ -
スケーリング則 13 Reducible loss [Kaplan et al.: Scaling Laws for
Neural Language Models, 2020] [Henighan et al.: Scaling Laws for Autoregressive Generative Modeling, 2020] モデルサイズ固定 (基本的に訓練データサイズと思ってよい) [Brown et al.: Language Models are Few-Shot Learners, 2020] (GPT-3モデルの解析) log(予測精度)=−𝛼 log 𝑛 + log(𝐶) -
深層学習 vs 浅層学習 (異なるスケーリング則) 学習すべき真の関数の形状によっては深層が有利になる 14 深 層 浅 層
縮小ランク回帰 特徴空間の次元 が低い状況は深 層学習が得意 区分滑らかな関数 不連続な関数の 推定は深層学習 が得意 Besov空間 滑らかさが非一 様な関数の推定 は深層学習が得 意 低次元データ データが低次元 部分空間上に分 布していたら深 層学習が有利 [Suzuki, 2019] [Schmidt-Hieber, 2019] [Nakada&Imaizumi, 2019][Chen et al., 2019][Suzuki&Nitanda, 2019] [Imaizumi&Fukumizu, 2019] 推 定 精 度 -
カーネル法と深層学習の違い 15 推定誤差 データサイズ 少ないデータサイズ では浅い学習が良い. 大きなデータサイズ では深層学習が良い. 深層学習 浅い学習
• スケーリング則自体は比較的古典的な理論からも導出できる. • しかし,これだけの「データ量」「モデルサイズ」「学習問題の複雑さ」 で実証されることはなかった. ref -
特徴学習の有無 • 線形モデル 16 • カーネルモデル • 深層モデル 非線形化 可変基底化
学習可能 固定 学習可能 学習可能 学習可能 固定 -
18 𝜑1 𝜑2 カーネル法 (特徴学習無し) は基底関数を固定している 𝜑(𝑥; 𝜃1 )を特徴量とするNN 𝜑(𝑥;
𝜃2 )を特徴量とするNN ターゲットの関数 半径 -
深いネットワークは指数関数的に表現力が高い19 𝐿’ 𝐿 = 深さを横幅で代替しようとすると 指数関数的な広さが必要 [Arora, Basu, Mianjy, Mukherjee:
Understanding Deep Neural Networks with Rectified Linear Units. ICLR2018.] 横幅 vs 深さ -
関数の平滑性の非等方性 22 不変な方向 変化する方向 𝑠1 , 𝑠2 , 𝑠3 :
滑らかさ (非平滑) 𝑠1 , 𝑠2 ≪ 𝑠3 (平滑) • 真の関数の滑らかさが方向に依存 • 多様体に直交する方向にはほぼ定数 (滑らかさ大) データが低次元多様体からはみ出る場合: [Suzuki&Nitanda: Deep learning is adaptive to intrinsic dimensionality of model smoothness in anisotropic Besov space. NeurIPS2021.] → 中間層で「重要な方向」を取り出すことで次元の呪いを回避. MNIST: 784 dim/ 13.4 intrinsic-dim [Facco et al. 2017] -
(超)高次元入力NNの学習理論 24 不変な方向 変化する方向 𝑠1 , 𝑠2 , 𝑠3 :
滑らかさ • 真の関数が方向によって異なる滑らかさを持つ状況では DNNは重要な方向を見つけ,次元の呪いを回避する. • 一方で,浅い学習法は次元の呪いを受ける. Hayakawa and Suzuki: Neural Networks 2020, 日本 神経回路学会論文賞. 関連研究: - 教師生徒設定における大域的最適化と次元 の呪いの回避 - 深層学習の浅層学習への優位性: Suzuki and Akiyama: ICLR2021, spotlight. : 非等方的Besov空間 (𝐵𝑝,𝑞 𝑠(ℓ) ). 真の関数のモデル: 非等方的Besov空間の元の合成関数 ➢ 滑らかさが方向によって異なる関数空間 ➢ 合成することで様々な形状を実現 (例:多様体上の関数: 一層目で座標を抽出,二層目がその座標上の関数) Def. (非等方的Besov空間) 真の関数の滑らかさが方向によって大きく 異なる状況で,ほとんどの方向に対して滑 らかならば次元の呪いを回避できる. → 「非等方的Besov空間」を用いた理論. Suzuki&Nitanda: Deep learning is adaptive to intrinsic dimensionality of model smoothness in anisotropic Besov space. NeurIPS2021, spotlight. -
(超)高次元入力NNの学習理論 25 不変な方向 変化する方向 𝑠1 , 𝑠2 , 𝑠3 :
滑らかさ • 真の関数が方向によって異なる滑らかさを持つ状況では DNNは重要な方向を見つけ,次元の呪いを回避する. • 一方で,浅い学習法は次元の呪いを受ける. Suzuki&Nitanda: Deep learning is adaptive to intrinsic dimensionality of model smoothness in anisotropic Besov space. NeurIPS2021, spotlight. Hayakawa and Suzuki: Neural Networks 2020, 日本 神経回路学会論文賞. 関連研究: - 教師生徒設定における大域的最適化と次元 の呪いの回避 - 深層学習の浅層学習への優位性: Suzuki and Akiyama: ICLR2021, spotlight. : 非等方的Besov空間 (𝐵𝑝,𝑞 𝑠(ℓ) ). 真の関数のモデル: 非等方的Besov空間の元の合成関数 ➢ 滑らかさが方向によって異なる関数空間 ➢ 合成することで様々な形状を実現 (例:多様体上の関数: 一層目で座標を抽出,二層目がその座標上の関数) Def. (非等方的Besov空間) 真の関数の滑らかさが方向によって大きく 異なる状況で,ほとんどの方向に対して滑 らかならば次元の呪いを回避できる. → 「非等方的Besov空間」を用いた理論. 直感 -
推定誤差の評価 26 深層 (次元の呪いを受ける) 浅層 浅い学習方法は一番滑らかでない方向の滑らかさ(𝒔𝟏)が 支配的で,次元の呪いを受ける. 主結果 (最小二乗推定量) ※今回は最適化手法に関しては議論せず,最適化はできるものと仮定する.
, Let 各方向への滑らかさの調和平均が収束レートを決める. 例: 𝑯 = 𝟏の時 浅い学習方法との比較 (informal): 少ない数の方向において𝒔𝒊 が 小さく(滑らかでない),その 他の方向には𝒔𝒊 が大きい(滑ら か)であるとき,次元の呪いを 回避できる. -
ニューラルネットの最適化理論 28 深層ニューラルネットワークをデー タにフィットさせるとは? 損失関数:データへの当てはまり度合い 𝑖番目のデータで正解していれば 小さく,間違っていれば大きく 𝑊: パラメータ 損失関数最小化
(Wは数十億次元) 通常,確率的勾配降下法で最適化 最適値 -
現状 • 非線形FNNの特徴学習に関する一般的な最適化 の保証は二層までしかなされていない. • 三層以上の解析もあるが,本質的には二層の解 析に帰着させる形になりがち. • 線形の多層NNの場合は,多くの解析がある. ➢任意の局所最適解が大域的最適解
[Kawaguchi, 2016; Lu&Kawaguchi, 2017; Yun, Sra&Jadbabaie, 2018] ➢深さは勾配法を加速させる (前処理付き勾配法と同様の効果) [Arora, Cohen, Hazan, 2018] ➢「陰的正則化」によって低ランクな解に収束する. 29 -
勾配法と陰的正則化 • 小さな初期値から勾配法を始めるとノルム最小 化点に収束しやすい→陰的正則化 30 [Gunasekar et al.: Implicit Regularization
in Matrix Factorization, NIPS2017] [Soudry et al.: The implicit bias of gradient descent on separable data. JMLR2018] [Gunasekar et al.: Implicit Bias of Gradient Descent on Linear Convolutional Networks, NIPS2018] [Moroshko et al.: Implicit Bias in Deep Linear Classification: Initialization Scale vs Training Accuracy, arXiv:2007.06738] 初期値 (原点付近) 解集合 最も「単純な」解 勾配法による最適化 多層の線形NNを判別問題でGDすると “スパースな解”が得られる. -
Implicit regularization (陰的正則化)31 • ニューラルネットワークの学習では様々な「陽的正則化」を用いる: バッチノーマリゼーション,Dropout,Weight decay,... • 実は深層学習の構造が自動的に生み出す「陰的正則化」も強く効いて いるという説.
例:線形ネットワーク (L2正則化学習) 任意の局所最適解は低ランクになる: モデルの複雑さが大幅に削減されている. (見た目のパラメータ数) 𝐿𝑚2→ 2𝑚 (実質的パラメータ数) -
低次元特徴量の問題 32 カーネル法: 𝑛 = Ω(𝑑𝑝) ニューラルネット: 𝑛 = O(𝑑)
情報理論的下限 (サンプル複雑度): • ガウシアンsingle indexモデル: ➢ 方向𝜃 ∈ ℝ𝑑とリンク関数𝜎∗ を推定する必要がある. • 意味のある方向は一方向 (𝜃方向) ⇒ 特徴学習 ➢ 𝜎∗ は次数𝑝かつ,情報指数 𝒌 の多項式とする: 必要な訓練データサイズ𝑛: [Ghorbani et al. 19; Donhauser et al. 21; Gavrilopoulos et al. 24;…] [Bach 17; Barbier et al. 19; Damian et al. 24;…] • NNの方が少ないデータ数で学習できる. • しかし,非凸目的関数を最小化する必要がある. “統計的複雑度 vs 計算量複雑度”のトレードオフ He𝑖: 次数𝑖のエルミート多項式 -
情報指数 • パラメータのランダム初期値における勾配: 34 He𝑖: 次数𝑖のエルミート多項式 Def (情報指数 𝑘 [Ben-Arous
et al. 2021]) 𝜎∗ の情報指数 (information exponent)は,𝛼𝑖 ∗ ≠ 0なる最小の次数である: = O(𝑑−(𝑘−1)/2) • シグナルの強さは情報指数で特徴づけられる. • SGDによる勾配の標準偏差=O(1/ 𝐵)⇒ミニバッチサイズ𝐵 = Ο(𝑑𝑘−1)がシグナルの 情報を得ることができる. 2層-NN: (⟨𝜃, 𝑤⟩を大きくしたい) -
SGDによる学習の複雑さ 36 定理 (informal [Ben Arous et al. 21;Bietti et
al. 22;Damian et al. 23]) SGDによって二層NNを訓練するなら,情報指数𝑘の𝑓∗ は, 以下の更新回数によって学習できる: 𝑛 ≃ 𝑑𝑘−1. • NN + SGDはカーネル法の複雑さ 𝑑𝑝 を優越する(特徴学習). • 𝑑𝑘−1 は情報理論的下限を達成していない. • SGDによる学習によって情報理論的下限を達成可能か? → バッチ再利用によって達成可能 • 情報指数は適切な指標か? → 情報指数を一般化した生成的指数 (generative exponent) が重要 疑問 ※ SGDの場合,更新回数=訓練データサイズ -
SQ, CSQ 下限 37 相関クエリ How many (noisy) correlation query
should be observed? • 相関統計的クエリ (correlation statistical query; CSQ): アルゴリズムが任意の関数𝜙と𝑦の間の相関の近似値 𝑞を取得可能な時 ➢ 考え方: 𝑛個のデータを用いて近似することで,𝜏 ∼ 𝑛−1/2とできる. 定理 (informal [Damian et al. 22; Abbe et al. 23; Damian et al. 24]) 多項式オーダー計算量の任意のアルゴリズムは,情報指数𝑘を持つ多項式𝑓∗ を学習 するのに以下のデータ数が必要: • CSQ 学習法: 𝒏 ≥ 𝒅𝒌/𝟐 • SQ 学習法: 𝒏 ≥ 𝒅 (上記の𝑛個のデータのア ナロジーを用いている) (with |𝜙| ≤ 1) • 統計的クエリ(statistical query; SQ): より一般的なクエリ ∼ 𝑛−1/2 ∼ 𝑛−1/2 -
計算量のまとめ 38 𝑑 ෨ 𝑂(𝑑) ෨ 𝑂(𝑑𝑘/2) ෨ 𝑂(𝑑𝑘−1) 𝑂(𝑑𝑝)
情報理論的 下限 (平均場ラン ジュバン) [Chen&Meka, 20] 平滑化SGD [Damian et al. 23] CSQ 下限 通常のSGD [Ben-Arous et al. 22] カーネル法 [Ghorbani et al. 21] SQ 下限 確率的勾配降下法: CSQ ガウシアンsingle indexモデル (𝑘: 情報指数, 𝑝: 多項式の最大次数) バッチ再利用を使えば,CSQ型アルゴリズムの下限を突破できる. 𝒏 = ෩ 𝑶(𝒅)の反復数かつサンプル複雑度で十分. [Dandi et al. 2024][Lee, Oko, Suzuki, Wu; NeurIPS2024] 非特徴学習の学習の複雑度 特徴学習をすることで到達可能な複雑度 -
バッチ再利用 • [Dandi et al. 2024] SGD + バッチ再利用 によってより高次の情報を取得
することができる. 40 高次の相関! (CSQに入らない) (if 𝑤(0) = 0) 定理 [Mondelli & Montanari 18; Barbier et al. 19; Chen & Meka 20] 任意の多項式𝜎∗ に対して,ある多項式 𝒯 が存在して: 多項式で変換することで情報指数が1もしくは2に減少する. ⇒生成的指数 (generative exponent) バッチを再利用することで𝑦ℓとの相関を生成することができる. -
主定理 41 定理 (バッチ再利用するSGDによる学習複雑度) ガウシアンsingle indexモデルの学習において,バッチ再利用SGDは 推定誤差 を達成するのに,以下の更新回数 で十分: •
Batch-reuse can break the CSQ lower bound. 非特徴学習の複雑度 特徴学習の複雑度 𝑑 ෨ 𝑂(𝑑) ෨ 𝑂(𝑑𝑘/2) ෨ 𝑂(𝑑𝑘−1) 𝑂(𝑑𝑝) 情報理論的 下限 (平均場ラン ジュバン) [Chen&Meka, 20] 平滑化SGD [Damian et al. 23] CSQ 下限 通常のSGD [Ben-Arous et al. 22] カーネル法 [Ghorbani et al. 21] バッチ再利 用したSGD 通常の確率的勾配降下法: CSQ -
主定理 42 定理 (バッチ再利用するSGDによる学習複雑度) ガウシアンsingle indexモデルの学習において,バッチ再利用SGDは 推定誤差 を達成するのに,以下の更新回数 で十分: •
Batch-reuse can break the CSQ lower bound. 非特徴学習の複雑度 特徴学習の複雑度 𝑑 ෨ 𝑂(𝑑) ෨ 𝑂(𝑑𝑘/2) ෨ 𝑂(𝑑𝑘−1) 𝑂(𝑑𝑝) 情報理論的 下限 (平均場ラン ジュバン) [Chen&Meka, 20] 平滑化SGD [Damian et al. 23] CSQ 下限 通常のSGD [Ben-Arous et al. 22] カーネル法 [Ghorbani et al. 21] バッチ再利 用したSGD 通常の確率的勾配降下法: CSQ 特徴学習によって 計算複雑度もサンプル複雑度も 得をする -
オーバーパラメトライゼーション 横幅が広いと局所最適解が大域的最適解になる. 44 • 二種類の解析手法 ➢ Neural Tangent Kernel (NTK)
➢ Mean-field analysis (平均場解析) … 狭い横幅 広い横幅 自由度が高いので,目的関数を減少さ せる方向が見つけやすい. 0 0 -
損失の地形 • 横幅の広いNNの訓練誤差には孤立した局所最 適解がない.(局所最適解は大域的最適解とつ ながっている) 45 [Venturi, Bandeira, Bruna: Spurious
Valleys in One-hidden-layer Neural Network Optimization Landscapes. JMLR, 2 34, 2019.] 定理 𝑛個の訓練データ 𝑥𝑖 , 𝑦𝑖 𝑖=1 𝑛 が与えられているとする.損失関数ℓは 凸関数とする. 任意の連続な活性化関数について,横幅がデータサイズより広い (𝑀 ≥ 𝑛)二層NN𝑓 𝑎,𝑊 (𝑥) = σ𝑚=1 𝑀 𝑎𝑚 𝜂(𝑤𝑚 ⊤𝑥)に対する訓練誤差 𝐿 𝑎, 𝑊 = 1 𝑛 σ𝑖=1 𝑛 ℓ(𝑦𝑖 , 𝑓 𝑎,𝑊 (𝑥𝑖 ))の任意のレベルセットの弧状連結 成分は大域的最適解を含む.言い換えると,任意の局所最適解は 大域的最適解である. こうはならない こうなる (つながっていない) ※とはいえ,勾配法で大域的最適解に到達可能かは別問題. -
粒子による描像 (平均場解析) 46 例: 他の粒子との相互作用がある 𝑀 → ∞ 多粒子化(平均場): 𝑀
→ ∞, 𝑡 → ∞の極限で粒子𝜃𝑗 の分布𝜇𝑡 は以下の分布に収束: 重要:分布𝝁に対しては凸関数!(if 損失が凸) 定理 (Hu, Ren, Šiška, and Szpruch, 2021; Mei, Montanari, and Nguye, 2018) -
粒子による分布の最適化 47 • 各ニューロンのパラメータを一つの粒子とみなす. • 各粒子が誤差を減らす方向に動くことで分布が最 適化される. 1つの粒子 [Nitanda&Suzuki, 2017][Chizat&Bach,
2018][Mei, Montanari&Nguyen, 2018] データへの当てはまりを 良くする方向に変化 M個の粒子が移動 (各粒子の移動方向:勾配方向) (分布) [Nitanda&Suzuki, 2017] -
平均場ランジュバン動力学 48 Example: 平均場ランジュバン動力学: 有限粒子での近似: (これは有限横幅NNのGLDになっている) Remark GLDの場合 モデル 目的関数
𝜇で微分 (局所的に線形化させ ていることに相当) (局所線形近似を最小化させる方向に移動) 対応する目的関数 -
非線形Fokker-Plakck方程式 • 平均場ランジュバン動力学は以下の非線形Fokker- Planck方程式に従う: 49 ベクトル場: 𝑏(𝑥, 𝜇𝑡 ) 質量:
𝜇𝑡 (𝑥) ※ 𝛿𝐹(𝜇𝑡) 𝛿𝜇 が𝜇𝑡 に依存するので「非線形」,普通のGLDは𝜇𝑡 について「線形」 これはℒを最小化させるWasserstein勾配流である. (𝑝𝜇𝑡 の定義) −𝑣𝑡 とおく すると, [連続の方程式] (∵連続の方程式) GLDの場合: 参考 -
MFLDの収束 50 近接点更新解: 𝑝𝜇𝑡 は一様に対数ソボレフ不等式 (定数𝛼)を満たし,ℓ𝑖 が凸であるなら, 定理 (Entropy sandwich)
for all 𝜈. [Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (2022)] (c.f., Mirror descent, exponentiated gradient) LSIは各ニューロ ンが有界なら成 り立つ by Holley--Stroock. (線形収束!) 損失を線形化して得られる解 𝜇𝑡 𝜇∗ 𝑝𝜇𝑡 -
McKean-Vlasov過程との関係 • 粒子間相互作用のある確率微分方程式 (McKean, Kac,…, 60年代) 51 𝑡 = 1
𝑡 = 2 𝑡 = 3 𝑡 = 4 Propagation of chaos: 粒子数無限大で各粒子が独立であるかのように振る舞う現象. Q: どれだけの数の粒子 があればお互い十分独 立な振る舞いをするの か? 粒子はお互いに複雑に 相互作用するので証明 は簡単ではない. -
離散粒子近似:Propagation-of-chaos52 時間 離散化 空間 離散化 損失関数の凸性と平滑性の仮定のもと, 𝑝𝜇 は対数Sobolev不等式を定数𝛼で満たすとする. 定理 [Chen,
Ren, Wang, 2022; Suzuki, Wu, Nitanda, 2023; Nitanda et al. 2025] : proximal Gibbs measure 既存研究では粒子数は時間に対して指数関数的に依存 [Mei et al., 2018; Javanmard et al., 2019; De Bortoli et al., 2020] 1. 𝐹: 𝒫 → ℝ is convex and has a form of 𝐹 𝜇 = 𝐿 𝜇 + 𝜆1 𝔼𝜇 𝑥 2 . 2. (smoothness) ∇𝛿𝐿 𝜇 𝛿𝜇 𝑥 −∇𝛿𝐿 𝜈 𝛿𝜇 𝑦 ≤ 𝐶(𝑊2 𝜇, 𝜈 + 𝑥 − 𝑦 ) and (boundedness) ∇𝛿𝐿 𝜇 𝛿𝜇 𝑥 ≤ 𝑅. Assumption: 線形収束 (空間離散化) (時間離散化) , -
一様対数Sobolev不等式 53 𝑋 𝑘 (1) 𝑋 𝑘 (2) 𝑋 𝑘
(𝑁) 𝒳𝑘 = 𝑋 𝑘 𝑖 𝑖=1 𝑁 ∼ 𝜇 𝑘 𝑁 : N粒子の同時分布 𝝁 𝒌 (𝑵) on ℝ𝒅×𝑵 のポテンシャル関数: ただし (Fisher divergence) ただし ➢ 有限粒子ダイナミクスは を最小化するWasserstein勾配流. (近似的) 一様対数Sobolev不等式 [Chen et al. 2022] Remind [Chen, F., Ren, Z., and Wang, S. Uniform-in-time propaga- tion of chaos for mean field langevin dynamics. arXiv:2212.03050, 2022.] -
Transformerの特徴学習 • Transformer 55 [Brown et al. “Language Models are
Few- Shot Learners”, NeurIPS2020] [Alammar: How GPT3 Works - Visualizations and Animations, https://jalammar.github.io/how-gpt3-works-visualizations- animations/] LLM (e.g., GPT3) [Vaswani et al.: Attention is All you Need. NIPS2017] [Dosovitskiy et al.: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv:2010.11929. ICLR2021] Vision tasks (e.g., ViT) -
Transformerの表現力 • Yun et al. (2020); Zaheer et al. (2020):
有限長の系列 から系列への写像に関して万能近似性がある. • Edelman et al. (2022): スパースBoolean関数 • Gurevych et al. (2022): 階層的合成モデルによる判別 問題の汎化誤差解析 • P’erez et al. (2019): TransformerのTuring完全性 • Wei et al. (2021): Boolean回路,Turingマシンの近似 理論とその汎化誤差 (Rademacher複雑度) 56 [Yun, Bhojanapalli, Rawat, Reddi and Kumar. Are Transformers universal approximators of sequence-to-sequence functions? ICLR2020] [Zaheer, Guruganesh, Dubey, Ainslie, Alberti, Ontanon, Pham, Ravula, Wang, Yang, and Ahmed. Big Bird: Transformers for Longer Sequences. In Advances in Neural Information NeurIPS2020.] -
Transformerのトークン選択能力 57 定理 (推定誤差) ➢ 入力が無限次元でも多項式オーダーの収束レート. (ほぼミニマックス最適) ⋯ 𝑥−1 𝑥0
𝑥1 𝑥2 ⋯ ⋯ 𝑌−1 𝑌0 𝑌1 𝑌2 ⋯ ⋮ ⋮ ⋮ ⋮ Self-attention FNN Transformerの性質 • かなり広い文脈幅を用いて次トー クン予測をしている. → 次元の呪い? • 入力に依存して重要なトークンを 選択できる. → 次元の呪いを回避! [Shokichi Takakura, Taiji Suzuki: Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Input. ICML2023] 動的な変数選択 (CNNとの 大きな違い) 多項式オーダー -
トークンの重要度 58 𝑌 ⋯ 𝑌−1 𝑌0 𝑌1 𝑌2 ⋯ ⋯
𝑥−1 𝑥0 𝑥1 𝑥2 ⋯ 各トークンは重要度が異なる 重要度をその変数に関する滑ら かさとして定義する. 滑らかさ: 𝑎−1 𝑎0 𝑎1 𝑎2 重要な変数(トークン)は少しの変化で大き く出力を変えてしまう⇒滑らかじゃない 小さい𝑎𝑖: 非滑らか⇒重要 大きい𝑎𝑖: 滑らか⇒重要じゃない [Letarte et al.: Importance of Self-Attention for Sentiment Analysis. 2018] ➢ Attention層はこの重要なトークンを 動的に選択する. • 重要度は入力列に依存 ∃Π: ℝ𝑑×(2𝐻+1) → ℝ𝑑×(2𝐻+1) (並べ替え) s.t. 𝐹 𝑗 ∘ 𝑋 = 𝑓 ∘ Π 𝑋[𝑗−𝐻:𝑗+𝐻] , ただし𝑓 ∈ ℱ𝑝,𝑞 𝛾 : 𝛾-平滑関数クラス. (Πは,重要度順に並べる並べ替え) -
入力によって異なる重要度 59 𝐹∘ 𝑌 𝑋 This is a pen. これはペンです.
𝜇𝑖 𝑋 : 入力𝑋におけるークン𝑖の重要度 Π 𝑋 = [𝑋𝜋 1 , 𝑋𝜋 2 ,… ] ただし𝜇𝜋 1 𝑋 > 𝜇𝜋 2 𝑋 > ⋯ 𝐹∘ 𝑌 𝑋 This cat is cute この猫はかわいい. 𝐹 𝑗 ∘ 𝑋 = 𝑓 ∘ Π 𝑋[𝑗−𝐻:𝑗+𝐻] ※ 𝜇𝑖 は相対的な位置の重要度と出来る (e.g., 𝜇𝑖−𝑗 (𝑋)). ⋯ 𝑥−1 𝑥0 𝑥1 𝑥2 ⋯ 𝑥0 𝑥1 𝑥2 𝑥−1 ⋯ Smoothness 𝛾-平滑関数クラス -
Soft-max Attentionの代替 60 > 能力 < 効率性 Attention 代替手法 (線形手法)
完全な代替は難しい Q.どのようなタスクなら代替可能? 主結果:入力依存な特徴抽出が必要な関数の推定では同等の能力 重要な場所が変化する画像の分類 [Takakura & Suzuki, 2023] 入力に応じた記憶の呼び起こし • 計算量 𝑂 𝐿2 (𝐿は入力長) • 基盤モデルのスタンダード • 計算量 𝑂 𝐿 or 𝑂 𝐿log(𝐿) • 能力の限界に関する指摘 言語等の離散データが不得意 -
代替手段 (線形手法) 61 (SSM) exp(𝑘𝑗 ⊤𝑞𝑖 ) = σ 𝑚=0
∞ 𝜙𝑚 (𝑘𝑗 )𝜙𝑚 (𝑞𝑖 ) ≈ σ 𝑚=0 𝑀 𝜙𝑚 (𝑘𝑗 )𝜙𝑚 (𝑞𝑖 ) = 𝜙(𝑘𝑗 )⊤𝜙(𝑞𝑖 ) 𝑦𝑗 = 𝜓2 (𝑢𝑗 ) ⊙ σ 𝑛=0 ∞ ℎ𝑛 ⋅ 𝜓1 (𝑢𝑗−𝑛 ) σ 𝑛=0 ∞ 𝜓1 (𝑢𝑗−𝑛 )⊤ ෨ 𝜓2 (𝑢𝑗 ) → ෨ 𝜓1 (𝑢𝑗−𝑛 ) 線形注意: SMM+gating: カーネル関数の有限和近似 -
SSMの表現力 62 ? 入力依存で重要なトークンを抽出可能する必要がある →重要度(入力依存) e.g. 自分 or 1つ前が最後のトークンと同じ 各トークンの重要度を前の層で計算すれば実は代替できる:
p 1 u 7 v 4 w 7 t 9 u 1 3 20 159 2 2 4 3 2 3 24 結果1:多層の FNN + SSM で Transformer を代替可能 結果2: 区分的𝛾-平滑関数の推定において,SSM はTransformerと同じ推定誤差を達成 事実: Copyingタスクにおいて1層のSSMはTransformerを代替しにくい [Jelassi et al.: Repeat After Me: Transformers are Better than State Space Models at Copying. 2024] [Naoki Nishikawa, Taiji Suzuki: State Space Models are Comparable to Transformers in Estimating Functions with Dynamic Smoothness. ICLR2025] しかし... (associative recall) -
63 [Armin W. Thomas, Rom Parnichkun, Alexander Amini, Stefano Massaroli,
Michael Poli: STAR: Synthesis of Tailored Architectures. ICLR2025, oral] 実験的にもSSMとTransformerの混合モデルは有効 - 完全な代替は出来ないが部分的代替は有効 -
学習レジームの多様化 66 事前学習 事前学習 事後学習 テスト時 推論 100% 45% 35%
20% 事前学習データの質向上 Data augmentation ここの重要度が上がっている 含,データの自動生成 ➢o3, AlphaProof, DeepSeek •Alignment •SFT (supervised fine- tuning) •Preference optimization •RLHF, RLAIF •Monte-Carlo Search •In-context learning (Few-shot prompting) •Chain-of-thought -
67 事前学習 事後学習 テスト時 推論 事前学習データの質向上 Data augmentation •Alignment •SFT
(supervised fine- tuning) •Preference optimization •RLHF, RLAIF •Monte-Carlo Search •In-context learning (Few-shot prompting) •Chain-of-thought -
69 • 解候補をモンテカルロサンプリングして良いものだけをピックアップ • 解候補の「良さ」を測るProcess Reward Verifier (PRM) も学習 →
枝刈り・推論を高速化 テストタイムの時間を多くとった方が性能向上 (たくさんサンプリングし た方が良い出力を見つけられる) → 見つかった良い結果を用いてモデルをfine-tuningすることも可能 → 候補の生成に費やした計算量も考慮すべき:新しいスケーリング則 -
テスト時スケーリングの理論 (BoN) • BoN戦略から生成される出力の分布と元分布との差 [Beirami et al. 2025] 70 𝜋ref
𝜋BoN 𝑟 (報酬, reward) • 𝑥: プロンプト入力, 𝑦: 回答 • Best-of-N (BoN) • BoNとKL正則化の同値性 [Yang et al., 2024; Mroueh, 2025] つまり,𝑁を大きくしていくことで,期待報酬が最大化されていく. しかし,BoNは報酬モデルに強く依存するため 𝑟 が間違っていれば過適合してしまう. (reward hacking) ➢ KL-divの意味で -
BoNと過剰適合 (reward hacking) • Rewardに誤差があるときのBoNのregret bound 71 • 理論上は,KL-ダイバージェンスではなく,𝜒2-ダイバージェンスを用いることで𝑁を外せる [Huang
et al. 2024] • 報酬関数の改善に能動学習的なサンプリングをすることで探索のコストを減らせる [Foster et al. 2025] 記号の定義は省略 𝑁を増やすことでより高い 報酬の出力が見つかる 自身の報酬モデル を過信することで 生じる損失 ➢最悪の状況では,この上限はタイト: ある𝑟, 𝜋ref , Ƹ 𝑟が存在して ∵ 𝜋ref と𝜋∗の密度比が大きい所の報酬関数の誤差はコントロールしにくい. [Huang et al.: Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment. ICML2025] 報酬モデルの誤差 -
論理推論パスの多様性は重要 72 [Dake Bu, Wei Huang, Andi Han, Atsushi Nitanda,
Bo Xue, Qingfu Zhang, Hau-San Wong, Taiji Suzuki: Diversity Matters: A Comparative Theory of Reward-based Finetuning and Inference-Scaling. 2025] • 論理タスクを解く段階で「簡単なステップ」と「難しいステップ」がある. • 事前学習モデルが「難しいステップ」を通る確率は低い.しかし,全てのタスクを解けるように なるには「難しいステップ」も通れるようにしておく必要がある. [理論の結果] • ある特定のタスクでKL-正則化なしで出力報酬 (outcome reward) を勾配法で直接最大化しようと すると「難しいステップ」を忘れる ⇒ 他のタスクが解けなくなる. • 解けるようになった問題のデータは棄却し,まだ解けない問題の学習データ上で学習することで 上記の問題は避けることができる.[カリキュラム学習] • KL-正則化を適切に入れることでも上記の問題を緩和できる.[モデルがたどれるパスの多様性を確保す ることで多様な問題を解く能力を確保する] -
73 Pre-training Post-training Test time inference •Enhancing quality of data
for pre-training •Data augmentation •Alignment •SFT (supervised fine- tuning) •Preference optimization •RLHF, RLAIF •Monte-Carlo Search •In-context learning (Few-shot prompting) •Chain-of-thought • [Chain-of-thought] Kim&Suzuki: Transformers provably solve parity efficiently with chain of thought. ICLR2025, oral. -
思考連鎖 (Chain-of-Thought) 74 • 思考の連鎖を訓練データに用いる • 思考の連鎖を例示してin-context learningさせる • 思考の連鎖を出力させる(e.g.,
think step by step) →精度向上,解釈性向上 • 結果だけでなく思考過程も出力/入力 [Wei et al.: Chain-of-Thought Prompting Elicits Reasoning in Large Language Models. 2022] -
数学への応用 76 • AlphaGeometry (DeepMind, 2023) • AlphaProof, AlphaGeometry2 (DeepMind,
2024) • 定理証明系の言語を利用:Lean [de Moura et al., 2015], Coq [Barras et al., 1997], Isabelle [Nipkow et al., 2002] • 定理を「形式化」して,証明をプログラムとして書き下す. • 証明の真偽は自動的に判定可能(単発の回答はもちろん真偽判定可能) ➢ 思考連鎖の訓練データを収集して学習 ➢ 思考連鎖を自動生成して証明が通ったものを訓練データにして学習 -
思考連鎖の理論 77 [Kim&Suzuki: Transformers provably solve parity Efficiently with chain
of thought. ICLR2025 (oral), arXiv:2410.08633] 𝑘-パリティ問題 ➢ 𝑥 = (𝑥1 , … , 𝑥𝑑 ) ∼ Unif( −1,1 𝑑) ➢ 𝑦 = 𝑥𝑖1 𝑥𝑖2 ⋯ 𝑥𝑖𝑘 = ς𝑗∈𝑝 𝑥𝑗 𝒅-次元要素中の𝒌個の要素の積が𝒚 多項式時間アルゴリズムのサンプル複雑度の下限: 𝑛 = Ω(𝑑𝑘−1) Q: 思考連鎖で改善できるか? 𝑘-パリティ問題のNNによる訓練に関する理論は豊富 Abbe et al. (2023); Refinetti et al. (2021); Ben Arous et al. (2022); Damian et al. (2022); Suzuki, Wu, Oko, Nitanda (2023). この判別問題を𝑛個のデータから学習: 𝒙𝒊, 𝒚𝒊 𝒊=𝟏 𝒏 -
𝑘-パリティ問題の階層 78 : Transformerに中間結果を逐一出力させる (think step by step) 訓練時に最終結果だけでなく中間的結果も出 力するようTransformerを訓練:
• 入力の各要素: 𝑥𝑖 (𝑖 = 1, … , 𝑑) • 中間結果: 𝑥𝑗 (𝑗 ≥ 𝑑 + 1) ➢ 𝑗 ≥ 𝑑 + 1の中間結果を次トークン予測で逐 次的に出力 ➢ 最終トークンが出力の𝑦 例:𝑥17 = 𝑥1 × 𝑥4 2パリティ問題に分解 -
𝑘-パリティ問題の階層 79 訓練時に最終結果だけでなく中間的結果も出 力するようTransformerを訓練: • 入力の各要素: 𝑥𝑖 (𝑖 = 1,
… , 𝑑) • 中間結果: 𝑥𝑗 (𝑗 ≥ 𝑑 + 1) ➢ 𝑗 ≥ 𝑑 + 1の中間結果を次トークン予測で逐 次的に出力 ➢ 最終トークンが出力の𝑦 𝑥1 𝑥2 𝑥16 𝑥17 𝑥23 𝑦 𝑥𝑚 ⋯ ⋯ Next token prediction ⋯ -
教師信号強制 (Teacher forcing) 80 𝑒1 𝑥1 Position encoding 𝑒2 𝑥2
𝑒𝑑 𝑥𝑑 ⋯ 𝑒𝑑+1 𝑥𝑑+1 𝑒𝑑+𝑘−1 𝑥𝑑+𝑘−1 ⋯ 𝑦 𝑒𝑚 𝑥𝑚 ⋯ 予測 • 𝑥 = (𝑥1 , … , 𝑥𝑑 ) ∼ Unif( −1,1 𝑑) • 𝑦 = 𝑥𝑖1 𝑥𝑖2 … 𝑥𝑖𝑘 = ς𝑗∈𝑝 𝑥𝑗 𝑥1 = 𝑥2 = -
教師信号強制 (Teacher forcing) 81 𝑒1 𝑥1 Position encoding 𝑒2 𝑥2
𝑒𝑑 𝑥𝑑 ⋯ 𝑒𝑑+1 𝑥𝑑+1 𝑒𝑑+𝑘−1 𝑥𝑑+𝑘−1 ⋯ 𝑦 𝑒𝑚 𝑥𝑚 ⋯ 予測 • 𝑥 = (𝑥1 , … , 𝑥𝑑 ) ∼ Unif( −1,1 𝑑) • 𝑦 = 𝑥𝑖1 𝑥𝑖2 … 𝑥𝑖𝑘 = ς𝑗∈𝑝 𝑥𝑗 𝑥1 = 𝑥2 = 勾配降下によって次のトークンを予測 するのに,過去のどのトークンを利用 するべきかを学習する必要がある. ⇒特徴学習 -
教師信号強制のもとでの学習可能性 82 訓練データサイズを 𝑛 = Ω 𝑑2+𝜖 , とし,学習率を𝜂 =
Θ(𝑑2+ Τ 𝜖 16)とすれば,1ステップの勾配降 下更新による学習により,テスト損失が以下のようになる: ො 𝑦test − 𝑦test ∞ ≤ O 𝑑− Τ 𝜖 8 . 定理 (教師信号強制のもとでのサンプル複雑度) w/o CoT with CoT 訓練データサイズ Ω(𝑑𝑘−1) O(𝑑2+𝜖) サンプル複雑度の比較 1ステップ勾配降下: ただし,𝑊0 = 𝑂. (勾配は近似的勾配で代替可能: ෩ ∇𝐿 s.t. ෩ ∇𝐿 − ∇𝐿 ≤ 𝑂(𝑑−2−𝜖/8)) 中間結果𝑥𝑚 はテスト時にはモデルの予測値 ො 𝑥𝑚 を用いる. -
数値実験 83 データ数の比較 思考連鎖 三つの異なる方策: - 普通の思考連鎖 - 教師信号強制 -
教師信号強制なし,ただし 段階的なカリキュラム学習 を実施 Grokking (訓練誤差が落ち切ってから テスト誤差が一気に下がる) -
84 Pre-training Post-training Test time inference •Enhancing quality of data
for pre-training •Data augmentation •Alignment •SFT (supervised fine- tuning) •Preference optimization •RLHF, RLAIF •Monte-Carlo Search •In-context learning (Few-shot prompting) •Chain-of-thought [In-context learning] Optimization & feature learning: • Oko, Song, Suzuki, Wu: Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In- Context. Advances in Neural Information Processing Systems 37 (NeurIPS 2024). pp. 77316--77365, 2024. • Naoki Nishikawa, Yujin Song, Kazusato Oko, Denny Wu, Taiji Suzuki: Nonlinear transformers can perform inference-time feature learning: a case study of in-context learning on single-index models. 2025. Minimax optimality: • Juno Kim, Tai Nakamaki, Taiji Suzuki: Transformers are Minimax Optimal Nonparametric In-Context Learners. Advances in Neural Information Processing Systems 37 (NeurIPS 2024). pp. 106667--106713, 2024. Best paper award, ICML 2024 Workshop on Theoretical Foundations of Foundation Models. -
In-context learning (文脈内学習) 85 In-Context Learning (ICL) [Brown et al.,
2020]. 良く事前学習されたモデルはテスト時のin-context learning (文脈内学習) でも良い性能を示す. Question ChatGPT -
In-context learning (文脈内学習) 86 Question ChatGPT In-Context Learning (ICL) [Brown
et al., 2020]. 良く事前学習されたモデルはテスト時のin-context learning (文脈内学習) でも良い性能を示す. -
ファインチューニングとの違い 87 通常のファインチューニングはパラメータを更新する. In-context learningでは更新しない. (e.g., RLHF) ※ 最近では,test time
trainingと言って,LoRAを用いてin-context learning時に少し ファインチューニングする方法も提案されている [Akyurek et al. arXiv:2411.07279]. -
事前学習によるメタ学習 88 Pretraining Test task Example Query Query Examples ICLはモデルパラメータを更新しない
→ メタ学習,学習の学習 事前学習時にたくさんのタスク を観測しておく. → タスク汎化 -
ICL (文脈内学習) の数学的定式化 89 事前学習 (𝑻タスク): ⋯ × 𝑇 ➢
𝑇個のタスクを観測 ➢ 各タスクで𝑛個の例示 を観測 テストタスク (文脈内学習): ⋯ 予測 • 真の関数𝑓∗ 𝑡はタスクごとにランダムに生成されるとする. • 𝑇タスク,各タスクで𝑛個の訓練データを観測. モデル: 𝑡 = 1, … , 𝑇: Task index -
Transformerの役割 • 事前学習 (pretraining): 特徴量 (表現) を学習 [𝑓∘] ➢Fourier, B-Spline
➢文脈 (𝑡) に非依存 ➢データを表現する「最も効率的」な基底を学習 → 中間層 • 文脈内学習 (in-context): 係数を学習 [𝛽𝑡 ] ➢文脈 (𝑡) に依存 ➢例示から現在の文脈𝛽𝑡 を推定 → 最終層のAttention 90 ✓ Guo et al. (2023), von Oswald et al. (2023) では,Transformerは浅い層で特徴量を 抽出して深い層で線形回帰 (or 勾配法) を行っていることを実験的に確認. -
ICLリスクとベイズ予測 91 期待予測リスク : 経験リスク : 𝑛個の観測データをもとに構成した推定量𝑓Θ の予測誤差 を最小化させる ⇒
なるべく最適な学習法を学習する:メタ学習 「学習法の学習」 これはベイズ推定量に他ならない: 𝑥に関する期待値 given 𝜃 この右辺を最小化する መ 𝑓が ベイズ推定量の定義 [ベイズリスク最小化] 損失 -
Transformerは線形回帰を実装できる 92 [Gang et al. 2022; Akyurek et al. 2023;
Zhang et al. 2023; Ahn et al. 2023; Mahankali et al., 2023; Wu et al. 2024] [Gang et al.: What Can Transformers Learn In-Context? A Case Study of Simple Function Classes. NeurIPS2022] [Zhang, Frei, Bartlett.: Trained Transformers Learn Linear Models In-Context. JJMLR, 2024] クエリ KQ行列 バリュー・キー プロンプトの例示 -
Transformerによる勾配法の実現 93 • Transformerは勾配法による文脈内学習を内部的に実装することができる. [勾配法の各更新が一層分に対応] • Transformerは交差検証法によるアルゴリズムの選択を内部的に実装できる. [Yu Bai, Fan
Chen, Huan Wang, Caiming Xiong, Song Mei: Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection. NeurIPS2023] ✓ 実際には,浅い層で特徴抽出をし,深い層で勾配法を実行しているらしい [Guo et al., 2023; von Oswald et al., 2023]. -
予測誤差の評価 94 経験リスク最小化: 定理 (ICL risk bound; Kim, Nakamaki, TS,
NeurIPS2024) 1. 2. Feature approximation error Pretraining generalization to estimate basis functions 3. In-context generalization gap 4. 𝑓𝑗 ∘ 𝑗=1 ∞ are “near” orthonormal 仮定 (informal) Covering number of DNN (関数空間の複雑さ) (基底の近似誤差) (基底は大きすぎない) (基底はほぼ正規直交) -
予測誤差の収束レート (簡略版) 95 • 例 (B-spline基底; 𝑓𝑗 ∘がB-spline→Besov/Sobolev空間): 𝑻が小さい: 記憶中の状況
𝑻が大きい: 記憶が完了し汎化できる状況 → テスト時推論のスケーリング則 「文脈 (𝛽𝑡 )」の推定誤差 ミニマックス最適 「最適な表現 (𝑓∘)」 +「最適な𝛽𝑡 の推定方法」 を獲得するための複雑さ -
タスク多様性と性能の関係 96 [Raventós, Paul, Chen, Ganguli: Pretraining task diversity and
the emergence of non-Bayesian in-context learning for regression. 2023 ] If # of pretraining tasks is enough, ICL coincides with optimal ridge regression. -
理論解析の詳細:真のモデル 97 where 𝛽𝑡 ∼ 𝑁(0, Σ) and 𝑓∘ 𝑥
∈ ℝ∞. Suppose that the true function admits a basis function decomposition: • B-Spline (Besov) • Fourier (Sobolev, 𝛾-smooth) Tensor product B-spline: 𝛾-smooth function class for 𝑑 = ∞ [Okumoto&Suzuki,ICLR2022], [Takakura&Suzuki, ICML2024] -
モデルの詳細 98 [Ahn et al.: Linear attention is (maybe) all
you need (to understand transformer optimization). arXiv:2310.01082] 2. 線形attention 予測 Query Key Value ⋯ 𝑌1 𝑌2 𝑌𝑛 𝑥1 𝑥2 𝑥𝑛 FNN Attention 𝑦1,𝑡 𝜙(𝑥1,𝑡 ) 𝑦𝑖,𝑡 𝜙(𝑥𝑖,𝑡 ) 𝑦𝑛,𝑡 𝜙(𝑥𝑛,𝑡 ) ? 𝜙(𝑥qr,𝑡 ) ⋯ ⋯ Query Key Value Prompt Attention 1. Soft-max attention 特徴量𝑓∘を中間層で表現: 深層ニューラルネット (非線形特徴マップ) ※ 実際は,各トークンは(𝜙 𝑥 , 𝑦)の組とすべきだが,理論研究では𝑄, 𝐾, 𝑉に特殊な形を想定して,このような形に「簡素化」することが多い. Γ ここではこっちを考える -
ICLリスク 99 期待ICLリスク: (where 𝑓∗ 𝑡(𝑥) = 𝛽⊤𝑓∘(𝑥)) 経験的ICLリスク :
→ 𝝓と𝚪について最小化 c.f. リッジ回帰: ≃ Γ 𝛽の推定量となる [Gang et al. 2022; Akyurek et al. 2023; Zhang et al. 2023; Ahn et al. 2023; Mahankali et al., 2023; Wu et al. 2024] Γを適切に選べばベイズ最適な推定量となる. -
事前学習の最適化保証ありの特徴学習 100 Optimization & feature learning: • Oko, Song, Suzuki,
Wu: Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In- Context. Advances in Neural Information Processing Systems 37 (NeurIPS 2024). pp. 77316--77365, 2024. • Naoki Nishikawa, Yujin Song, Kazusato Oko, Denny Wu, Taiji Suzuki: Nonlinear transformers can perform inference-time feature learning: a case study of in-context learning on single-index models. ICML2025. -
問題設定(再掲) 101 事前学習 (𝑻タスク): ⋯ × 𝑇 ➢ 𝑇個のタスクを観測 ➢
各タスクで𝑛個の例示 を観測 テストタスク (文脈内学習): ⋯ 予測 • 真の関数𝑓∗ 𝑡はタスクごとにランダムに生成されるとする. • 𝑇タスク,各タスクで𝑛個の訓練データを観測. モデル: 𝑡 = 1, … , 𝑇: Task index -
教師モデルの詳細 ガウシアン single index モデル: 部分空間𝒮と基底関数He𝑖 を事前学習において 学習しておくことが重要 ただし,𝜎∗ 𝑡
はリンク関数で真の方向 𝛽𝑡 はランダムに生成される: 𝜷𝒕 ただし,𝑐𝑖 𝑡 はランダムに生成されていて以下を満たす: は𝑟次元線形部分空間𝒮の単位球面上で一様に分布. 𝑟 ≪ 𝑑を仮定. ⇒ 情報指数 = 𝒌. 𝒮 有効な特徴量の空間は低次元 𝝈∗ 𝒕 -
教師モデルの詳細 ガウシアン single index モデル: 部分空間𝒮と基底関数He𝑖 を事前学習において 学習しておくことが重要 ただし,𝜎∗ 𝑡
はリンク関数で真の方向 𝛽𝑡 はランダムに生成される: 𝜷𝒕 ただし,𝑐𝑖 𝑡 はランダムに生成されていて以下を満たす: は𝑟次元線形部分空間𝒮の単位球面上で一様に分布. 𝑟 ≪ 𝑑を仮定. ⇒ 情報指数 = 𝒌. 𝒮 有効な特徴量の空間は低次元 𝝈∗ 𝒕 𝛽𝑡 𝜎∗ 𝑡 -
• FNN層 : 生徒モデル (2層Transformer) 104 [Ahn et al.: Linear
attention is (maybe) all you need (to understand transformer optimization). arXiv:2310.01082] • 線形attention: 予測 Query Key Value 𝑦1,𝑡 𝑓𝑊,𝑏 (𝑥1,𝑡 ) 𝑦𝑖,𝑡 𝑓𝑊,𝑏 (𝑥𝑖,𝑡 ) 𝑦𝑛,𝑡 𝑓𝑊,𝑏 (𝑥𝑛,𝑡 ) ? 𝑓𝑊,𝑏 (𝑥qr ) ⋯ ⋯ Query Key Value Prompt 線形attention (𝜎: ReLU) (線形回帰) 𝑦𝑛+1 𝑥1 𝑦1 FNN Attention 𝑥2 𝑦2 𝑥𝑛+1 ∗ c.f., soft-max attention -
In-Context Learning (ICL) risk 105 期待予測誤差: 経験損失: → 𝑾, 𝒃
(特徴写像) と 𝚪 (attention) について最小化 (note that 𝑦𝑖,𝑡 = 𝑓∗ 𝑡 𝑥𝑖,𝑡 + 𝜖𝑖,𝑡) (大サンプルサイズ極限: 𝑛 → ∞ かつ 𝑇 → ∞) 勾配法による学習で 予測誤差を小さくできるか? ≃ Γ (事前情報) 線形Attentionはリッジ回帰を実装できる: [Gang et al. 2022; Akyurek et al. 2023; Zhang et al. 2023; Ahn et al. 2023; Mahankali et al., 2023; Wu et al. 2024] -
最適化アルゴリズム (勾配法) 106 •ステージ1: One-step 勾配降下. •ステージ2: 𝚪の最適化. 初期化: 𝒘
𝑗 (0) ∼ Unif 𝕊𝑑−1 , 𝑏𝑗 = 0, Γ 𝑗,𝑗 0 = Unif {±1} (diagonal). 𝑏𝑗 をランダムに再初期化:𝑏𝑗 ∼ Unif −1,1 . ステージ1で学習した𝑾をもとに,Γを最適化 (Γ については凸): 𝑾をone-stepの勾配降下で更新: ➢ 2層NNの学習におけるone-step勾配降下の解析を援用 [Damian et al. 22; Ba et al. 22]. 部分空間𝒮が見つかる Attentionを訓練することで, 𝛽𝑡 を抽出できるようになる. -
勾配法による特徴学習 107 𝒮 • 1ステップ勾配降下の更新 (+正則化)により,ランダム初期値𝑤 𝑗 (0)が部分空 間𝒮に射影される形になる (大きさの自由度を除き).
𝒘 𝒋 (𝟏) • 十分多くのニューロンがあれば, 𝑤 𝑗 1 𝑗=1 𝑚 は部分空間𝒮を張る (1st -stage). • 同じく十分多くのニューロンがあれば,それらの線形結合で対象の関数 𝜎𝑡 ∗(⟨𝛽𝑡 , 𝑥⟩)を十分よく近似できる (2nd-stage + test prompt). • 𝑊の学習:部分空間𝒮 が得られる. • Γの学習:基底関数の 係数を文脈から推定す る方法を学習. 𝒘 𝒋 (𝟎) -
主結果 108 定理 (文脈内学習の予測誤差) 𝑛∗ をテストタスクの例示数とする. 今,事前学習データが十分大きく, 𝑇1 = Θ(𝑑𝑘+1)
and 𝑛 = ෩ Ω(𝑑𝑘), を満たし,𝑚 ≫ 𝑛 かつ𝑇2 ≫ 𝑛を満たすなら予測誤差は以下のように 抑えられる: 𝑚: NNの横幅, 𝑇1: 事前学習ステージ1のタスク数 (𝑊の学習用), 𝑇2: 事前学習ステージ2のタス ク数 (Γの学習用), 𝑛: 事前学習の各タスクで提示される例示数. • 事前学習無しでは,カーネル法なら𝑛∗ = Ω(𝑑𝑝),バッチ再利用する勾配法を 用いればNNで𝑛∗ = Ω(𝑑)必要である. • しかし,事前学習後の文脈内学習の複雑度は𝑑に依存しない (𝑛∗ = poly (𝑟)). ➢事前学習時にコンパクトな特徴量を得ていることの恩恵. ➢ここで,特徴学習が失敗しているとうまく汎化しない. -
主結果 109 定理 (文脈内学習の予測誤差) 𝑛∗ をテストタスクの例示数とする. 今,事前学習データが十分大きく, 𝑇1 = Θ(𝑑𝑘+1)
and 𝑛 = ෩ Ω(𝑑𝑘), を満たし,𝑚 ≫ 𝑛 かつ𝑇2 ≫ 𝑛を満たすなら予測誤差は以下のように 抑えられる: 𝑚: NNの横幅, 𝑇1: 事前学習ステージ1のタスク数 (𝑊の学習用), 𝑇2: 事前学習ステージ2のタス ク数 (Γの学習用), 𝑛: 事前学習の各タスクで提示される例示数. ➢事前学習時にコンパクトな特徴量を得ていることの恩恵. ➢ここで,特徴学習が失敗しているとうまく汎化しない. • 事前学習無しでは,カーネル法なら𝑛∗ = Ω(𝑑𝑝),バッチ再利用する勾配法を 用いればNNで𝑛∗ = Ω(𝑑)必要である. • しかし,事前学習後の文脈内学習の複雑度は𝑑に依存しない (𝑛∗ = poly (𝑟)). 事前学習無し 事前学習有り 手法 Kernel NN (SQ) ICL サンプル複雑度 𝑑𝑃 𝑑 ⋅ polylog(𝑑) 𝑟4𝑃 事前学習 --- --- 𝑇1 = 𝑑𝑘+1, 𝑛 = 𝑑𝑘 -
GPT-2による数値実験 110 Test error Number of instruction examples ICL Non-ICL
• 𝑟を固定 • 𝑑だけ変化 NN without pre-training is affected by 𝑑. ICL by transformer is not affected by 𝑑. (thanks to feature learning in pre-training) -
数値実験2 111 Fixing d, changing r Fixing r, changing d
GPT2 model with 12-layers (∼22M parameters) Only 𝑟 affects the test error while 𝑑 does not. -
これで十分か? テスト時のサンプル複雑度: 𝑟O(𝑃) ⇒ これは特徴量を次元𝑟まで減らした後のカーネル法と同じ誤差 二つの問題点: A) できれば𝑟𝑃は𝑟𝑘まで減らしたい. B) さらにできれば,情報理論的下限෩
Θ(𝑟)まで減らしたい. 実は... ソフトマックスattention を使うことで, 𝒓ge(𝝈∗) まで減らせる. ⇒ これはテスト時推論においても特徴学習できることを意味する. 112 ge(𝜎∗): 生成的指数(多項式の場合1 (奇関数),2 (偶関数)) 事前学習無し 事前学習有り 手法 Kernel NN (CSQ or SQ) ICL サンプル複雑度 𝑑𝑃 𝑑 ⋅ polylog(𝑑) 𝑟4𝑃 事前学習 --- --- 𝑇1 = 𝑑𝑘+1, 𝑛 = 𝑑𝑘 -
問題設定 113 モデル: 𝑡 = 1, … , 𝑇: Task
index 仮定: where • Remark: ここでは,𝜎∗はタスクによらず一定とする. 定義 (生成的指数 [Damian et al. 2024]) 𝜎∗ の生成的指数は,以下のように定義される: 𝜎∗ が多項式なら: -
Soft-max Transformerモデル 114 • Softmax attention層 (単一ヘッド): 𝑦qr 𝑥1 𝑦1
FNN Softmax attention 𝑥2 𝑦2 𝑥qr ∗ • 全結合層: (𝜎: ReLU) • 非線形変換 𝑦 ↦ 𝑦exp(𝑦/𝜌) が情報指数を生成的指数まで下げる. • さらに,Attention層がテスト時に⟨𝑥qr , 𝛽𝑡 ⟩を抽出できることが示せる. -
結果のまとめ 115 定理 (Soft-max attentionによる文脈内学習の予測誤差) 1. (最適化誤差) 𝑇1 = ෩
Ω(𝑟2𝑑IE(𝜎∗)+2), 𝑛 = ෩ Ω (𝑟2𝑑IE(𝜎∗)+2), 𝑇2 = ෩ Ω(𝑟3ge(𝜎∗)/2), なら,訓練誤差は𝑜𝑑 1 . 事前学習無し 事前学習有り 手法 Kernel NN ICL/線形注意 ICL/非線形注意 サンプル複雑度 𝑑𝑃 𝑑Θ(ge 𝜎∗ ) 𝑟4𝑃 𝑟3ge(𝜎∗)/2 事前学習 --- --- 𝑇1 = 𝑑𝑘+1, 𝑛 = 𝑑𝑘 𝑇1 = 𝑟2𝑑𝑘+2, 𝑛 = 𝑟2𝑑𝑘+2 2. (テスト時のサンプル複雑度) 𝑛∗ = ෩ Ω (𝑟3ge(𝜎∗)/2)なら,テスト時の予測誤差は𝑜𝑑 1 . 事前学習時の特徴学習により𝑑 → 𝑟へ表現の次元圧縮 テスト時にもAttentionが特徴学 習をして𝑃 → ge(𝜎∗)へ改善 ge(𝜎∗): 生成的指数 (多項式の場合1 (奇関数),2 (偶関数)) -
拡散モデル 117 「An astronaut riding a horse in a photorealistic
style」 「Teddy bears shopping for groceries in the style of ukiyo-e」 SORA (OpenAI, 2024) DALL·E: [Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever: Zero-Shot Text-to-Image Generation. ICML2021.] DALL·E2:[Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, Mark Chen: Hierarchical Text-Conditional Image Generation with CLIP Latents. arXiv:2204.06125] -
離散拡散モデルと言語生成 118 [Nie et al. Large Language Diffusion Models. arXiv:2502.09992]
[Google Deepmind: Gemini Diffusion, 2025] -
拡散モデルの定式化 119 順過程:ターゲットの分布を正規分布に変換していく (OU-過程). 逆過程:正規分布 (ノイズの分布) から逆にたどってターゲットの分布に逆変換していく. [Vahdat, Kreis, Kautz:
Score-based Generative Modeling in Latent Space. arXiv:2106.05931] (𝑌𝑡 ∼ 𝑋 𝑇−𝑡 ) [Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Song et al., 2020; Ho et al., 2020; Vahdat et al., 2021] Estimated from data Reverse (target distribution) (Standard Normal) Forward (Wasserstein GF) KL-divergence from 𝜇∗ -
順過程 120 順過程: ただし,𝜇𝑡 = exp −𝑡 , 𝜎𝑡 2
= 1 − exp −2𝑡 . OU-過程 GLDの一般論より,順過程は指数関数的に標準正規分布に近づく. [Vahdat, Kreis, Kautz: Score-based Generative Modeling in Latent Space. arXiv:2106.05931] 𝑝𝑡 を𝑋𝑡 の確率密度関数とする. 形がわかっている! 𝒙𝟎 が与えられれば𝒙𝒕 の サンプリングも可能 元の分布 標準正規分布 OU-過程 -
逆過程 122 逆過程: [Haussmann & Pardoux, 1986] 事実:𝑌𝑡 の分布=𝑋ത 𝑇−𝑡
の分布 順過程を逆にたどることによって,(ほぼ)正規分布に従うノイズ を徐々に修正して元の画像の分布を再現できる. (𝑡 ∈ [0, ത 𝑇]) すなわち,𝑌𝑡 ∼ 𝑝 𝑇−𝑡 𝑌0 ∼ 𝑝 𝑇 こっちから始める こっちで終わる -
拡散モデルの図示 123 Reverse process (target distribution) (Gaussian distribution) Forward process
(Wasserstein GF) KL-divergence from 𝜇∗ -
OU過程からの修正 124 𝑇 OU-過程: OU-過程を逆過程で用いれば,常にガウス分布のまま. 当然,ターゲットの分布𝑝0 には収束しない. 修正した 確率過程: どうやって“修正”
すればよいか? (𝑣𝑡 をどう求める?) 𝑣𝑡 OU-process Reverse process (modification of OU) 𝑝0 -
最小エネルギー修正としての定式化 125 OU-過程 ただし,෨ 𝑌ത 𝑇 ∼ 𝑝0 , 𝑌0
∼ ෨ 𝑄0. ただし,𝑌ത 𝑇 ∼ 𝑝0 , 𝑌0 ∼ 𝑝ത 𝑇. with 𝑌0 ∼ 𝑁(0, 𝐼). 定理 (Vargas, Grathwohl and Doucet, 2023; adapted) • 通常用いられる逆過程はOU-過程からのKL-divを最小化している( ෨ 𝑌ത 𝑇 ∼ 𝑝0 のもと) . • さらに,それは修正項の「エネルギー」∫ 𝒗𝒕 𝟐𝒅𝒕 (+ 初期分布間のKL) を最小化 している. 上記の形で書ける任意の確率過程 ෨ 𝑌𝑡 (with ෨ 𝑌ത 𝑇 ∼ 𝑝0)に対し以下が成り立つ: (条件:ターゲット分布に収束) 𝑌𝑡 𝑡=0 ത 𝑇 の標本路の分布 通常の逆過程 別の逆過程 -
最小作用の修正 126 𝑣𝑡 OU-process (𝑷𝐫𝐞𝐟 ) Reverse process (modification of
OU) 𝑝0 𝒗𝒕 : OU-過程からの修正分 𝑌0 𝑌ത 𝑇 最小エネルギー ≤ exp(−2ത 𝑇) -
スコアの推定 127 逆過程: ⇒ 𝑌𝑡 ∼ 𝑝 𝑇−𝑡 [Haussmann &
Pardoux, 1986] (未知) ⇒ スコア関数𝛻log(𝑝𝑡 )をできるだけ正確に推定できれば良い:スコアマッチング 近似モデル (生成モデル): (未知) (𝑝 𝑇 は𝑁(0, 𝐼)に十分近い) 𝑌ത 𝑇 を生成画像として用いる. (𝑡 ∈ [0, ത 𝑇]) (𝑡 ∈ [0, ത 𝑇]) 定理 (Girsanov’s theorem) ≤ exp(−2ത 𝑇) -
スコアマッチング 128 観測値 (𝑛データ点, 𝐷𝑛 = 𝑥𝑖 𝑖=1 𝑛 ):
経験スコアマッチング損失: 陽に求まる (正規分布の密度より) 条件付分布はOU過程からサンプリングできる (正規分布) -
is sufficiently smooth on the edge of the support Assumption
129 Assumption 1 The true distribution 𝑝0 is supported on −1,1 𝑑 and with 𝑠 > Τ 1 𝑝 − Τ 1 2 + as a density function on −1,1 𝑑. Assumption2 Very smooth Besov space Besov space (𝐵𝑝,𝑞 𝑠 (Ω)) Smoothness Spatial inhomogeneity Reference -
(𝑌𝑡 ∼ 𝑋 𝑇−𝑡 ) 拡散モデルの統計理論 130 Stable diffusion, 2022.
Forward process Backward process どちらも(ほぼ)ミニマックス最適 [Yang & Barron, 1999; Niles-Weed & Berthet, 2022]. 経験スコアマッチング推定量: (for any 𝛿 > 0). 定理 Let 𝑌 be the r.v. generated by the backward process w.r.t. Ƹ 𝑠, then (Estimator for 𝑊1 distance requires some modification) (𝑠: 密度関数の滑らかさ) [Kazusato Oko, Shunta Akiyama, Taiji Suzuki: Diffusion Models are Minimax Optimal Distribution Estimators. ICML2023, oral] (2% of all submissions) -
データの低次元構造 131 ℝ𝑑 ℝ𝑑′ Theorem (Estimation error by W1-distance) 任意の𝛿
> 0に対し,DNNによる経験誤差最小化元 Ƹ 𝑠 (若干の修正有) は以下の予測誤差を達成する: 内在的次元𝒅′にし か依存しない! これもまたミニマックス最適 (up to 𝛿) [Niles-Weed & Berthet (2022)]. MNIST: 784 dim/ 13.4 intrinsic-dim [Facco et al. 2017] ・多様体への拡張 [Azangulov, Deligiannidis, Rousseau: arXiv:2409.18804] -
離散拡散モデルの理論 132 例: 0,1 𝐷上の分布を学習 • 順過程: • 逆過程: ただし,
仮定 𝑢𝑗: 𝑄の𝑗番目の固有ベクトル (固有ベクトルによる分解) 直接多項分布として推定 (𝑠 > 1) (真の分布) 次元の呪い を受ける 深層スコアマッチングによる学習(拡散モデル) 次元の呪いを解消 (重要な基底に絞って学習: 特徴学習) [Wakasugi, Suzuki: State Size Independent Statistical Error Bound for Discrete Diffusion Models. 2025] -
まとめ 133 事前学習 事後学習 テスト時 推論 事前学習データの質向上 Data augmentation アラインメント
教師有りファインチューニン グ Preference optimization RLHF, RLAIF Monte-Carlo Search In-context learning (Few- shot prompting) Chain-of-thought 汎化 = 情報圧縮 = 特徴学習 (⇒ 生存確率増大) 中間層における表現の獲得 • 次元の呪いの解消 • 思考連鎖による学習効率の向上 • テスト時推論 (文脈内学習) の効率向上 特徴学習の最適化理論 • 計算量・サンプル複雑度ともに恩恵あり Attentionによる動的なトークン選択 • 無限次元入力でも学習可能 (SSMでも一部代替可能) • 論理タスクの適切な特徴表現,その獲得のインセンティブ ➢ Next token predictionだけで情報の圧縮は十分か? -
Direct Preference Optimization • DPO: fine-tuning method for generative models
such as LLMs. 135 Fine-tuning data: • For each prompt 𝑐 ∼ 𝑝(𝑐), generate 𝑦1 , 𝑦2 ∼ 𝑝SFT 𝑦 𝑐 (independently). • Get preference 𝑦𝑤 ≻ 𝑦𝑙 between 𝑦1 , 𝑦2. (human feedback) 1. 2. (Bradley-Terry model) (computation of normalization constant is not required) [Rafailov et al. 2024] -
拡散モデルのファインチューニング 136 「An astronaut riding a horse in a photorealistic
style」 「Teddy bears shopping for groceries in the style of ukiyo-e」 SORA (OpenAI, 2024) Diffusion model DALL·E: [Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever: Zero-Shot Text-to-Image Generation. ICML2021.] DALL·E2:[Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, Mark Chen: Hierarchical Text- Conditional Image Generation with CLIP Latents. arXiv:2204.06125] Optimizing distribution [Rafailov et al. 2024] • Post training: e.g., Preference optimization • Bayesian inference • Reinforcement learning E.g.: DPO, Bayes filtering 𝜇ref: Pretrained diffusion model -
We obtain the density ratio between 𝜇ref and Ƹ 𝜇.
Algorithm 137 We can generate data from 𝜇ref. Dual averaging method: • 𝜇ref: Pretrained diffusion model where For 𝑘 = 1, … , 𝑁 − 1: Phase 1: Optimization. where ← Approximated by neural networks in our implementation -
Convergence analysis 138 where For 𝑘 = 1, … ,
𝑁 − 1: ← Approximated by neural networks in our implementation The algorithm converges in 𝑂(1/𝑁) for both convex and non-convex objective. [Convex] • ℒ 𝜇(𝑁) − ℒ 𝜇opt = 𝑂(1/𝑁) [Non-convex] • min 𝑘∈[𝑁] Var 𝜇(𝑘) 𝛿ℒ 𝜇(𝑘) 𝛿𝜇 = 𝑂(1/𝑁) Theorem (convergence of DA) We only have a density ratio. We don’t have the density. How to generate sample from the target distribution? ⇒ Doob’s h-transform. -
Doob h-transform 139 Phase 2: Sampling from ෝ 𝝁 ∝
𝐞𝐱𝐩(−ෝ 𝒈) 𝝁𝐫𝐞𝐟. Doob ℎ-Transform (Doob, 1957; Rogers & Williams, 2000) Correction term Reverse process for reference model Reference model (Gaussian distribution) Corrected process Optimal model Correction reference (𝜇ref) Corrected process Ƹ 𝜇 𝑌0 𝑌ത 𝑇 𝜇ref Reverse process of 𝜇ref: See also Vargas, Grathwohl, Doucet (2023) & Heng, De Bortoli, Doucet (2024), Uehara et al. (2024) for more details. -
Minimum energy characterization141 Reference model where ത 𝑌ത 𝑇 ,
෨ 𝑌ത 𝑇 ∼ ො 𝜇, 𝑌0 ∼ 𝑝init. with 𝑌0 ∼ 𝑃0 ref. Theorem (Vargas, Grathwohl and Doucet, 2023; adapted) The doob h-transform can be characterized by “minimum energy” ∫ 𝒗𝒕 𝟐𝒅𝒕 solution. (+ KL between initial distributions). For any process ෨ 𝑌𝑡 with ෨ 𝑌ത 𝑇 ∼ ො 𝜇, it holds that (Both process converge to the target distribution) Path measure of 𝑌𝑡 𝑡=0 ത 𝑇 ∵ the construction of ℎ-transform -
Numerical comparison 142 (Existing work) Comparison with Diffusion-DPO • Diffusion-DPO
optimize only an upper-bound. -
Numerical experiment 143 The model is trained to generate a
couple of images facing a same direction. The loss during DA for tilt correction. “Objective”: DPO objective. The targetpoint was [2.5, 0]. “Regularized Objective”: “Objective” + βDKL(q∥pref ), β = 0.01 Right: Tilt-corrected Head CT image generation. Trained model -
直感的説明 145 非滑らか 滑らか • ガウスカーネルを用いた関数近似 (カーネル法・非適応的) 同じ幅の基底の線形結合 → 効率悪い
• NNによる関数近似 場所によって解像度(幅)の違う基底を生成可能 → 効率良い • 深層学習は場所によって 解像度を変える適応力が ある. →学習効率が良い • 浅い学習は様々な関数を 表現できる基底をあらか じめ十分用意して“待ち構 えている”必要がある. →学習効率が悪い -
仮定 𝑓∘ ∈ 𝐵𝑝,𝑞 𝑠 ( 0,1 𝑑): 真が“Besov空間”に入っている. 「浅い」学習との比較
146 ≫ (𝑛: sample size,𝑝: uniformity of smoothness,𝑠: smoothness) カーネルリッジ回帰等: 線形推定量 (非適応的手法) 深層学習 ミニマックス最適性の意味で理論上これ以上改善 できない精度を達成できている. 平均二乗誤差 E መ 𝑓 − 𝑓∗ 2 がサンプルサイズが増えるにつれ減少するレート [Suzuki, ICLR2019] 一様な解像度 適応的解像度 最適ではない 最適 • Wavelet shrinkageより弱い条件 • 基底を用意せず最適化するだけでOK 推定誤差 (平均二乗誤差) : -
[参考] 線形推定量 147 例 • Kernel ridge estimator • Sieve
estimator • Nadaraya-Watson estimator • k-NN estimator 線形推定量: 観測値𝑌 = 𝑦𝑖 𝑖=1 𝑛 に対して線形な推定量. 線形 Kernel ridge regression: “浅い” 学習法 正則化付き最小二乗推定量 (特徴マップ) 固定 学習可能 固定 グラム行列 (カーネル関数) (see also [Imaizumi&Fukumizu, 2019]) -
Training without teacher forcing 149 𝑒1 𝑥1 Position encoding 𝑒2
𝑥2 𝑒𝑑 𝑥𝑑 ⋯ 𝑒𝑑+1 ො 𝑥𝑑+1 𝑒𝑑+𝑘−1 ො 𝑥𝑑+𝑘−1 ⋯ 𝑦 𝑒𝑚 ො 𝑥𝑚 ⋯ Prediction • 𝑥 = (𝑥1 , … , 𝑥𝑑 ) ∼ Unif( −1,1 𝑑) • 𝑦 = 𝑥𝑖1 𝑥𝑖2 … 𝑥𝑖𝑘 = ς𝑗∈𝑝 𝑥𝑗 𝑥1 = 𝑥2 = (where we set ො 𝑥𝑗 = 𝑥𝑗 for 1 ≤ 𝑗 ≤ 𝑑) Without teacher forcing, the model needs to generate CoT chains end- to-end during training, causing error accumulation and complicating dynamics. -
Training strategy 150 • GD with error correction: where 𝑊0
= 𝑂. r[⋅]: rounds to the nearest-integer. To alleviate error accumulation, we apply an “error correction” strategy. • Attention mask: To impose stronger autoregressivity, we make each token only depend on previous levels. W/t mask With attention mask • Stage-wise unlocking of training chains (curriculum learning): ➢ If (internal) tokens on some level are uninformative (≃ 1𝑛; not changed from initialization), zero out its output since all subsequent reasoning will be wrong. ➢ This induces curriculum learning: each 2-parity level is ‘unlocked’ sequentially. -
Learnability without teacher forcing 151 With the training data size
𝑛 = Ω 𝑑2+𝜖 , and the learning rate 𝜂 = Θ(𝑑2+ Τ 𝜖 16), training without teacher forcing achieves a loss ො 𝑦test − 𝑦test ∞ ≤ exp −Ω 𝑑𝜖/16 , after 𝑡 = log2 (𝑘) steps GD updates. Theorem (Sample/computational complexity without teacher forcing) w/o CoT with CoT required data-size Ω(𝑑𝑘−1) O(𝑑2+𝜖) Data-size comparison (the gradient can be replaced by an approximate oracle ෩ ∇𝐿 s.t. ෩ ∇𝐿 − ∇𝐿 ≤ 𝑂(𝑑−2−𝜖/8)) • We still have generalization guarantee for training w/o teacher forcing. -
153 Pre-training Post-training Test time inference •Enhancing quality of data
for pre-training •Data augmentation •Alignment •SFT (supervised fine- tuning) •Preference optimization •RLHF, RLAIF •Monte-Carlo Search •In-context learning (Few-shot prompting) •Chain-of-thought • [Chain-of-thought] Kim&Suzuki: Transformers provably solve parity efficiently with chain of thought. ICLR2025. -
• Reasoning capabilities drastically improve by allocating more compute during
inference time, e.g. running search against a verifier or trained reward model (Jaech et al., 2024; Kimi et al., 2025; Snell et al., 2024; Wu et al., 2024; Guo et al., 2025) • The search trace can be utilized to refine the pretrained model or distill its reasoning patterns into more efficient models (Zhang et al., 2024; Busbridge et al., 2025) 154 How can the benefits of test time scaling methods be rigorously understood? -
Idea: Model long CoT generation as a Markov chain over
abstract reasoning states • Distinguish between easy/trivial reasoning steps (e.g., rearranging terms in an equation) and hard/crucial reasoning steps (e.g., applying an abstract theorem) Task: Find a path from 𝑋in (problem statement) to 𝑋out (conclusion or end-state, e.g. QED) 155 -
156 Assumption of transition probability • Within cluster: Θ(1/𝑀) •
Between clusters: O 𝜖 ≪ O 1/𝑀 1/𝑀 𝜖 Theorem • The transition probability can be properly estimated by GD. • The average hitting time from 𝑋in to 𝑋out is # of clusters = 𝐾 # of nodes in each cluster = 𝑀 -
157 Assumption of transition probability • Within cluster: Θ(1/𝑀) •
Between clusters: O 𝜖 ≪ O 1/𝑀 1/𝑀 𝜖 Theorem • The transition probability can be properly estimated by GD. • The average hitting time from 𝑋in to 𝑋out is # of clusters = 𝐾 # of nodes in each cluster = 𝑀 -
159 Theorem • The tree search algorithm can identify the
cluster structure and the sparse edge with high-probability. • By running PPO-Clip with 𝑐clip = 𝜖max /𝜖, the base model transition probability 𝑝𝜖 can be updated to 𝑝𝜖max with high probability when 𝜖max = ෩ O(1/𝑀). • PPO-Clip [Schulman et al., 2017]: Maximize the PPO loss 𝐿PPO with a reward መ 𝐴 such that መ 𝐴 𝑋0 , 𝑋1 = 1 if 𝑋0 , 𝑋1 ∈ 𝐸𝑠 (sparse edge) and 0 otherwise: Before PPO After PPO Time complexity of search: -
Extension to logical reasoning task • Given (𝑋in , 𝑋out
), the goal is to output both ➢A valid path 𝑋0:𝑇 from 𝑋in to 𝑋out , and ➢its logical value 𝑟 𝑋0:𝑇 = 𝛼 𝑋0 ⋅ 𝛼 𝑋0 , 𝑋1 ⋅ ⋯ ⋅ 𝛼 𝑋𝑇−1 , 𝑋𝑇 where 𝛼 𝑋𝑡 , 𝑋𝑡+1 = 1 if 𝑋𝑡 , 𝑋𝑡+1 ∉ 𝐸𝑠 . 160 Theorem (hardness without global information) Let 𝑓𝜃 nbd 𝑀𝑝 𝑋in , 𝑋out , 𝐴 be any parametric model with polynomially bounded gradients that can freely search a local neighborhood of the generated CoT. Then any algorithm 𝜃(𝐴) that makes at most poly (𝐾) queries to the 𝐾−𝜔(1)- approximate oracle of gradient satisfies w.p. 1 − 𝐾−𝜔(1) for 𝑀 sufficiently large. • Logical reasoning task without global information requires exponential computation w.r.t. 𝐾. (proof: derive an SQ-lower bound for this class) • The global information obtained at the pretraining phase can mitigate this to polynomial time.