MRT: Learning Compact Representations with Mixed RWKV-Transformer for Extreme Image Compression

TL;DR

This work tackles extreme image compression by learning compact 1-D latent representations through a Mixed RWKV-Transformer (MRT) that combines global long-range modeling via linear-attention RWKV with local intra-window processing via Transformer blocks. A dedicated RWKV Compression Model (RCM) compresses these 1-D latents, supported by a two-stage training strategy that aligns latent codes to discrete targets and then fine-tunes decoding with pixel- and perceptual-supervised losses. Empirically, MRT achieves superior rate-distortion performance at ultra-low bitrates (below $0.02$ bpp), delivering significant bitrate savings over 2-D architectures on Kodak and CLIC2020 (e.g., $43.75 ext{%}$ and $30.59 ext{%}$ on DISTS) while preserving perceptual quality. The approach demonstrates the practical potential of 1-D latent representations and global-local hybrid modeling for efficient image compression, with a dedicated 1-D entropy model enabling effective compression of the learned latents. Future work aims to reduce computational complexity while maintaining RD gains.

Abstract

Recent advances in extreme image compression have revealed that mapping pixel data into highly compact latent representations can significantly improve coding efficiency. However, most existing methods compress images into 2-D latent spaces via convolutional neural networks (CNNs) or Swin Transformers, which tend to retain substantial spatial redundancy, thereby limiting overall compression performance. In this paper, we propose a novel Mixed RWKV-Transformer (MRT) architecture that encodes images into more compact 1-D latent representations by synergistically integrating the complementary strengths of linear-attention-based RWKV and self-attention-based Transformer models. Specifically, MRT partitions each image into fixed-size windows, utilizing RWKV modules to capture global dependencies across windows and Transformer blocks to model local redundancies within each window. The hierarchical attention mechanism enables more efficient and compact representation learning in the 1-D domain. To further enhance compression efficiency, we introduce a dedicated RWKV Compression Model (RCM) tailored to the structure characteristics of the intermediate 1-D latent features in MRT. Extensive experiments on standard image compression benchmarks validate the effectiveness of our approach. The proposed MRT framework consistently achieves superior reconstruction quality at bitrates below 0.02 bits per pixel (bpp). Quantitative results based on the DISTS metric show that MRT significantly outperforms the state-of-the-art 2-D architecture GLC, achieving bitrate savings of 43.75%, 30.59% on the Kodak and CLIC2020 test datasets, respectively.

Figures (14)

• Figure 1: Visualization of original images and heatmaps of latent representations. 2-D architecturesillm compress 65536 pixels into 256 latent features, while redundancy exists between these features. Our 1-D architectures generate more compact latent representations by producing 32 1-D latent features from 65536 pixels.
• Figure 2: Qualitative comparison of reconstruction quality across different compression architectures. Our proposed MRT demonstrates superior visual fidelity compared to VTM VVC and existing 2-D architectures including LALIC lalic, DiffEIC Diffeic and MS-ILLM illm. MRT achieves better preservation of fine details even at lower bitrates, while conventional 2-D architectures exhibit significant degradation despite operating at higher bitrates.
• Figure 3: (a) Overview of the proposed MRT, which consists of a MRT encoder, RCM, a MRT decoder and a pixel generator. MRT encoder applies a stack of 1-D transformation layers to extract compact 1-D latent representations, which are then compressed. The decoder mirrors this process to reconstruct the image. (b) The 1-D transform alternates between Bi-RWKV and windowed ViT blocks to capture both global and local dependencies. (c) The Bi-RWKV block comprises spatial and channel mix modules, each utilizing layer normalization, linear projections, and BiWKV attention for efficient sequence modeling.
• Figure 4: Effective receptive field visualization for different architectural configurations on kodim09. Columns: our proposed model with global Bi-RWKV blocks, model without cross-window dependency modeling, and model with global ViT blocks. Rows: gradient thresholds (0.0001 and 0.001).
• Figure 5: (a) RWKV compression model architecture. MLP$(c_1, c_2)\uparrow$ and MLP$(c_2, c_1)\downarrow$ denote dimensionality expansion and reduction, respectively. AE, AD, and Q represent arithmetic encoding, arithmetic decoding, and quantization. LFQ-E, LFQ-D stand for look-up free (LFQ) encoding and decoding. 1-D SCCTX denotes the 1-D Spatial-Channel Context Module. (b) Two-stage training strategy. Stage 1: auxiliary encoder aligns quantized codes with discrete targets via cross-entropy loss. Stage 2: RCM and MRT decoder are optimized end-to-end using pixel-level and perceptual losses.