RWKVQuant: Quantizing the RWKV Family with Proxy Guided Hybrid of Scalar and Vector Quantization

TL;DR

RWKVQuant addresses the challenge of post-training quantization for the RWKV family, where non-linear operators and broadly uniform weight distributions hinder traditional SQ or VQ methods. It introduces a proxy-guided hybrid SQ+VQ strategy, combining a coarse-grained information-entropy proxy with a fine-grained Taylor-moment proxy to decide per-weight quantization, and adds codebook optimization tailored to RWKV's element-wise multiplication. Empirical results show ~3-bit quantization with minimal accuracy loss and significant speed/memory benefits across language and vision tasks, with RWKV models exhibiting a higher propensity for scalar quantization than comparable LLaMA models. The work demonstrates practical deployment gains for RWKV on resource-constrained devices and provides a framework for further PTQ improvements in RWKV architectures.

Abstract

RWKV is a modern RNN architecture with comparable performance to Transformer, but still faces challenges when deployed to resource-constrained devices. Post Training Quantization (PTQ), which is a an essential technique to reduce model size and inference latency, has been widely used in Transformer models. However, it suffers significant degradation of performance when applied to RWKV. This paper investigates and identifies two key constraints inherent in the properties of RWKV: (1) Non-linear operators hinder the parameter-fusion of both smooth- and rotation-based quantization, introducing extra computation overhead. (2) The larger amount of uniformly distributed weights poses challenges for cluster-based quantization, leading to reduced accuracy. To this end, we propose RWKVQuant, a PTQ framework tailored for RWKV models, consisting of two novel techniques: (1) a coarse-to-fine proxy capable of adaptively selecting different quantization approaches by assessing the uniformity and identifying outliers in the weights, and (2) a codebook optimization algorithm that enhances the performance of cluster-based quantization methods for element-wise multiplication in RWKV. Experiments show that RWKVQuant can quantize RWKV-6-14B into about 3-bit with less than 1% accuracy loss and 2.14x speed up.

Figures (5)

• Figure 1: Accuracy-model size curve. Results of zero-shot accuracy are evaluated on the LAMBADA dataset radford2019language. Our proposed RWKVQuant outperforms the individual utilization of SQ and VQ methods for all sizes of models.
• Figure 2: Model Structure of RWKV-7. It contains several blocks and each block has two key modules: Time Mixing and Channel Mixing.
• Figure 3: Zero-shot accuracy when applying different quantization methods to specific weights. For the weight in each sub-figure, $\mathrm{SQ}_{\mathrm{acc}}$ denotes the performance when SQ is applied, $\mathrm{VQ}_{\mathrm{acc}}$ denotes the performance when VQ is applied, while all other weights are quantized with VQ. $P_c$ and $P_f$ are coarse- and fine-grained proxy, while $\mu_c$ and $\mu_f$ are their corresponding thresholds.
• Figure 4: Effectiveness of clipping for batch integration.
• Figure 5: Comparison of SQ proportion between RWKV and LLaMA Models.