ZipRWKV-Speech (Research Stage)

[English] | 中文版

ZipRWKV-Speech is a Speech-LLM based Automatic Speech Recognition (ASR) system. Its core architecture utilizes Zipformer as the speech encoder, integrated with RWKV7 as the language model backbone.

[!IMPORTANT]

Research Phase Disclaimer: This project is currently in the experimental/development stage and serves primarily as a personal implementation record. The codebase draws from several open-source projects (such as NeMo, K2, and RWKV), extracting and simplifying key components to create a lightweight yet high-performance SpeechLLM framework. Each module contains its own README to explain the implementation logic.

---

🏗️ System Architecture

• Speech Encoder: Zipformer (from K2/Icefall), providing efficient downsampling and feature extraction.

• LLM Backbone: RWKV7, combining the inference efficiency of RNNs with the training performance of Transformers.

• Data Pipeline: A dynamic bucketing system based on Lhotse.

🚀 Development Roadmap

Current progress on code organization and implementation:

• Data Pipeline (Lhotse-based)
  • Support for NeMo-style Manifest loading.
  • Implementation of DynamicBucketingSampler for dynamic Batch Size adjustment.
  • Integration of Cutset.mux for weighted multi-source data mixing.
  • On-the-fly data augmentation (Speed, Volume, Noise, SpecAugment).
  • test dataset code and conf.yaml.
  • noise manifest prepare script.
• Model Architecture
  • Zipformer Encoder integration and test code.
  • RWKV7 and PEFT (Parameter-Efficient Fine-Tuning) integration.
  • projector (MLP and rwkv)
• Training Implementation
  • PyTorch Lightning Training Module.
• Checkpoints & Evaluation
  • Release of pre-trained model weights.

K2 Installation Guide

Environment Requirements

• Python 3.10
• CUDA 12.8
• PyTorch 2.8

Installation Steps

1. Download K2 Wheel

Download the K2 wheel file that matches your environment configuration:

wget https://k2-fsa.github.io/k2/cuda.html
wget k2-1.24.4.dev20250807+cuda12.8.torch2.8.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

2. Fix the Downloaded Filename

Sometimes the + character in the filename may be converted to %2B after download. You need to rename it back to +:

# If the filename contains %2B, change it to +
# Original filename: k2-1.24.4.dev20250807%2Bcuda12.8.torch2.8.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
# New filename: k2-1.24.4.dev20250807+cuda12.8.torch2.8.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

3. Install K2

Install the wheel package using uv pip:

uv pip install "k2-1.24.4.dev20250807+cuda12.8.torch2.8.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl"

Zipformer Encoder

This project uses AudenAI/auden-encoder-tta-m10 as the encoder for Zipformer.

LLM-Based ASR: Audio-Text Embedding Fusion Concepts

Key Concepts

Concept 1: Audio Downsampling and Token-Level Alignment

Encoder Output to Tokens

After audio encoding (Fbank → Encoder), the encoder output features are processed through a projector that performs downsampling:

• Encoder output shape: (B, T, encoder_out_dim) where T = audio_seconds × 25
• After projector downsampling: (B, T//2, lim_dim)
• Resulting temporal resolution: ~12.5 Hz
  • Each token represents 80ms of audio
  • Typical Chinese character duration: 0.2~0.3 seconds
  • Tokens per character: 2~4 tokens

Concept 2: Prompt Design and LLM Compatibility

Model-Specific Prompt Formats

Different pre-trained LLM models require different prompt structures, which are already defined during their training:

• RWKV7: Simple format like User: and Assistant:
• Qwen: Special tokens such as <|im_start|> and <|im_end|>
• ASR-specific prompts: Common Chinese prompts include "請轉錄這段{lang}語音" (Please transcribe this {lang} audio)

Additional information: Some papers have also proposed a prompt projector module to address the impact of prompts. Reducing Prompt Sensitivity in LLM-based Speech Recognition Through Learnable Projection

Prompt Projector Alternative

Recent papers have proposed using a simpler approaches work well in practice:

• A straightforward special token like <|en|> (or <|zh|> for Chinese) is often sufficient
• This lightweight approach reduces computational overhead while maintaining effectiveness
• Language specification helps the model adapt its decoding strategy appropriately

Concept 3: Fusion and LLM Forward Pass

Embedding Concatenation Strategy

The key to the entire system is simple concatenation:

1. Downsampled audio features: Shape (B, T//2, lim_dim)
2. Text embeddings: Various prompt tokens and assistant markers
3. Concatenation: Place audio features immediately after the User: embedding
4. Combined input: [User_emb] + [Audio_tokens] + [Assistant_emb] + [Label_emb] with total shape (B, seq_len, lim_dim)

Training vs. Inference Difference

• Training:

  • Concatenate orange portion + green portion
  • Feed the entire sequence to the LLM for forward pass
  • Compute loss on the green portion to train the model

• Inference:

  • Only the orange portion is needed
  • The green portion is generated by the LLM's autoregressive decoding
  • No need to provide target text at inference time

Summary

The elegance of this approach lies in its simplicity:

1. Downsampling creates naturally-aligned tokens (~80ms each) that correspond well with phonetic units
2. Prompt design is straightforward and model-dependent; language indicators are sufficient
3. Fusion is simply concatenation of embeddings in a specific order
4. Training/Inference asymmetry leverages the LLM's inherent ability to generate text autoregressively

By combining downsampled audio features with carefully formatted text prompts, the system enables LLMs to perform effective end-to-end ASR without complex architectural modifications.

LLM-Based ASR: Projector Architecture Guide

Overview

In LLM-based ASR systems, the projector serves as a critical bridge between the audio encoder and the LLM. Its role is to downsampe and transform audio features into the embedding space of the LLM. This guide introduces common projector architectures and practical implementation insights.

Common Projector Architectures

When designing projectors for LLM-based ASR, three main architectures are commonly used, typically combined with downsampling mechanisms:

1. MLP (Multi-Layer Perceptron)

• Simple fully-connected layers
• Lightweight and computationally efficient
• Direct mapping from encoder output to LLM embedding space

2. Transformer Encoder

• Example: Fun-ASR implementation
• Adds sequential modeling capacity
• Can capture temporal dependencies in audio features
• Slightly higher computational cost than MLP

3. Q-Former

• Cross-attention mechanism for feature alignment
• Advanced feature interaction between audio and LLM embedding space
• More complex but potentially better feature alignment

Practical Implementation: RWKV-ASR Case Study

Architecture Details

Based on the RWKV-ASR implementation:

• Projector: 2-layer RWKV (simplified from original implementation)
  • Using 2-layer RWKV proved sufficient for effective feature projection
  • Maintains consistency with the base LLM architecture

Training Strategy: Two-Stage Approach

Stage 1: Frozen Encoder & LLM

• Freeze the audio encoder completely
• Freeze the LLM parameters completely
• Train only the projector
• Allows efficient feature space alignment without catastrophic forgetting

Critical Implementation Detail

Attention Mask Handling

A key correction to the original RWKV-ASR implementation is the addition of attention_mask during forward pass:

# Original issue: Missing attention_mask
# Impact: Causes hallucinations from padding tokens

# Fix: Include attention_mask parameter
output = projector(features, attention_mask=mask)

Hallucination Reduction Results:

• Before fix: ~20% hallucination rate
• After fix: 1-2% hallucination rate (on clean test set)
• Note: Hallucination increases on extremely noisy audio

The attention mask prevents the model from attending to padding positions, which was a major source of spurious token generation.

Emerging Approach: MoE-Based Projector

Motivation

Recent papers have introduced Mixture of Experts (MoE) mechanisms into projectors:

• Language-Specific Alignment: Different languages have different acoustic-to-semantic alignment patterns
• Specialized Experts: Assign different MLP experts for different languages
• Routing Mechanism: Dynamically select appropriate expert based on language identifier

Advantages

• Better handling of language-specific phonetic-phonemic mappings
• Improved cross-lingual ASR performance
• Flexible scaling for new languages without retraining the entire projector

Implementation Considerations

• Requires explicit language labels during training or let the router choose and adjust via balance loss.
• Add language-specific routing tokens (e.g., <|en|>, <|zh|>)
• Can be combined with standard downsampling strategies

Summary

The projector in LLM-based ASR is a crucial yet often overlooked component:

1. Architecture Choice: MLP, Transformer Encoder, or Q-Former, each with trade-offs in complexity and performance
2. Training Efficiency: Freezing encoder and LLM allows rapid iteration on projector design
3. Implementation Details Matter: Proper attention mask handling can dramatically reduce hallucinations
4. Language Diversity: MoE-based projectors offer a promising direction for multilingual ASR

By carefully designing the projector and addressing implementation details, systems can achieve robust audio-to-LLM embedding alignment while maintaining computational efficiency.

🙏 Acknowledgements

We borrowed a lot of code from the following excellent projects:

• rwkv7
• rwkv_asr
• speechllm
• Auden
• Zipformer_Lightning
• icefall
• RWKV-PEFT