Register Tokens for Bounded-State Reasoning in Diffusion Language Models
This repository contains a diffusion language model post-trained with register tokens for bounded-state reasoning, as described in the paper Register Tokens for Bounded-State Reasoning in Diffusion Language Models.
The model is based on the Dream architecture and uses a fixed-size set of register tokens to carry reasoning state across generation chunks. Between chunks, the generated text can be cleared while preserving the register values, enabling bounded-state continuation after context resets.
Code and inference examples are available in the GitHub repository.