ChiTako/unofficial-photon-repro-ja-8b-a1b-v3-nvfp4-hier

Unofficial PHOTON reproduction -- Japanese 8.1B-A1.2B

Unofficial. This is an independent, hobby reimplementation of the PHOTON architecture described in arXiv:2512.20687, trained from scratch by a private individual. It is not released by, affiliated with, or endorsed by the paper's authors or any organisation, and no weights, data or code from any official PHOTON release were used.

Read this before using it. It was trained on roughly 0.36B tokens against 8.06B parameters -- about 0.3 tokens per active parameter, where the Chinchilla-optimal figure is ~20. It is two orders of magnitude undertrained. This checkpoint demonstrates that the architecture and the training recipe work at this scale; it is not a competitive model and should not be compared to one.

A Japanese language model on the PHOTON hierarchical autoregressive architecture, with the transformer stacks replaced by fine-grained Mixture-of-Experts using auxiliary-loss-free load balancing.

What is different about this version

• —Expert capacity follows the token flow. Rows a stack sees per micro-batch order as L1 decoder (2S) > L2 decoder (S/2) > L1 encoder (S/4) > L2 encoder (S/16), and expert counts now follow that order. The previous 8B config put 256 experts on the last one -- 68% of the model -- and its load-imbalance metric (MaxVio) climbed monotonically past 30, meaning that capacity was inert.
• —llm-jp-tokenizer v4 (196,608). Adopted because logit-level distillation needs the teacher's exact vocabulary, and because v4 carries the openai-harmony control tokens that make the agent format possible. It is not better at compressing Japanese -- measured on real Japanese Wikipedia it needs 1.7915 chars/token against v3's 1.8320. d_token dropped 2048 -> 1024 to pay for the doubled vocabulary at the same embed+head cost.
• —Distillation, in two forms. Sequence-level from frontier open models via their published transcripts (Kimi K3, DeepSeek-V4-pro, poolside Laguna S 2.1, DeepSeek-V3.2, GLM-5.1) plus llm-jp-4's own SFT corpus. Logit-level KD needs an exact tokenizer match, which is the reason the vocabulary is llm-jp v4.
• —Japanese is protected on purpose. Every frontier agentic corpus measured 0.0% Japanese, so the mid-training mixture caps English agentic data at 0.18 and carries 0.74 Japanese. See docs/distillation.md.

How good is it, honestly

Held-out Japanese Wikipedia (evaluated beyond record 700,000, which is past everything this run consumed), against the two 250M models from the same code base. Perplexity is per token and these models use different vocabularies, so the only column that compares across rows is bits per character:

This model is the worst of the three. Scaling parameters 32x while scaling tokens 1.7x buys nothing: 0.04 tokens per parameter against the 250M v4's 0.8. The architecture work in this version is real and the routing measurements hold, but the binding constraint on this project is the token budget, and no amount of architectural care substitutes for it.

Size

The FLOP-equivalent figure is the one that governs speed: a level-l encoder runs once per C_<=l tokens, so its cost is amortised. This model thinks with 8B parameters and costs about what a 0.6B dense model costs.

Generation protocols

• —HierGen keeps encoder state at every level. Exact -- it reproduces the training-time distribution.
• —RecGen keeps only the top-level KV cache and feeds the decoder cascade its own reconstructions, cutting KV traffic by roughly 14x. It is exact when recursive consistency holds, which is what the training objective L_token + alpha * L_rec optimises for.

Usage

from photon_jp.model.photon import PhotonForCausalLM
from photon_jp.model.config import PhotonConfig
from photon_jp.infer.generate import PhotonGenerator, GenerationConfig

cfg = PhotonConfig.load("model_config.json")
model = PhotonForCausalLM(cfg)
# load model.safetensors, then:
gen = PhotonGenerator(model, "cuda")
out = gen.generate(input_ids, GenerationConfig(mode="recgen", max_new_tokens=256))

Agent use

The bundled tokenizer speaks openai-harmony, so tool calls round-trip through the chat template. Two things a harness must get right:

1. 1.Stop on `<|call|>` as well as `<|return|>`. generation_config.json lists both. Stopping only at <|return|> lets the model run past its own tool call and invent the tool's reply.
2. 2.Parse the channels. analysis is private reasoning, commentary with a to= recipient is a tool call, final is the user-visible answer.

from photon_jp.infer.agent import parse_harmony, append_tool_result

turn = parse_harmony(generated_text)
if turn.wants_tool:
    call = turn.tool_calls[0]                  # .name, .arguments (parsed JSON)
    messages = append_tool_result(messages, call, my_tools[call.name](**call.arguments))
else:
    print(turn.final)                          # turn.reasoning holds the analysis

Quantised variants

All five were evaluated on the same held-out slice with the same 100 sequences, which matters: the first pass measured bf16 with 200 sequences and the quantised variants with 100, and the sampling difference alone made fp8 look better than bf16. Re-measured on equal footing the ordering is what physics requires.

NVFP4 costs 1.7% perplexity for 3.15x less space. FP8 costs 0.2% for 1.9x. MXFP4's coarser block and power-of-two scale cost 2.3% — the price of being readable outside Blackwell.

-nvfp4-hier is reported as measured rather than as intended: keeping the level-2 encoder at FP8 lowered the RecGen/HierGen KL (280 -> 276) but left perplexity identical to plain NVFP4, so the extra 0.55 GiB bought nothing this evaluation can see. The idea behind it is sound — that stack runs once per 16 tokens, so 8-bit there is nearly free in FLOPs, and weight absorption for the latent KV cache multiplies its MLA matrices together — but it is not supported by this measurement.

Routers, embeddings, the LM head and all 1-D tensors stay in bf16 in every variant (409M parameters). Top-k routing over 192 experts turns on margins finer than 4-bit resolution, and a wrong pick runs a different expert rather than degrading gracefully.

A caveat on RecGen agreement: it is low (23-30%) in every variant including bf16, so quantisation is not the cause. RecGen has failed to reproduce HierGen in all four model sizes built from this code base; it is a property of the training objective, not of the precision. See docs/findings.md 4b.

Accounting

==============================================================================
PHOTON-JP parameter report
==============================================================================
vocab=196,608  D0=1024  L=2  C_<=L=16  ctx=8192

stack                    total        active   amort     amort.act
------------------------------------------------------------------------------
L1.encoder            4342.1M       293.6M       4        73.4M
L1.decoder            1165.8M        56.9M       1        56.9M
L2.encoder            1333.8M       295.7M      16        18.5M
L2.decoder             756.8M       105.6M       4        26.4M
------------------------------------------------------------------------------
embedding              201.3M
lm_head                201.3M
chunk/convert           50.4M
mtp                     10.8M
==============================================================================
TOTAL          : 8.062 B   (7.660 B non-emb)
ACTIVE / token : 1.205 B   (0.802 B non-emb)
AMORTISED      : 0.628 B  (FLOP-equivalent dense size)
sparsity       : 6.69x

fwd FLOPs/token @ ctx=8192: 0.81 GFLOP (matmul 0.75, attn 0.055)
KV cache mode    : MLA latent (weight-absorbed)
KV cache  HierGen: 6.469 KiB/token   (other mode: 13.500)
KV cache  RecGen : 0.469 KiB/token   (other mode: 7.500)
  -> 51.8 MiB for a full 8192-token context (HierGen, per sequence)
==============================================================================