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] が token t+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が欲しい…

参考