PhotonBERT

A bidirectional encoder that adapts the PHOTON hierarchy to BERT-style masked-language-model pretraining.

PhotonBERT-30M-1B — 29,662,888 parameters, 1 B training tokens, Phase 1 only, ModernBERT tokenizer (50,368 vocab), sequences up to 1,024 tokens, English.

The original work

PHOTON (Ichikawa et al., 2025) replaces a Transformer's horizontal token scan with a vertical, multi-resolution one. Four modules form the hierarchy:

• context chunker — concatenates C_l adjacent states into one coarse state;
• context encoder — a Transformer over the reduced-rate chunk sequence;
• context converter — expands one coarse state into R_l conditioning vectors;
• context decoder — reconstructs each chunk independently from that prefix, so local attention length is bounded by R_l + C_l regardless of sequence length.

The decoder-only model is trained with L = L_token + α·L_rec + β·L_context, where L_rec aligns each top-down reconstruction with its bottom-up state and L_context makes chunk states predictable from their causal prefix. Reported models use L = 2, C₁ = C₂ = 4, and α = β = 0.

PhotonBERT keeps that hierarchy and trains it with MLM:

masked token embeddings                T × H
  ├─ chunker (4) → level-1 context encoder      T/4 × H
  ├─ chunker (4) → level-2 context encoder      T/16 × H
  ├─ converter → level-2 chunk-local decoder    T/4 × H
  ├─ converter → level-1 chunk-local decoder    T × H
  └─ MLM / classification / retrieval head

Unchanged from the paper: L = 2, C₁ = C₂ = 4, the four-module decomposition (Eqs. 2–8), X̂⁽ᴸ⁾ := X⁽ᴸ⁾ as the start of the top-down pass, the R_l + C_l attention bound, and bottom-up states as reconstruction teachers.

Differences from the paper

Every departure is intentional. PhotonBERT is a re-derivation for masked bidirectional encoding, not a reproduction of the generative model.

Architecture

• Causal attention → bidirectional at every level; MLM removes the target from the input, so full context leaks nothing.
• The converter conditions chunk g on its own coarse state, not the previous one (Eq. 6); that shift exists only for autoregressive causality.
• The local decoder receives only the converter prefix plus learned position queries, not the teacher-forced lower stream of Eq. 7 — feeding it to a bidirectional decoder would let residuals copy the cosine target unchanged. Accepted cost: every token prediction flows through the coarsest stream alone.
• R_l = 4 rather than the R_l = 2 implied by Table 6, to widen the prefix that now carries all chunk detail.
• One width D = n_embd and one tied embedding table, instead of D₀=416 / D₁=D₂=1664 with a lossless l=1 concat chunker and a second decoder-side embedder.
• The chunker projects then normalises without bias; GQA is available for the larger presets (30m is MHA, as is the paper).

Objective — L = L_MLM + α·L_rec

• L_token becomes standard 80/10/10 MLM.
• β·L_context (Eq. 12) is not implemented: it is defined over a causal prefix and has no bidirectional analogue. The paper sets β = 0 anyway.
• α = 0.3 for the released checkpoint (the code default is 0.1), instead of the paper's α = 0 — the top-down path is now the only route to the MLM head, so the reconstruction term is what keeps the levels aligned.
• L_rec is a per-level mean of 1 − cos summed with equal weight, not Eq. 11's single denominator (which weights the token level 4×).
• Reconstruction teachers are stop-gradient, and the token-level teacher embeds the clean IDs so [MASK] is not reconstructed toward [MASK].

Training — FineWeb-Edu 90.25 % / DCLM 4.75 % / StarCoderData 5 %, ModernBERT tokenizer, ≤1024 tokens, 1 B non-padding tokens, AdamW + cosine at lr 2e-3, 512 sequences per step, bf16, and an MLM mask ratio decaying linearly from 0.15 to 0.10 — versus 134 B Pile, Llama tokenizer, 2048 tokens, Adam at lr 3e-4, batch 256, 600 M / 1.2 B parameters.

Usage

The checkpoint ships its own modelling code, so trust_remote_code=True is required.

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

repo = "RikkaBotan/PhotonBERT-30M-1B"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True).eval()

text = "The capital of France is [MASK]."
batch = tokenizer(text, return_tensors="pt")
with torch.no_grad():
    logits = model(**batch).logits
index = (batch["input_ids"][0] == tokenizer.mask_token_id).nonzero()[0, 0]
print(tokenizer.decode(logits[0, index].topk(5).indices))

AutoModel and AutoModelForSequenceClassification are mapped as well. Sequences are padded internally to a multiple of 16 (the hierarchy's compression ratio) and trimmed back, so any input length is accepted. The masked-LM output head lives in mlm_head.safetensors alongside model.safetensors; from_pretrained warns if it is absent, in which case fill-mask predictions are meaningless.

Intended use and limits. This is a small research encoder for fine-tuning and for studying the PHOTON hierarchy under a masked objective — not a production-ready model. It is English-only, trained on 1 B tokens of web and code text, and inherits the biases of FineWeb-Edu, DCLM and StarCoderData. Because every token prediction is routed through a 16× compressed stream (see Differences from the paper), token-level tasks such as CoLA are noticeably weaker than the sentence-level ones.

Key settings.

chunk_sizes, prefix_lengths, encoder_layers and decoder_layers must be the same length, and the head dimension must be even for RoPE. The hierarchy pads internally to prod(chunk_sizes) and trims back; packed segments are aligned to 16 so a top-level chunk never crosses a document boundary. Only checkpoints carrying architecture_version: photonbert-v1 are loadable.

Results

RikkaBotan/PhotonBERT-30M-1B on GLUE dev, from scripts/evaluate_glue.py (glue_results_30m_1b.json). RTE, MRPC and STS-B start from an MNLI-tuned encoder (MNLI intermediate transfer, on by default); ± is the standard deviation over the seeds actually run for that task, out of [19, 8364, 717, 10536, 90166].

The average takes the mean of both reported metrics for MRPC, STS-B and QQP. Read these against the budget: 29.7 M parameters, 1 B tokens, Phase 1 only, and an MLM head fed exclusively through a 16× compressed stream.

Citation

@article{ichikawa2025photon,
  title={PHOTON: Hierarchical Autoregressive Modeling for Lightspeed and
         Memory-Efficient Language Generation},
  author={Ichikawa, Yuma and Takagi, Naoya and Nakagawa, Takumi and
          Kanazawa, Yuzi and Sakai, Akira},
  journal={arXiv preprint arXiv:2512.20687},
  year={2025}
}