PHOTON を読んで、手元でミニチュア再現する — 475 倍をスループット×メモリに分解し、チャンク階層を自分で測る

富士通が発表した新アーキテクチャ PHOTON¹は、「GPU リソース当たり最大 475 倍のスループット」を謳います。この数字を鵜呑みにするのも、全否定するのも簡単です。今回はそのどちらもせず、論文を読み、コアとなるアーキテクチャをミニチュアで自作し、自分の手元で測ります。

トークン化から始まり、attention、KV キャッシュ、サンプリング、ゼロからの学習と積み上げてきたこの連載も、今回で最終回です。

この連載の現在地 — 積み上げた語彙で見出しを分解する

• ここまで分かっていること — トークン化と自己回帰（Part 1）、self-attention と O(T²)（Part 2）、KV キャッシュと memory-bound な decode（Part 3）、self-consistency による test-time compute の経済（Part 4）、ゼロからの学習コスト（Part 5）を、すべて自分の手元で測ってきました
• 今回新しく説明できること — PHOTON のアーキテクチャ（bottom-up encoder + top-down decoder のチャンク階層）を読み解き、「475 倍」をスループット×メモリに分解したうえで、そのミニチュア版を自作して Part 5 のモデルと同条件で比較します
• 今回はやらないこと — 論文の忠実な再現（134.2B トークンの学習・DGX H200・公式コード）、RecGen という追加の decode 最適化（「10³倍」という別の数字の内訳）、量子化・蒸留・分散学習・MoE や状態空間モデルとの比較
• PHOTON のどの主張につながるか — すべてです。最終回として、ここまでの伏線を「なぜ階層チャンクが O(T²) と KV traffic の両方を減らせるのか」という1つの説明に収束させます

PHOTON が解こうとした問題

Part 3 で確かめたように、Transformer の decode は KV キャッシュの読み出し帯域に律速される memory-bound な処理です。context が伸びるほど毎ステップ読み書きする KV キャッシュが線形に膨らみ、decode 速度はそのサイズに縛られていきます。

PHOTON の論文はこの状況を「horizontal token-by-token scanner」と表現しています。どのステップでも、直前までのトークン全履歴を水平に舐め続けるという意味です。論文はここで発想を転換します。生成を token-by-token の水平スキャンとしてではなく、粗い状態から細部へ降りていく vertical, multi-resolution context scanning として設計し直せないか、という問いです。

言語には元々このような階層があります。文字が単語を作り、単語が文を作り、文が文章を作ります。会話の一貫性も、全ての細かいトークンを毎回読み直すのではなく、粗い議論の流れを保持しながら必要な時だけ細部に降りることで保たれています。PHOTON はこの自然言語の階層性を、推論アーキテクチャの側に持ち込みます。

アーキテクチャの仕組み — bottom-up encoder と top-down decoder

PHOTON は2つの部品からなります。

Hierarchical Encoder（bottom-up）は、token 列を粗い潜在表現の列に圧縮します。まず Context Chunker が複数トークンをまとめて1つのチャンク表現にします（論文は concat して線形射影するというシンプルな方法を採用しています）。続く Context Encoder がそのチャンク表現の列を自己回帰 Transformer で文脈化します。これは Part 1〜5 で見てきた自己回帰 Transformer そのものですが、粒度が「トークン」ではなく「チャンク」になっている点が違います。

Hierarchical Decoder（top-down）は、粗い潜在表現から token レベルの表現を再構成します。Context Converter が1つのチャンク潜在を複数本の conditioning vector に展開し、Context Decoder がそれを窓として、チャンク内のトークンを局所的な causal attention でデコードします。この局所デコーダの attention span はチャンクの窓幅に固定され、系列全体の長さ T には依存しません。これが Part 2 で見た O(T²) の壁を回避する仕掛けです。

並列にできるのは「チャンク内のトークン」ではなく「チャンク間」であることに注意してください。学習時（teacher forcing）は正解トークンが全て既知なので、複数チャンクをバッチにまとめて一度に処理できます。しかし生成時は、チャンク内のトークンは1つ前のトークンに依存する自己回帰列であり、vanilla の token-by-token 生成と同じく1つずつ順に生成するしかありません。論文もこの局所デコーダを「autoregressively in each chunk（各チャンク内で自己回帰的に）」と明記しています¹。PHOTON の並列性は、独立した複数のチャンクや複数の系列をまとめて処理できることにあります。

論文の主結果はこの階層を2段重ねます（チャンク長 4, 4）。今回作るミニチュア実装 photon_mini.py は、この階層を1段に簡略化します。トークン列を8個ずつ束ねて1段だけ粗い潜在を作り、そこから token レベルの表現を再構成する — MEGABYTE（Meta, 2023）²と同型の、global/local という2層構造です。学習目的に含まれる recursive consistency の補助損失も実装していません。これは手抜きではなく、論文の主結果自体が「階層構造そのものの効果を単離する」ためにこの補助損失をオフにした設定を採用しているのと同じ条件です。

「475倍」を分解する

論文の Table 1 は、1.2B パラメータのモデルを decode-heavy（短い入力・長い継続生成）という条件で比較しています。数字をそのまま並べます。

TPM（throughput per memory）は論文が定義する合成指標で、TPM = スループット ÷ メモリです。表を素直に読むと、スループットは 1.00 から 43.80 へ約44倍、メモリは 0.390 から 0.036 へ約11倍削減、その掛け算として TPM が 2.56 から 1216.67 へ約475倍になっています。論文本文にも、1.2B モデルの decode-heavy 設定で TPM が 475 倍改善すると明記されています¹。

ここで大事な注意が2つあります。

1つ目。475倍は「1.2B モデル・decode-heavy」という特定条件の数字です。モデルサイズや入出力比率が変われば倍率も変わります（600M では約417倍、900M では約1195倍と、論文の同じ表の中でも一定ではありません）。PHOTON が常に475倍速いという読み方はできません。

2つ目。論文の Abstract と結論は「最大10³倍」というさらに大きな数字も掲げています。これは475倍とは別の条件から来ています。RecGen という追加の decode スケジュール（bottom-up の再エンコードを省略し、最上位レベルの状態だけを更新し続ける手法）を組み合わせた場合の最大値で、論文の Table 2 から計算すると 600M モデルで最大1,856倍に達します。論文自身も RecGen の品質面は teacher-forcing による厳密な尤度評価が難しく将来課題だと正直に留保しています¹。この記事では、アーキテクチャ本体の効果である475倍（HierGen）に焦点を当て、RecGen は扱いません。

品質面も直視します。同じ Table 1 から WikiText perplexity は 600M で約34%、1.2B で約21%悪化しています。zero-shot 精度（HellaSwag・SciQ・ARC-Easy の平均）も、600M で3.8ポイント、1.2B で6.8ポイント低下しています。論文自身は moderate degradation と表現していますが、20〜34%という数字はわずかと切り捨てられる規模ではありません。

忠実な再現はできない

PHOTON は The Pile（134,217,728,000 トークン）を NVIDIA DGX H200 で1エポック学習したモデルです¹。個人の Mac 1台では、この学習量・この計算資源を再現することはできません。公式コードも公開されていません。

だから、この記事が再現するのは論文の倍率ではなく、アーキテクチャの方向性と機構です。チャンク階層がメモリを削減する仕組み、スループットを改善する仕組み、その代償として品質がどう動くか。この3点を、TinyShakespeare という小さなデータセットの上で、Part 5 と全く同じ条件で確かめます。

photon_mini.py を作る

photon_mini.py は Part 5 の tiny_gpt.py を companion リポジトリ内で直接 import します。その decoder ブロック（causal self-attention + MLP）を、bottom-up encoder と top-down decoder の両方の部品として再利用します。新しく実装したのは4つの部品です。チャンク化（ContextChunker）、チャンク列の文脈化（ContextEncoder）、conditioning vector への展開（ContextConverter）、チャンク幅に窓を固定した局所デコード（ContextDecoder）、そしてそれらを束ねる PhotonMini クラスです。

パラメータ数は vanilla（10,795,776）にほぼ一致するよう埋め込み次元を較正しました（10,770,700、差0.23%）。開発は Part 1〜5 と同じ規律で進めました。まず selftest() を実機なしの CPU 上で通します。検証したのは、forward の出力形状、過学習で loss が実際に下がること、そして階層的な causality です。最終チャンクの内容を変えても先頭チャンクの出力 logits は変わらない一方、先頭チャンクを変えると最終チャンクの出力は変わる — これは情報が bottom-up に確かに伝搬しているかを確認する、この実装ならではのチェックです。全項目パスを確認してから、初めて実機実走に進みました。

MEGABYTE との違いにも触れておきます。MEGABYTE はトークナイザを使わず、生のバイト列を patch 単位でモデリングすることが目的のアーキテクチャです²。一方 PHOTON は、バイトから単位を学習するのではなく、すでにサブワードトークン化されたモデルに、さらに上位の階層的文脈を追加して推論を高速化することが目的です。階層で長距離依存を圧縮するという設計思想は共有していますが、対象とする粒度と目的が異なります。

実測 — PPL・スループット・メモリを同条件比較する

photon_mini.py を Part 5 と全く同じ設定（TinyShakespeare、5,000 iteration、batch size 64、fp32、seed 1337）で MPS 実走しました。さらに3 seed を確保するため、seed 1338・1339 でも追加実走しています。以下、学習(PPL)と生成(スループット・メモリ)の2つの実験に分けて見ていきます。

学習では PPL はほぼ互角、しかし過学習の挙動が違う

best PPL（各 run が到達した最良値）で比べると、PhotonMini は 4.415、vanilla は 4.389 とほぼ同水準です。差は0.6%で、PhotonMini 自身の seed 間ばらつき 0.032 の範囲に収まります。パラメータ数を揃えたミニチュア実験では、チャンク階層は品質を大きく犠牲にしていません。

一方で、最終 iteration（4,999）まで学習を続けた場合の挙動が3 seed 全てで違いました。Part 5 の vanilla は val loss が iteration 2,250（PPL 4.389）で底を打ったあと反転悪化し、iteration 4,999 では PPL 4.936 まで悪化しています。過学習です。PhotonMini は3 seed とも、この反転がほとんど見られませんでした。final PPL は 4.433 で、best との差はわずか0.4%です。

これは事前に想定していなかった一次観察です。理由として考えられるのは、チャンク化による情報のボトルネック（8トークンを1つのベクトルに圧縮してから書き戻す）が暗黙の正則化として働き、小さいデータセットへの過適合を遅らせている可能性です。ただし3 seed・1つのアーキテクチャ変種・1つの小さな char-level データセットでの観察に過ぎず、階層構造は過学習しにくいという一般的な主張はできません。

学習時間とメモリも、チャンク階層側が優位でした。学習は約2倍速く終わり（20.9分 対 40.8分）、ピークメモリは半分以下（0.517GiB 対 1.267GiB）です。ただしこれは PHOTON 論文の主張である decode 時のメモリ削減とは別の効果で、学習時の活性化メモリ・計算量の違いから来ています。論文が測っているのは推論時の KV キャッシュであり、学習時のメモリではありません。この違いを混同しないよう、次の節で切り分けます。

生成: スループットは控えめな改善にとどまった

--measure-efficiency で、context 長 16〜192 の5点について、未学習のモデル（アーキテクチャの計算・メモリ特性そのものを測るため、学習済み重みは不要です）で HierGen（PhotonMini）と vanilla の decode を比較しました。各点は5回反復した中央値です。

Phase 0 の段階で立てた事前予想 — 「スループットは MPS や Python のオーバーヘッドに埋もれやすい」 — は、実測でそのまま裏付けられました。スループットは1.39〜1.62倍の改善にとどまり、論文の44倍とは桁が違います。

チャンク内のトークンが逐次生成であること自体は vanilla と同じ制約であり、ここでの遅さの原因ではありません（前節で見た通り、チャンク内は自己回帰的で並列化できません）。差が縮む主な理由は、このミニチュアが数千万トークン規模ではなく数百トークン規模の nano 実験であることです。論文の44倍は長い context・大きな KV キャッシュが memory bandwidth を圧迫する状況（Part 3 で見た memory-bound な領域）で効いてきますが、本実験の context 長（16〜192）はその領域にまだ届いていません。加えて、この実装は Python のインタプリタ・オーバーヘッドや、独立したチャンク・系列をまとめて処理するバッチ化を持たない教育用のミニチュアであり、本番のカーネル実装であればさらに縮められる余地があります。

生成: KV bytes は一桁後半〜十倍規模の削減が出る

同じ計測から、PhotonMini（HierGen）が理論上必要とする KV キャッシュサイズと、vanilla の KV キャッシュ理論値を比較します。どちらも実際にキャッシュを保持する実装はしていません。本文にも触れた通り、Part 5 の学習ループと同じく素朴に forward するだけです。そのため両者とも Part 3 と同じ理論式（2 × n_layer × n_head × head_dim × T × dtype_bytes）で計算しています。PhotonMini 側の内訳は2つです。context encoder が保持し続ける global cache（チャンク数 M に応じて増える）と、decode 中の1チャンクぶんだけ保持される local decoder cache（チャンク境界ごとに破棄される定数サイズ）の合計です。

KV bytes（メモリ）は、context 長が伸びるほど削減率が上がる傾向で、8.1〜12.8倍の削減が出ています。local decoder cache が定数サイズのため、context が短いほど PhotonMini 側の総メモリに占める割合が相対的に大きくなり、削減率が下がります。context が伸びるほど global cache（context 長に応じて増える部分）が支配的になり、削減率は論文と同じ「context が伸びるほど有利」という方向に近づいていきます。

スループットとメモリを掛け合わせないことについて: 論文の TPM はスループットとメモリを同一の実装・同一の実測系で一貫して測った値です。この記事のスループット実測は、vanilla・PhotonMini とも実際には KV キャッシュを持たない素朴な forward の繰り返しであるのに対し、メモリの数字は「キャッシュを実装したら必要になる理論値」です。測定条件が異なる2つの数字を掛け合わせても、論文の TPM に相当する一貫した量にはなりません。そのため、この記事ではスループット改善率とメモリ削減率をそれぞれ独立した発見として報告し、両者を掛け合わせた「TPM」の値は提示しません。これも「475倍」を分解してきたこの記事らしい結論です。掛け算で出す数字は、掛け合わせる2つの数字が同じ土俵で測られていて初めて意味を持ちます。

擬似再現A×Bをつなぐ

Part 4 の self-consistency（擬似再現A）は、同じクエリを N 回サンプリングして多数決を取ることで、test-time compute を品質に変換する仕組みでした。今回の擬似再現B は、アーキテクチャの階層構造でスループット×メモリを変換する仕組みです。両者は別の機構であり、数字を掛け合わせて PHOTON を再現したという主張はしません。

PHOTON の実運用は、RecGen とマルチクエリ集約を組み合わせて、階層構造による品質劣化を補償する設計思想を持っています。self-consistency と PHOTON のマルチクエリ集約は仕組みが異なりますが、test-time compute で品質を買い戻すという同じ経済を、別の機構で体感するという意味では地続きです。この記事では、この2つを物語としてのみ接続します。

連載を振り返る — トークンからPHOTONまでの6歩

この連載の到達点は、PHOTON を再現できたことではありません。トークン・attention・KV キャッシュ・サンプリング・学習という語彙を積み上げたことで、GPU あたり最大475倍という見出しを、スループットとメモリという別々の効果に自分で分解できるようになりました。その一部を自分の手元の小さな実験で確かめられる、ここまで来ました。次に「N倍高速化」という見出しに出会ったとき、どこから疑うべきかが、もう分かるはずです。

手元に残るもの

• 動くコード — photon_mini.py（bottom-up encoder + top-down decoder のチャンク階層、Part 5 の tiny_gpt.py を継承）
• 測った数字 — best PPL 4.415 ± 0.032（vanilla の4.389とほぼ互角）、スループット改善1.4〜1.6倍、KV bytes 削減8.1〜12.8倍。両者は測定条件が異なるため掛け合わせた「TPM」は提示していません。いずれも3 seed・1つの小さな char-level データセット・1つの簡略化アーキテクチャ変種での観察であり、一般化はしていません
• メンタルモデル — 「N倍高速化」という見出しは、スループットとメモリという別々の効果の積であることが多いです。分解すれば、どちらが本物でどちらが条件依存かが見えてきます

トークンから PHOTON まで、6つの実装と6つの実測を積み上げました。ここから先は、あなたが次に読む「N倍」のニュースを、この語彙で分解する番です。

コードはサンプルコード（companion repo）の part-06 タグで再現できます。clone と各 Part のタグは、companion repo かこの記事末尾のシリーズナビからたどれます。

参考文献