DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models
This work addresses a gap in research on RNN state representation for AI/ML practitioners, offering a novel direction and architectural reference, though it appears incremental in building on existing RNN and diffusion methods.
The paper tackles the problem of editing and utilizing the internal state of recurrent neural networks (RNNs) like RWKV as a knowledge representation, proposing the DREAMSTATE framework with a conditional Diffusion Transformer to generate and edit states, and a hybrid architecture that combines RNNs with global context adaptability to dynamically adjust parameters, demonstrating stable training via multi-objective loss.
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.