DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models

TL;DR

This paper introduces DREAMSTATE, a framework that models and edits RNN states using diffusion techniques, and proposes a hybrid architecture combining RNNs with global context-aware parameter generation, advancing RNN interpretability and adaptability.

Contribution

It presents the DREAMSTATE framework for modeling RNN states with diffusion models and introduces a hybrid RNN architecture with dynamic, context-aware parameters.

Findings

DREAMSTATE effectively models RNN states and enables editing.

The hybrid architecture achieves stable training with multi-objective loss.

Experimental results validate the model's ability to adapt to global context.

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.