Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models

Liu Xiao

TL;DR

This work addresses the limited editability of the internal RWKV state by proposing DREAMSTATE, which models the state as a probabilistic variable with a conditional diffusion transformer to generate $p(S|c)$ and enables direct state editing. It then introduces a hybrid architecture where a Diffusion Transformer dynamically synthesizes the RWKV WKV parameters conditioned on global context, fused with static parameters via a learnable interpolation, mitigating structural noise in recurrence. The approach is validated through state manifold visualization (t-SNE) and controlled generation experiments, plus a multi-objective training regime showing stable joint optimization of language modeling and parameter diffusion. Overall, the paper demonstrates that combining probabilistic state modeling with context-aware parameter synthesis can yield controllable, expressive recurrent world models with practical training stability and inference capabilities.

Abstract

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.

DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models

TL;DR

This work addresses the limited editability of the internal RWKV state by proposing DREAMSTATE, which models the state as a probabilistic variable with a conditional diffusion transformer to generate and enables direct state editing. It then introduces a hybrid architecture where a Diffusion Transformer dynamically synthesizes the RWKV WKV parameters conditioned on global context, fused with static parameters via a learnable interpolation, mitigating structural noise in recurrence. The approach is validated through state manifold visualization (t-SNE) and controlled generation experiments, plus a multi-objective training regime showing stable joint optimization of language modeling and parameter diffusion. Overall, the paper demonstrates that combining probabilistic state modeling with context-aware parameter synthesis can yield controllable, expressive recurrent world models with practical training stability and inference capabilities.

Abstract

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.
Paper Structure (22 sections, 5 equations, 3 figures, 1 table)

This paper contains 22 sections, 5 equations, 3 figures, 1 table.

Table of Contents

  1. INTRODUCTION
  2. Methodology: From Static Recurrence to Dynamic Synthesis
  3. Recurrence as Attention and the Problem of Structural Noise
  4. Part 1: Generative State Synthesis (DREAMSTATE)
  5. Modeling the State Manifold.
  6. State Representation and Training.
  7. Part 2: Dynamic Parameter Synthesis as Attention Control
  8. Generative Parameter Fields.
  9. Architectural Design.
  10. Multi-Objective Training.
  11. Experiments
  12. Experimental Setup
  13. Results for State Synthesis
  14. State Manifold Visualization.
  15. Controlled Inference.
  16. ...and 7 more sections

Figures (3)

  • Figure 1: The proposed hybrid architecture for dynamic parameter synthesis. A Diffusion Transformer (DiT) tracks global features from a variable-length context to generate context-specific WKV parameters, which are then fused with static parameters to modulate the RWKV recurrence.
  • Figure 2: t-SNE visualization of RWKV-7 states from different text prompts.
  • Figure 3: Training loss for the dynamic parameter synthesis architecture. The stable decrease of both the language modeling loss with the parameter diffusion loss demonstrates the viability of the multi-objective training.