Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

WuNeng: Hybrid State with Attention

Liu Xiao, Li Zhiyuan, Lin Yueyu

TL;DR

WuNeng tackles the expressivity bottleneck of Transformer-based models by combining RWKV-7 state-driven heads with standard multi-head attention, enabling richer representations without large parameter increases. The approach uses cross-head interactions and a multi-token state processing mechanism to fuse attention and RWKV state, guided by the ARWKV training paradigm. Early results show WuNeng-7B achieving 10–15% improvements over strong baselines on benchmarks like MMLU and GSM8K and faster attention alignment convergence, suggesting substantial gains in reasoning and sequence generation tasks. The work points to a practical path toward more expressive yet efficient large language models and lays groundwork for future extensions to longer contexts, MoE, and multimodal settings.

Abstract

The WuNeng architecture introduces a novel approach to enhancing the expressivity and power of large language models by integrating recurrent neural network (RNN)-based RWKV-7 with advanced attention mechanisms, prioritizing heightened contextual coherence over reducing KV cache size. Building upon the hybrid-head concept from Hymba, WuNeng augments standard multi-head attention with additional RWKV-7 state-driven heads, rather than replacing existing heads, to enrich the model's representational capacity. A cross-head interaction technique fosters dynamic synergy among standard, state-driven, and newly introduced middle heads, leveraging concatenation, additive modulation, and gated fusion for robust information integration. Furthermore, a multi-token state processing mechanism harnesses the continuous RWKV-7 state to capture intricate, sequence-wide dependencies, significantly boosting expressivity. Remarkably, these enhancements are achieved with minimal additional parameters, ensuring efficiency while empowering the model to excel in complex reasoning and sequence generation tasks. WuNeng sets a new standard for balancing expressivity and computational efficiency in modern neural architectures.

WuNeng: Hybrid State with Attention

TL;DR

WuNeng tackles the expressivity bottleneck of Transformer-based models by combining RWKV-7 state-driven heads with standard multi-head attention, enabling richer representations without large parameter increases. The approach uses cross-head interactions and a multi-token state processing mechanism to fuse attention and RWKV state, guided by the ARWKV training paradigm. Early results show WuNeng-7B achieving 10–15% improvements over strong baselines on benchmarks like MMLU and GSM8K and faster attention alignment convergence, suggesting substantial gains in reasoning and sequence generation tasks. The work points to a practical path toward more expressive yet efficient large language models and lays groundwork for future extensions to longer contexts, MoE, and multimodal settings.

Abstract

The WuNeng architecture introduces a novel approach to enhancing the expressivity and power of large language models by integrating recurrent neural network (RNN)-based RWKV-7 with advanced attention mechanisms, prioritizing heightened contextual coherence over reducing KV cache size. Building upon the hybrid-head concept from Hymba, WuNeng augments standard multi-head attention with additional RWKV-7 state-driven heads, rather than replacing existing heads, to enrich the model's representational capacity. A cross-head interaction technique fosters dynamic synergy among standard, state-driven, and newly introduced middle heads, leveraging concatenation, additive modulation, and gated fusion for robust information integration. Furthermore, a multi-token state processing mechanism harnesses the continuous RWKV-7 state to capture intricate, sequence-wide dependencies, significantly boosting expressivity. Remarkably, these enhancements are achieved with minimal additional parameters, ensuring efficiency while empowering the model to excel in complex reasoning and sequence generation tasks. WuNeng sets a new standard for balancing expressivity and computational efficiency in modern neural architectures.
Paper Structure (13 sections, 12 equations, 3 figures, 1 table)

This paper contains 13 sections, 12 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Hybrid-Head Architecture
  5. Cross-Head Interaction
  6. Multi-Token State Processing
  7. Evaluation
  8. Experimental Setup
  9. Stage 1: Attention Alignment
  10. Stage 2: Knowledge Distillation
  11. Stage 3: SFT and DPO
  12. Discussion
  13. Conclusions

Figures (3)

  • Figure 1: Illustration of the WuNeng hybrid-head architecture, showcasing the integration of standard multi-head attention and RWKV-7 state-driven heads.
  • Figure 2: Visualization of the multi-token state processing mechanism in WuNeng. The figure depicts how the RWKV-7 state $S_t$ is updated with multi-token context from the input sequence , and subsequently modulates the attention mechanism to capture rich, sequence-wide dependencies for enhanced contextual coherence.
  • Figure 3: Stage 2 alignment loss curves comparing WuNeng-7B (green) and Qwen2.5-7B-Instruct (blue) during knowledge distillation with word-level KL-Divergence. WuNeng-7B converges to a lower loss ( 0.08) after 3k steps, demonstrating superior alignment of attention patterns with the teacher model, which contributes to its 10%--15% better benchmark performance in Stage 3 (Section 4.3).