PHOTON という階層型言語モデルを、公開論文を手掛かりに実装してきました。公式実装の移植ではなく、論文の数式と付録表から構造を組み直し、vanilla Transformer や Block Transformer と同じパイプラインで比較するための検証実装です。
本記事は、生成AIを利用しながら作成しています。リポジトリと整合が取れるように記載しているつもりですが、誤っている部分があれば適宜修正します。
コードと計測結果は verify-photon で公開しています。この記事で扱うのは、ローカルの main ブランチが fdcd225 に到達した時点の内容です。
main まで進めて、結果は単純な成功談にはなりませんでした。推論時に持ち続ける状態は確かに小さくなり、RTX 4090 上の同時リクエストでも明確な差が出ました。一方、tiny モデルの品質は vanilla に届きませんでした。RecGen の整合性をほぼ完全に合わせても、生成が良くなるとは限りませんでした。
この記事では、良かった数字だけでなく、途中で読み直した仕様、崩壊した学習結果、あえて空欄にしたベンチマークも含めて、どこまで確認できたかを整理します。
このリポジトリは PHOTON 論文の非公式・独立実装であり、Fujitsu および論文著者とは関係ありません。公式・第三者実装のソースコードは参照していません。
最初に、検証できた範囲を区切る
「実装が動いたこと」と「論文の結果を再現したこと」は分けて考えました。
| 今回確認したこと | 今回確認していないこと |
|---|---|
| 因果性を保つ階層構造と学習経路 | 論文と同じ 134B token の学習 |
| 600M / 900M / 1.2B の付録表に載る各モジュールの parameter count | 論文の headline PPL・zero-shot 結果 |
| HierGen / RecGen の生成手順と cache の成長量 | 著者側の学習データ、tokenizer との byte identity |
| RTX 4090 上の throughput と persistent state | 論文の絶対 TPM 値や 475 倍の multi-query 結果 |
| continuous batching での PHOTON / Block / vanilla 比較 | 9-query aggregation や decoder-only の weight residency 削減 |
今回の到達点は、PHOTON の推論効率を生む仕組みは確認できたものの、論文の品質結果まで再現したわけではない、というところにあります。以降の数字も、この境界内で読む必要があります。
PHOTON は何を変えるのか
通常の Transformer は、生成した token ごとに token-level の KV cache を積み上げていきます。context が長くなるほど、各 step で読む状態も増えます。
PHOTON は token 列を複数の時間解像度に圧縮し、上位の低頻度な状態から下位の表現を局所的に復元します。
token stream x0..x3 | x4..x7 | x8..x11 | x12..x15 | ... 長さ T
│ │ │ │ C1 = 4
level 1 stream z0 z1 z2 z3 ... 長さ T / 4
└────────┴─────────┴──────────┘ C2 = 4
level 2 stream h0 ... 長さ T / 16
bottom-up: chunking + global causal encoder
top-down: parent latent -> converter -> bounded local decoder
今回の中心設定は C=[4, 4] なので、最上位 stream は 16 token に 1 回だけ更新されます。RecGen では、次の最上位入力を生成 token の再 encode ではなく decoder 側の reconstruction から作ります。そのため、成長し続ける状態を最上位 KV に絞れます。
ここで「constant memory」と言い切るのは正確ではありません。下位には bounded buffer があり、最上位 KV は O(T/16) で成長します。通常の flat KV が O(T) で伸びるのに対し、成長率を圧縮している、と捉えるのが近いです。
論文を実装に落とす
最初に固定したのは、モデルの大きさではなく計算契約
階層化すると、shape が合って loss が下がるだけでは安心できません。chunk 内の未来 token が上位表現を経由して漏れれば、見かけ上は学習できてしまいます。
そこで、学習を回す前にまず次のコア契約をテストで固定しました。
- token
jを変更しても、strictly-past な reconstruction state は位置j以前で変わりません（公開 logits では位置jより前が変わりません） - converter は「直前の parent latent」だけを見ます
- local decoder は、予測対象自身の teacher state を見ません
- recursive loss の重み
α=0でも、token CE の勾配が最上位 encoder まで届きます - 公開 logits は Hugging Face の CausalLM 契約、つまり
logits[:, t]が tokent+1を予測する形にします
学習 path を追加した Phase 2 では、chunk 境界に揃わない長さと right padding の境界 test も加えました。最初からすべてを見通せたわけではなく、実装範囲が広がるたびに契約を増やしています。
特に logits の契約は地味ですが重要でした。内部の reconstruction は「位置 t を、それより前の情報から予測する」という並びになります。一方、Hugging Face の shifted CE や評価器は 1 token ずれた公開契約を期待します。最終位置だけは「次の位置を 1 step 生成する計算」で補い、通常の CausalLM と同じインターフェースに揃えました。
現在の main 相当コードでは、これらを含む test suite が通っています。
414 passed
実行環境によっては PyTorch の NVML 初期化 warning が出ることもありますが、test の成否とは別です。
付録表を読んで、最初の設計を見直す
初期実装では、token 側と上位側の次元差を projection でつなぐ想定でした。しかし、論文の付録表をモジュール単位で照合すると、その構成では parameter count が合いませんでした。
表から読み取れたのは次の 3 点です。
- level 1 の chunker に parameter row がないため、ここは parameter-free の pure concatenation と考えるのが自然です
- 次の chunker は、連結後の RMSNorm と bias 付き Linear を採用すると表の値に一致します
- token 側の
416 -> 1664の次元差は projection ではなく、付録の独立したEmbedderで埋めます
この読み直しを paper_600m()、paper_900m()、paper_1200m() に反映し、付録表に記載された全モジュールの parameter count を row ごとにテストしました。converter の prefix 長 R=2 も、本文に値がないため付録の parameter count から逆算しています。ただし、parameter count だけでは normalization の粒度や wiring の全詳細までは一意に決まりません。ここでは全 table row と一致する最も単純な構成を採用しました。
なお、モデル総数を付録の合計と完全一致させたわけではありません。表にはありませんが、計算上必要な learnable start latent があるため、実装の総数はその分だけ増えます。ここを削って「完全一致」とするより、差の由来を残す方を選びました。
“HierGen” は一つではない
実装途中で最も大きかった読み直しは、論文中の “HierGen” が二つの異なる手順を指していたことです。
-
hiergen_teacher_forced: 学習時の conditional と整合する reference です。下位 chunk 境界ごとに再 encode します -
hiergen_paper: ACL Appendix B.2, Definition B.2 の手順です。一つの top-level state から meta-context 全体を生成し、その token を再 encode して次の top-level 入力を作ります -
recgen: Definition B.3 の手順です。meta-context 内はhiergen_paperと同じですが、次の top-level 入力に decoder 側の reconstruction を使います
meta-context 境界に揃った prompt から始めると、hiergen_paper と recgen は最初の meta-context では同じ token を生成します。差が生じ得るのは、最上位を更新し、strict right shift の遅延を通過した後です。この性質も test にしました。
過去の JSON では hiergen が teacher-forced 版を指していました。後から数字の意味を書き換えず、hiergen は legacy alias として残し、定義を明記した field を追加しました。ベンチマークを読む際は、bare な “HierGen” ではなく、どちらの schedule かを見る必要があります。
生成 API も明示的に分けています。以下は tiny を学習して checkpoint を作った後の概念例で、input_ids は tokenized prompt とします。
from photon import PhotonForCausalLM
model = PhotonForCausalLM.from_pretrained("runs/tiny-photon/final").eval()
output = model.generate(
input_ids=input_ids,
generation_mode="hiergen_paper", # or "hiergen_teacher_forced", "recgen"
max_new_tokens=128,
parallel_chunks=True,
)
品質側で起きたこと
tiny の品質評価で分かったこと
main には、既存の 4 checkpoint に対する FP32 の local quality suite も載せました。standalone PPL evaluator と、lm-eval==0.4.12 による HellaSwag / SciQ / ARC-Easy の zero-shot evaluator は、共通の checkpoint / tokenizer 読み込み契約から実行しています。
4 checkpoint の記録上の学習条件は WikiText-2 raw、GPT-2 tokenizer、sequence length 512、batch 16、1,500 steps で、Trainer の記録では約 5.10 epoch でした。ただし、学習時の dataset / tokenizer revision までは固定されていません。下表の PPL は別 protocol で、固定 revision の test split を特殊 token や文書 separator なしで一つの stream に連結し、context 2048・stride 1024 で評価しています。
| checkpoint | parameters | WikiText PPL | 3-task mean |
|---|---|---|---|
| tiny PHOTON | 31.2M | 478.2283 | 0.24754 |
| tiny vanilla | 37.0M | 299.3707 | 0.24349 |
| tiny_paper PHOTON | 62.3M | 434.7606 | 0.25752 |
| tiny_paper vanilla | 56.3M | 309.5485 | 0.25137 |
PHOTON は二つの組み合わせで PPL が悪く、3-task mean は vanilla より 0.004、0.006 だけ高い結果でした。ただし各 task の値は、HellaSwag が 0.252–0.259、SciQ が 0.207–0.229、ARC-Easy が 0.269–0.285 で、いずれも均等ランダム選択の基準（ほぼ 0.25）に近い値です。しかも単一 seed で、bootstrap による誤差も算出していません。この差を PHOTON の品質優位の根拠にはできません。
表の
WikiText PPLは WikiText-103 raw と WikiText-2 raw の両方で同じ値でした。調べると、固定した dataset revision では config の fingerprint は異なるのに、test split は同じ 283,287-token stream と同じ SHA-256 になっていました。そこで片方を消さず、両方の record と stream hash を残しました。
また、学習中の共通評価では PHOTON / vanilla の PPL 比が tiny の 2.38 倍から tiny_paper の 2.09 倍へ縮みました。リポジトリ独自の分析規約 photon-analytical-matmul-v1 で sequence length 512 の forward FLOPs を数えると、PHOTON / vanilla は tiny で 0.499、tiny_paper で 0.784 でした。つまり、tiny PHOTON は vanilla の約半分の forward FLOPs で PPL が 2.38 倍、tiny_paper は約 78% の forward FLOPs で PPL が 2.09 倍でした。一方、tiny_paper PHOTON 自体は tiny の 1.571 倍の forward FLOPs と、ほぼ 2 倍の parameter を使っています。これは equal-step・single seed の観測であり、FLOPs matched でも parameter matched でもありません。付録どおりの配線が差を縮めた、と因果帰属できる比較ではありません。
これらは historical checkpoint の学習 provenance が end-to-end で固定されていないため、
legacy-tiny local protocolと呼んでいます。論文の 600M headline と並べるための数字ではなく、Phase 6 の評価 path が最後まで動くことと、tiny での傾向を確認するための記録です。JSON には weight hash がありますが、checkpoint 本体は Git 管理していないため、clone だけで掲載値を bit-exact に再評価できるわけではありません。
整合性がほぼゼロでも、生成は良くならない
RecGen が hiergen_paper に近い conditional を保つには、decoder 側で再構成した bottleneck と、生成 token を bottom-up に再 encode した bottleneck が近い必要があります。この距離を縮める recursive loss の重みが α です。
tiny モデルで α=0.3 にすると、学習時の top-bottleneck cosine distance は約 0.784 から 0.0003 まで下がりました。WikiText-2 の 32 window（prompt 64 + horizon 192）で測った RecGen と二つの HierGen の stepwise KL もほぼゼロになりました。
ここだけを見れば成功に見えます。しかし同時に、次の結果が出ていました。
- forced-trajectory の teacher-forced NLL:
5.781 -> 6.076 - eval PPL: 約
445 -> 538 - greedy 生成の repeated 4-gram fraction:
0.979 - greedy 生成の distinct-1:
0.0026
ここで PPL は、sequence length 512 の packed data を使った Trainer evaluation の値です。前節の context 2048・full test stream の PPL とは評価 protocol が異なります。
実際の greedy sample は同じ短い句を繰り返していました。サンプリング（temperature 0.8、top-p 0.95）では同じ n-gram collapse は出ていませんが、それだけでこの設定を良い生成モデルとは判断できません。
つまり、reconstruction error が小さいだけでは、RecGen の有用な生成品質を保証できません。今回の記録では、near-zero の mismatch と崩壊した greedy output が同時に観測されました。今後 α を選ぶなら、token quality、free-running の反復、latent diversity も判断条件に含める必要があります。
この比較にはもう一つ注意があります。後段の gap 評価 JSON は clean commit ですが、
α=0.3の学習 JSON 自体はgit_dirty=trueでした。その学習 source tree を clean commit に結び付けた証拠にはできないため、main に残る legacy record を後から再点検した警告材料として扱っています。
推論側で確認できたこと
RTX 4090 での continuous serving
推論効率は、モデルごとに別の都合のよい loop を使うのではなく、同じ request lifecycle で比較しました。各 request は prompt、sampling 設定、RNG、EOS、生成長を個別に持ち、同じ cache signature の request だけを一つの cohort にまとめます。完了した row は compact され、後から到着した request も state shape が合えば join できます。
計測全体では concurrency 1 / 8 / 32 を sweep し、各 cell は warmup 1 回の後、fresh engine で 2 回測りました。prefill / decode 時間と scheduler steps は 2 回の平均、logical state は 2 回のうち大きい peak、peak VRAM は 2 回を通した最大 allocation を採用しています。次の表は、そのうち concurrency 32 の 4 つの cell を抜き出したものです。paper600m 構成、RTX 4090、bf16、base prompt 128 token、base output budget 256 token で、実際の paired request には arrival delay に応じた不均一な budget も含まれます。速度・状態量の検証なので、weight は random です。各 architecture は論文付録にある自身の構成を使っており、parameter 数は PHOTON 646.4M、Block 629.8M、vanilla 610.9M でした。
表の TPM(state) は throughput per memory、つまり aggregate decode throughput を request あたりの persistent state（GiB）で割った値で、単位は K tok/s/GiB です。aggregate throughput を使うため、絶対値は concurrency とともに増えます。
| model / schedule | decode tok/s | state MiB/request | TPM(state) K tok/s/GiB | peak VRAM |
|---|---|---|---|---|
| vanilla KV | 2,824 | 39.003 | 74.2 | 3.03 GiB |
| Block Transformer | 4,885 | 5.100 | 981 | 1.46 GiB |
| PHOTON teacher-forced HierGen | 7,755 | 3.290 | 2,413 | 1.41 GiB |
| PHOTON RecGen | 9,096 | 0.878 | 10,604 | 1.40 GiB |
なお、serving JSON の mode: "hiergen" は、現在の命名では hiergen_teacher_forced に当たる legacy alias であり、論文 Definition B.2 の hiergen_paper を測った行ではありません。
この条件では RecGen は vanilla に対して、decode throughput が 3.22 倍、request あたりの persistent state が 44.4 倍小さくなり、state-based TPM は 143 倍になりました。全モデルの scheduler width は同じで、mean / max は 31.63 / 32 でした。
この表にはいくつか限定があります。
- random weight なので、品質を含む優位性を示すものではありません
- request 到着を cache signature に合わせた best-case の cohort formation です
- TPM の分母は、このリポジトリで定義した persistent generation state です。論文の KV-only 定義と絶対値を直接比較できません
- serving の出力は token-sequential で、後述する offline の chunk folding はまだ統合していません
また、offline と serving では workload 長だけでなく vanilla の state 分母も同一ではないため、二つの節の絶対 TPM は横断比較しません。
それでも、PHOTON の方が小さいモデルだから速い、という結果ではありません。上の 3 モデルでは PHOTON が最も大きいです。それでも state と decode throughput の両方に差が出た点は、階層化の機構を確認する材料になりました。
長い decode と、未計測の一つのセル
offline benchmark では、prompt 128 + 2048 token 生成、batch 32 の decode-bound 条件も測りました。tiny 構成の historical record では、vanilla が 5,767 tok/s, 31.88 MiB/sample、RecGen が 12,475 tok/s, 0.55 MiB/sample でした。状態は約 58 倍小さく、このリポジトリの定義による TPM は約 125 倍でした。
paper600m でも PHOTON と Block は batch 32 まで測れましたが、vanilla の decode-bound / batch 32 は空欄にしています。
理由は、修正前の 18-layer vanilla で 4 回試したところ、WSL2 の GPU paravirtualization 層で residency exhaustion が起きたためです。最後は CUDA illegal memory access と host crash まで進みました。その後、論文表に合わせて vanilla を 16 layers に修正しましたが、「修正版なら落ちないはず」と一般化して再試行するのはやめました。修正版の prefill-bound / batch 32 は peak 16.5 GiB で完走していますが、decode-bound は未計測のままです。
空欄を推定値で埋めなかったことも、今回の検証結果の一部だと考えています。
独立 chunk の一括実行で約 2 倍
hiergen_paper と RecGen は、一つの meta-context 内に互いに独立な下位 chunk を持ちます。greedy 生成に限れば、これらを batch 次元へ fold してまとめて実行できます。
同じ clean record 内で逐次版と比較すると、paper600m の decode-bound / batch 32 は次のようになりました。
-
hiergen_paper:7,655 -> 16,169 tok/s（2.112 倍） - RecGen:
9,175 -> 18,264 tok/s（1.991 倍）
paper600m の各 cell は 1 iteration です。絶対性能の精密な推定ではなく、同じ記録内で execution shape の差を見る paired A/B として扱っています。
persistent state の shape と byte 数は変わりません。浮動小数点演算の順序は変わるため logits の bit identity は要求せず、greedy token の一致と、logits / cache state の数値的近さをテストしています。
partial prefix や短い suffix は、同じ request の中で逐次処理を併用します。sampling、EOS-aware generation、teacher-forced の両表記、L=1 では warning を出して既存の逐次 path に戻します。都合の悪い入力を黙って無視して速い path だけを測らないようにしました。
paper-scale は「実行できる形」まで
論文の main training 条件は、Pile-uncopyrighted 177,008,913 documents、134,217,728,000 tokens、context 2048、global batch 256、256,000 optimizer steps、LLaMA tokenizer、α=0、DGX H200 です。この規模の学習は実行していません。
代わりに、次の tooling までを main に入れました。
- revision-pinned で resume 可能な corpus preparation
- atomic shard、全 shard の SHA-256、tokenizer snapshot と manifest
- 600M / 900M / 1.2B の未実行 training config
- versioned な FLOPs 計算と baseline の FLOPs matching
- data / config / code の semantic identity と、別枠の runtime audit を記録する
run_plan.json - FP32 の token-weighted PPL evaluator
- dataset revision と
lm-eval==0.4.12を固定した zero-shot evaluator - training / PPL / zero-shot の result file を上書きしない create-exclusive publication
ここでも「再現可能」と「著者と同一」は分けました。公開候補の monology/pile-uncopyrighted は、固定 revision の row count が論文記載と一致しません。論文は prepared corpus の hash も、LLaMA tokenizer の正確な artifact も公開していません。
そのため通常の corpus preparation は count mismatch で fail closed にしています。公開候補を approximate source として使う場合は、準備時の PHOTON_ALLOW_APPROXIMATE_SOURCE=1 と、学習時の --allow_approximate_data を別々に指定する必要があります。動かすために暗黙で条件を緩めるのではなく、結果 manifest に差を残すためです。
main の最後に詰めた Hugging Face 互換性
最後に直したのは、新しい architecture の中身ではなく、保存した checkpoint 単体で動くかどうかという点でした。
save_pretrained() した checkpoint には custom code と auto_map を含め、元の checkout を import path から外した別プロセスでも読み込めるようにしました。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
checkpoint_or_hub_id,
trust_remote_code=True,
)
AutoConfig、AutoModel、AutoModelForCausalLM のすべてを同じ条件の別プロセスで確認しています。return_dict=False、最終 token-level LM-head 入力を返す output_hidden_states、vocabulary resize 時の decoder-side Embedder も契約に加えました。inputs_embeds は encoder / decoder が対称な config だけで使え、独立した decoder Embedder を持つ paper preset では input_ids が必要です。
最終 commit では、利用側が config.use_cache=True や return_dict=False、各種 output flag を設定していても、named generation、serving、bottleneck 計測の内部 full pass に漏れないよう、内部 forward option を明示的に固定しました。
一方、PHOTON の multi-rate cache は標準の transformers.Cache と同じものではありません。standard past_key_values は黙って無視せず明示的にエラーにし、stock generate() は cache-free full recompute、generation_mode=... を指定したときだけ HierarchicalCache を使うように境界を残しました。
ローカルで確認する
最小限の setup と test は次の通りです。
git clone https://github.com/arumada-5u/verify-photon.git
cd verify-photon
git checkout fdcd225a98d04b58d7f2802f790c1b44d7c84417
uv venv
uv pip install torch --torch-backend=auto
uv pip install -e ".[dev,train,eval]"
uv run pytest tests/ -q
掲載した環境は Python 3.12.3、PyTorch 2.11.0+cu128、Transformers 5.13.1 です。上の install は project の許容範囲を解決する手順で、同じ version を固定する lockfile ではありません。test は CPU でも実行できますが、以下の学習と --device cuda の benchmark には CUDA GPU と dataset / tokenizer の download が必要です。
学習・ベンチマークの実行例
tiny の学習は architecture だけを切り替えて実行できます。
uv run python training/train.py \
--config training/configs/tiny.yaml --model_type photon
uv run python training/train.py \
--config training/configs/tiny.yaml --model_type vanilla
uv run python training/train.py \
--config training/configs/tiny_paper.yaml --model_type photon
生成 schedule の gap と throughput は、それぞれ独立した JSON に出力します。
uv run python evaluation/recgen_gap.py \
--model runs/tiny-photon/final \
--config training/configs/tiny.yaml \
--output results/recgen-gap-tiny-rerun.json \
--num_windows 32 --prompt_len 64 --horizon 192 \
--num_gen_prompts 8 --seed 0
uv run python evaluation/throughput.py \
--scale tiny --device cuda --dtype bf16 \
--batch-sizes 1 8 32 --regimes pf de --iters 2 \
--out results/tpm-tiny-rerun.json
uv run python evaluation/serving_throughput.py \
--scale tiny --models photon block vanilla \
--concurrencies 1 8 32 --prompt-lengths 128 --new-tokens 256 \
--warmup-iters 1 --iters 2 \
--out results/serving-tpm-tiny-rerun.json
上の recgen_gap.py、throughput.py、serving_throughput.py は、現時点では既存の出力先を拒否しません。公開済み record を誤って上書きしないよう、再実行時は新しい出力名を指定しています。
まとめ
main までで確認できたのは、PHOTON の効率化が少なくとも実装上の見かけだけではないことです。因果性を保ったまま、RecGen では成長する状態を最上位のストリームに絞ることができ、同じ serving 条件では flat KV より小さい状態と高い decode throughput を観測しました。
ただし、それは論文の主要結果を再現したことと同義ではありません。tiny では品質差が残り、整合性 loss を小さくするだけでは collapse も防げませんでした。134B-token 学習を走らせていない以上、論文の PPL も 475 倍も自分の結果としては書けません。
今回いちばん大きかった収穫は、成功した数字そのものより、主張をどこで区切るべきかが明確になったことでした。次は、重複のない十分な学習データを用意し、FLOPs と seed を揃えたうえで、品質と状態量を同じ実験の中で比べたいと考えています。
つよつよGPUが欲しい…