DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models

Modern Recurrent Neural Networks (RNNs), such as RWKV, are distinguished by their powerful short-range modeling capabilities and efficient fixed-size states, which constitute a core advantage over standard Transformers. However, there is a significant lack of research into their internal state as an editable knowledge representation. To fill this gap, we first explore the representational properties of the RWKV state by proposing the DREAMSTATE framework. This framework utilizes a conditional Diffusion Transformer (DiT) to directly model the probability manifold of the state, enabling its generation and editing. The structural nature of this representation is validated through t-SNE visualizations and controlled generation experiments. After successfully uncovering and modeling the state’s representational potential, we further propose a novel hybrid architecture that combines the local advantages of RNNs with global context adaptability. This architecture features a parallel DiT that processes a variable-length global context to dynamically generate and adjust the core recurrent module’s WKV parameters, transforming the fixed recurrence mechanism into a context-aware dynamic function. Experiments demonstrate that this hybrid model can be trained stably via a multi-objective loss, validating its design feasibility. Our work not only opens a new research direction for RNN state representation but also provides a concrete architectural reference for future model design. The code is publicly available at: https://huggingface.co/2dgx41s/DreamState.

💡 Research Summary

The paper introduces a two‑stage diffusion‑based framework, DREAMSTATE, that treats both the hidden state and the core WKV parameters of the RWKV‑7 recurrent language model as probabilistic variables that can be learned, sampled, and edited. In the first stage, the authors collect (text, final‑state) pairs by running a pretrained RWKV‑7 model over a large corpus. They flatten the multi‑head state tensors into a single vector and feed it to a conditional Diffusion Transformer (DiT). The DiT is trained with the standard DDPM loss to predict added noise conditioned on the text embedding, thereby learning the conditional distribution p(S | c). t‑SNE visualizations of real states reveal distinct clusters corresponding to different persona prompts (e.g., programmer, poet), confirming that the state lies on a structured manifold. Controlled generation experiments show that initializing the recurrent model with a state sampled from DREAMSTATE (or interpolated between two states) can steer the model toward a desired mode, such as storytelling or blended concepts, demonstrating practical editability of the hidden state.

In the second stage, the authors address the “structural noise” caused by the static WKV parameters (the weighted‑key‑value matrices) that govern the recurrence. They propose a parallel DiT that processes a variable‑length global context to generate a context‑specific parameter vector θ_gen, which contains the flattened W, K, V matrices. This generated vector is linearly interpolated with the original static parameters θ_static using a learnable mixing coefficient α, yielding the final parameters θ_WKV‑final. The whole system is trained end‑to‑end with a multi‑objective loss: λ₁ L_LM (standard next‑token cross‑entropy) plus λ₂ L_param (the DDPM loss for the parameter diffusion). Training on a subset of the Pile shows stable convergence of both losses, and the hybrid model achieves a modest perplexity reduction compared to the baseline static‑parameter RWKV‑7.

Key contributions are: (1) the DREAMSTATE framework that models the recurrent hidden state as a generative diffusion object, validated by visualization and controlled inference; (2) a novel architecture that dynamically synthesizes the core recurrence parameters from global context, effectively canceling the static‑parameter structural noise; (3) a joint training scheme that balances language modeling and parameter generation; and (4) empirical evidence that both state‑level and parameter‑level diffusion improve controllability and performance.

The work opens a new research direction for recurrent large language models: rather than treating recurrence as a fixed, opaque mechanism, it can be made adaptive and generative through diffusion models. Future work may scale the DiT, explore richer context encoders, and integrate state and parameter diffusion into a single unified generative process, potentially enabling full “generate‑and‑control” loops for long‑range, multi‑modal, and highly structured generation tasks.