Test-Time Registers as Global Priors for Tokenized Image Generation

Cheng-Yao Hong^1,*, Yifan Wang^1,*, Yuewei Lin², Chenyu You^1
1 Stony Brook University    ² Brookhaven National Laboratory
^* Equal contribution.

---

Overview

RegToken extracts register structure from vision-transformer activations and uses it to construct compact global prior tokens. The release provides CPU-testable NumPy utilities for NFN-based localization, register-neuron scoring, TokenRank and WriteMass analysis, spectral diagnostics, and the projection-and-conservation construction update. Core analysis APIs operate on caller-supplied activation, attention, and value arrays. Model-specific OpenCLIP and DINOv2 hook pipelines are not bundled. It also provides a minimal checkpoint-backed VQ-LL-32 TiTok runtime for one-image encoding and decoding.

Available Functionality

• NFN scoring and candidate-layer ranking from Eqs. (3)–(5)
• high-norm token localization and register-neuron scoring from Eq. (6)
• row-stochastic transition validation, TokenRank, and WriteMass from Eq. (7)
• projection-and-conservation construction from Eqs. (8)–(9)
• continuous TiTok latent validation, codebook lookup, and slot preparation
• VQ-LL-32 checkpoint loading, one-image encode/decode, and image saving
• optional caller-aligned prior injection and selected-slot optimization
• spectral, frequency, artifact-validation, and CPU analysis helpers

See the method mapping for the equation-to-API index.

Repository Structure

src/regtoken/              Core RegToken utilities
src/regtoken/adapters/     TiTok latent and prior interfaces
src/regtoken/runtime/      Checkpoint-backed VQ-LL-32 runtime
examples/                  One-image runtime examples
scripts/                   CPU analysis and dry-run scripts
tests/                     Unit and runtime tests
docs/                      Project website

Installation

Python 3.10 or newer and NumPy 1.24 or newer are required.

python -m pip install .

The tensor adapter requires PyTorch 2.1 or newer. The checkpoint-backed runtime also requires Pillow, OmegaConf, Einops, Hugging Face Hub, and Safetensors:

python -m pip install '.[adapter]'
python -m pip install '.[runtime]'

Quick Start

Run the CPU dry run from the repository root:

PYTHONPATH=src python scripts/dry_run.py

Run the test suite with:

PYTHONPATH=src python -m unittest discover -s tests/unit -v

Basic API usage:

import numpy as np

from regtoken import construct_regtoken, tokenrank, write_mass

attention = np.array([[0.75, 0.25], [0.50, 0.50]], dtype=np.float64)
values = np.array([[3.0, 4.0], [0.0, 2.0]], dtype=np.float64)

mass = write_mass(attention, values)
rank = tokenrank(attention)
result = construct_regtoken(
    values,
    outlier_patch_indices=[0],
    register_channels=[0],
    scale=2.0,
    alpha=0.5,
)

VQ-LL-32 TiTok Runtime

The runtime uses a caller-provided token-opt checkout and defaults to yucornetto/tokenizer_titok_l32_imagenet at the revision pinned in regtoken.runtime. Model files are retrieved through the Hugging Face cache when they are not already local; neither TiTok source nor checkpoint weights are stored in this repository.

Run a checkpoint-backed round trip:

PYTHONPATH=src python examples/titok_roundtrip.py \
  --token-opt-root /path/to/token-opt \
  --input input.jpg --output decoded.png --device cpu

examples/titok_with_prior.py accepts an aligned 12-dimensional .npy prior, plus explicit slot_index and gamma values. examples/titok_prior_opt.py demonstrates a callback-based Prior+Opt interface with an illustrative scalar callback. It optimizes only the selected slot and keeps the tokenizer, decoder, and all other slots fixed.

Prior injection expects a caller-supplied prior in TiTok's 12-dimensional feature space. Automatic alignment, bundled paper priors, and full quantitative ImageNet reproduction are not included. Slot index 0 is an adapter default; TiTok does not assign native global or register semantics to that slot.

Citation

@inproceedings{hong2026regtoken,
  title     = {Test-Time Registers as Global Priors for Tokenized Image Generation},
  author    = {Hong, Cheng-Yao and Wang, Yifan and Lin, Yuewei and You, Chenyu},
  booktitle = {European Conference on Computer Vision},
  year      = {2026}
}

Machine-readable citation metadata is provided in CITATION.cff.

Acknowledgements

Our register-neuron analysis was informed by the public test-time-registers implementation. We also thank the authors of token-opt and TiTok for releasing their code and models. These external sources and models are not included in RegToken; this acknowledgement is academic attribution, and RegToken's MIT License does not apply to them.

License

This project is released under the MIT License.