U-RWKV: Unified & U-Shaped RWKV Models

Updated 4 July 2026

U-RWKV is a family of architectures that extend RWKV’s linear-time global mixing into unified and U-shaped vision pipelines for tasks like low-light restoration and medical segmentation.
It adapts canonical sequence modeling through innovations such as LAN, DARM, and SSF, providing targeted solutions for dynamic degradation and directional feature fusion.
These models balance efficiency and performance, achieving competitive quantitative metrics with low computational overhead in both restoration and segmentation applications.

U-RWKV is not a single canonical variant in the original RWKV literature; rather, it denotes a cluster of later RWKV-derived architectures whose names emphasize either unified modeling or U-shaped vision segmentation design. The original RWKV model, introduced for sequence modeling, already presented itself as a unification of Transformer-like parallel training and RNN-like efficient autoregressive inference (Peng et al., 2023). Subsequent papers then reused related nomenclature in domain-specific settings, most notably "URWKV: Unified RWKV Model with Multi-state Perspective for Low-light Image Restoration" (Xu et al., 29 May 2025), "U-RWKV: Lightweight medical image segmentation with direction-adaptive RWKV" (Ye et al., 15 Jul 2025), and "Med-URWKV: Pure RWKV With ImageNet Pre-training For Medical Image Segmentation" (Zhou, 12 Jun 2025). In that literature, U-RWKV most often refers to RWKV-based U-shaped vision systems that preserve RWKV’s linear-time global mixing while adapting it to dense prediction.

1. Terminological scope and naming

The term “U-RWKV” does not appear in the original language-model paper "RWKV: Reinventing RNNs for the Transformer Era" (Peng et al., 2023). The survey "A Survey of RWKV" likewise states that it does not explicitly mention a variant called “U-RWKV,” and instead treats the closest “unification” effort as a generalized implicit-attention view of gated-linear recurrent sequence models (Li et al., 2024). Within later application papers, however, closely related names are used in distinct but overlapping ways: URWKV for low-light image restoration, U-RWKV for lightweight medical segmentation, and Med-URWKV for a pure RWKV U-Net with pre-trained encoder reuse.

This naming pattern indicates that U-RWKV is best understood as a family label in practice, not as a single formally standardized architecture. In one usage, “U” means Unified, referring to the handling of coupled degradations in low-light restoration. In another, it denotes a U-shaped encoder–decoder segmentation framework. A related but separate interpretation appears in ultra-compact vision RWKV work such as SCRWKV, although that paper does not itself use the name U-RWKV (Zhang et al., 14 May 2026).

Variant	Domain	Distinguishing characteristics
URWKV	Low-light restoration	Multi-state perspective; LAN, SQ-Shift, SSF; 2.25M parameters, 18.34G FLOPs
U-RWKV	Medical image segmentation	DARM and SASE; 2.97M parameters, 7.28 GFLOPs
Med-URWKV	Medical image segmentation	Pure RWKV U-Net with pretrained VRWKV encoder; 14.33M parameters
RWKV-UNet	Medical image segmentation	IR-RWKV encoder and CCM skip fusion; 17.4M parameters, 14.64G FLOPs

A recurrent misconception is that U-RWKV names a single official successor to RWKV-4/5/6. The literature provided here does not support that reading. Instead, the original RWKV paper frames RWKV itself as the relevant unification, while later vision papers specialize the mechanism for restoration and segmentation (Peng et al., 2023).

2. Canonical RWKV basis

All U-RWKV variants inherit the core RWKV mechanism: a residual stack with a time-mixing sub-block and a channel-mixing sub-block. RWKV replaces quadratic token-token self-attention with a channel-wise, time-decayed accumulation that is exact and can be written as either a parallel scan for training or a recurrent update for inference. The central time-mixing equations are

$r_t = W_r \cdot (\mu_r \odot x_t + (1 - \mu_r) \odot x_{t-1}),$

$k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$

and

$wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$

The output is then gated by receptance,

$o_t = W_o \cdot (\sigma(r_t) \odot wkv_t),$

while the channel-mixing branch uses a gated MLP-like transformation with squared ReLU. In recurrent form, RWKV maintains numerator and denominator accumulators,

$a_t = e^{-w} \odot a_{t-1} + e^{k_t} \odot v_t, \qquad b_t = e^{-w} \odot b_{t-1} + e^{k_t},$

so that per-token autoregressive inference is $O(d)$ in both time and memory (Peng et al., 2023).

This recurrent reformulation is the technical substrate from which later vision U-RWKV systems are derived. The original paper states RWKV (ours): Time $O(Td)$ , Space $O(d)$ (inference) and reports models up to 14 billion parameters, trained for one epoch on The Pile (≈330B tokens), with competitiveness against similarly sized Transformers across 12 standard NLP tasks. It also identifies characteristic trade-offs: the per-channel accumulators create an information bottleneck relative to full quadratic self-attention, and prompt ordering can materially affect behavior, as illustrated by a zero-shot RTE F1 shift from 44.2% to 74.8% after reordering instructions (Peng et al., 2023).

For U-RWKV applications, the main significance of canonical RWKV is architectural rather than lexical. The later models reuse its two defining properties: linear-time global mixing and small recurrent state, while modifying token shift, decay control, scan direction, or skip-fusion strategy to suit 2D visual structure.

3. URWKV in low-light image restoration

In the restoration literature, the most explicit “unified” usage is URWKV, a U-shaped encoder–decoder for low-light image enhancement and joint low-light enhancement and deblurring. Its design addresses dynamically coupled degradations, including noise amplification under brightness enhancement, motion blur induced by long exposure or camera shake, color cast/distortion, contrast loss, non-uniform illumination, and strong light-dark contrast. The network has three stages in encoder and three stages in decoder; each encoder stage contains $N_1=3$ URWKV blocks, each decoder stage contains $N_2=2$ URWKV blocks, and the base channel width is $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 0 (Xu et al., 29 May 2025).

The core URWKV block uses a multi-state spatial mixing sub-block and a multi-state channel mixing sub-block. Three modules distinguish it from canonical Vision-RWKV. First, Luminance-adaptive Normalization (LAN) modulates normalization parameters using inter-stage luminance states. With aggregated luminance representation $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 1, the modulator is

$k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 2

and the normalized output is

$k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 3

Second, multiple intra-stage states are accumulated by EMA,

$k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 4

with $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 5 in experiments, and then passed to Multi-state Quad-directional Token Shift (SQ-Shift). Third, State-aware Selective Fusion (SSF) replaces naive skip connections by dynamically aligned fusion across encoder stages, using adaptive spatial weighting rather than fixed add/concat (Xu et al., 29 May 2025).

The reported efficiency is a central part of the model definition. URWKV uses 2.25M parameters and 18.34G FLOPs, compared with Restormer (26.11M, 140.99G) and MIRNet (31.76M, 785G). The training setup uses PyTorch, a single NVIDIA Tesla A40 GPU, Adam with $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 6 and $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 7, an initial learning rate of $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 8 cosine-annealed to $k_t = W_k \cdot (\mu_k \odot x_t + (1 - \mu_k) \odot x_{t-1}), \qquad v_t = W_v \cdot (\mu_v \odot x_t + (1 - \mu_v) \odot x_{t-1}),$ 9, and flipping and rotation of paired images. Inference supports arbitrary input shape without cropping/resizing (Xu et al., 29 May 2025).

Quantitatively, URWKV reports 23.11 dB / 0.874 SSIM on LOL-v2-real, 26.36 dB / 0.944 SSIM on LOL-v2-syn, 31.24 dB / 0.911 SSIM on SDSD-indoor, 29.99 dB / 0.887 SSIM on SDSD-outdoor, 26.08 dB / 0.936 SSIM on MIT-Adobe FiveK, and 27.27 dB / 0.890 SSIM on LOL-blur; on SID it reaches 23.11 dB / 0.673 SSIM, reported as second-best and closely trailing Retinexformer. Ablations on LOL-v2-real attribute the full 23.11 dB / 0.874 result to the joint use of LAN + SSF and to Multi-state + Q-shift, while naive skip configurations or single-state designs are weaker (Xu et al., 29 May 2025).

Within the U-RWKV family, URWKV is the clearest example of “unified” meaning not merely architectural reuse but explicit modeling of multiple degradation states. Its significance lies in extending RWKV beyond sequence or generic vision backbones into a scene-adaptive restoration pipeline where luminance modulation, state aggregation, and selective multi-stage fusion are first-class components.

4. U-RWKV in lightweight medical image segmentation

In medical image segmentation, U-RWKV denotes a lightweight U-shaped encoder–decoder built around two modules: the Direction-Adaptive RWKV Module (DARM) and the Stage-Adaptive Squeeze-and-Excitation Module (SASE). The stated motivation is that U-Net variants have limited global Effective Receptive Fields (ERFs), whereas Transformer alternatives increase ERF at $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 0 complexity. U-RWKV instead targets long-range modeling at $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 1 computational cost (Ye et al., 15 Jul 2025).

DARM operates on encoder feature maps $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 2 using QuadScan and Dual-RWKV. QuadScan extracts four directional sequences—left-to-right, right-to-left, top-to-bottom, and bottom-to-top—then reconstructs 2D features and averages them. Dual-RWKV runs RWKV in forward and reversed order without weight sharing. The paper’s spatial mixing writes the bidirectional visual RWKV operator as

$wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 3

with learned $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 4. This formulation is then combined with channel mixing after reconstruction to produce context-enriched features for the decoder. SASE adapts squeeze-and-excitation behavior by stage: shallow stages use lightweight pointwise convolutions and inverted residual structures with SERatio set to 4, while deep stages split channels into 8 groups and apply depthwise separable convolutions, making SERatio effectively $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 5 (Ye et al., 15 Jul 2025).

The model is trained for 280 epochs on a single NVIDIA RTX 3090 GPU. Synapse uses the official split; the remaining datasets use a 70/30 train/val split. RWKV initialization comes from Vision-RWKV, and the training settings follow CMU-Net and CMUNeXt. The paper does not enumerate optimizer, loss, or augmentation, and explicitly notes that these details are not specified (Ye et al., 15 Jul 2025).

On five 2D datasets—BUSI, Kvasir, CVC-ClinicDB, ISIC 2017, ISIC 2018—the main model reports 2.97M parameters, 7.28 GFLOPs, and Avg Dice 82.27, with per-dataset Dice of 82.34, 88.17, 90.58, 90.13, and 87.26 respectively. The tiny variant U-RWKV-s uses 0.46M parameters and 1.02 GFLOPs with Avg Dice 80.06. On Synapse multi-organ segmentation, U-RWKV reaches Dice 80.64 and HD95 26.61. ERF analysis reports a high-contribution area ratio at threshold 0.99 of 0.992 for U-RWKV, compared with 0.568 for U-Net and 0.946 for ViT/VMamba bottleneck variants (Ye et al., 15 Jul 2025).

Ablations assign the strongest contribution to DARM. On average IoU over BUSI, Kvasir, ISIC’17, removing DARM reduces performance to 74.97, removing SASE to 75.29, and removing Dual RWKV to 76.68, whereas the full model reaches 77.62. Single-direction scans are weaker than the full directional combination, indicating that directional bias is materially relevant in 2D medical structure modeling (Ye et al., 15 Jul 2025).

This version of U-RWKV therefore uses RWKV not as a generic backbone replacement but as a direction-sensitive global aggregator inserted into a conventional U-shaped segmentation hierarchy. Its principal claim is that linear-time scan-based global context can enlarge ERF and improve Dice while remaining lightweight enough for resource-constrained deployment.

5. Pure and hybrid U-shaped RWKV variants

A broader U-RWKV lineage in medical imaging includes both pure RWKV and hybrid CNN–RWKV systems. Med-URWKV defines U-RWKV explicitly as a U-Net-style segmentation architecture built entirely from RWKV blocks. Its encoder is a pretrained VRWKV-Tiny backbone taken from upernet_vrwkv_adapter_tiny_512_160k_ade20k, itself trained first on ImageNet and then fine-tuned on ADE20K with a ViT-Adapter. Only the encoder is retained; the bottleneck, decoder, and segmentation head are initialized randomly. The encoder is frozen for the first 5 epochs and then unfrozen for end-to-end training. Inputs are resized to 512×512 RGB, training uses AdamW with initial learning rate $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 6, weight decay $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 7, batch size 8, 200 epochs, and a poly schedule on a single GeForce RTX 3090 24GB GPU (Zhou, 12 Jun 2025).

Med-URWKV reports 14.33M parameters and competitive or superior results across seven datasets. The tabulated scores include BUSI 77.98 / 67.37, ISIC-2017 84.91 / 76.44, ISIC-2018 89.85 / 82.91, Kvasir-SEG 90.93 / 85.41, GLAS 91.02 / 84.24, and TDD 89.13 / 81.26 in DSC / IoU. Its most striking analysis concerns pretraining: on BUSI, the best DSC is 77.98 with pretraining versus 51.17 without, and even at 10 epochs the gap is 31.04 vs 19.47, indicating that encoder reuse is not merely a convenience but a major determinant of convergence and final accuracy (Zhou, 12 Jun 2025).

By contrast, RWKV-UNet is not pure RWKV. It keeps CNN efficiency in early stages and decoder design while injecting Vision-RWKV Spatial Mix into the encoder through the Inverted Residual RWKV (IR-RWKV) block and strengthening skip connections through the Cross-Channel Mix (CCM) module. The encoder uses pure CNN stages first, then switches to IR-RWKV in later stages; the decoder remains CNN-only. End-to-end variants are RWKV-UNet-T: 3.3M params, 3.78G FLOPs, RWKV-UNet-S: 9.9M params, 7.75G FLOPs, and RWKV-UNet (Base): 17.4M params, 14.64G FLOPs. The base model achieves Synapse average DSC 84.02% and HD95 15.70, ACDC DSC 92.17, and strong binary segmentation results such as BUSI 81.92±0.82, CVC-ColonDB 92.27±2.48, and GLAS 92.35±0.07. At 512×512 on Synapse, it reports 85.52 DSC at 58.6G FLOPs, versus TransUNet 84.36 DSC at 168.7G FLOPs (Jiang et al., 14 Jan 2025).

Taken together, these systems show that “U-RWKV” can denote two distinct architectural philosophies. In the strictest sense used by Med-URWKV, it means a pure RWKV U-Net. In a broader applied sense, it also covers hybrid U-shaped models in which RWKV supplies linear-time global cooperation while convolutional operators retain local inductive bias. The literature does not impose a single boundary between these usages.

The survey literature places RWKV inside a wider class of attention-free or gated-linear recurrent models and points to “A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence Models” as the closest formal unification beyond the original architecture (Li et al., 2024). Within the RWKV lineage itself, later variants such as RWKV-5 (Eagle) and RWKV-6 (Finch) are described as architectural unifiers: Eagle introduces multi-head matrix-valued states with stabilized decay $wkv_t = \frac{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} \odot v_i + e^{u+k_t} \odot v_t }{ \sum_{i=1}^{t-1} e^{-(t-1-i)w+k_i} + e^{u+k_t} }.$ 8, and Finch makes the decay dynamic through data-dependent token shift (“ddlerp”) and LoRA-style adaptation (Li et al., 2024). These developments matter for U-RWKV because vision variants frequently inherit the same design pressure: broaden state expressiveness without losing RWKV’s linear-time character.

A related but separate branch is SCRWKV, an Ultra-Compact Structure-Calibrated Vision RWKV for crack segmentation. Although the paper does not use the name U-RWKV, it is explicitly described as Ultra-Compact and Vision-RWKV, with 1.22M parameters, 28MB, 22.78 GFLOPs at 512×512, and 46 FPS (0.0216 s/frame) on a resource-constrained server. Its backbone combines AMCM, GBST, and Dy-WKV with DSCD, and on TUT it reports F1=0.8428 and mIoU=0.8512 (Zhang et al., 14 May 2026). This suggests that ultra-compact, topology-aware RWKV segmentation is an active adjacent direction, even when the exact U-RWKV label is absent.

The principal limitations remain consistent across the family. At the canonical level, RWKV’s per-channel recurrent accumulation can be less expressive than full token-token attention for pinpoint recall of distant details, and the original paper notes greater prompt-order sensitivity than GPT-family models (Peng et al., 2023). At the survey level, RWKV exhibits challenges on extremely long-context and in-context learning benchmarks such as LooGLE, RULER, and S3EVAL, and performance can degrade with very long inputs and complex long-range dependencies (Li et al., 2024). In task-specific U-RWKV systems, limitations are domain-shaped rather than purely algorithmic: URWKV trails the best LLIE baseline on SID (Xu et al., 29 May 2025); the medical U-RWKV paper notes inference is slightly slower than pure CNNs like UNeXt and lists 3D segmentation as future work (Ye et al., 15 Jul 2025); Med-URWKV highlights the need to study larger pretrained encoder scales and RWKV-specific spatial mixing tailored to segmentation (Zhou, 12 Jun 2025).

Overall, U-RWKV denotes a technically coherent but nomenclaturally heterogeneous research direction: RWKV-derived models embedded in unified restoration or U-shaped dense-prediction pipelines. What links these systems is not a single official specification, but a common engineering premise—replace quadratic attention with RWKV-style linear global mixing, then adapt the resulting recurrence to multi-stage vision structure, directional scanning, luminance-aware normalization, or pretrained encoder reuse.