音频Tokenization与离散表示

1. 概述

音频Tokenization是将连续音频信号转换为离散token序列的过程，是现代音频深度学习的基础。与文本tokenization类似，音频tokenization将高维连续的音频波形压缩为有限词汇表中的离散表示，使得自回归模型和语言建模技术能够应用于音频领域。

1.1 为什么需要音频Tokenization？

传统音频处理采用手工设计的特征（如MFCC、滤波器组），而深度学习时代需要可学习的端到端表示。音频Tokenization的核心动机：

统一多模态表示：将音频与其他模态（文本、图像）统一到离散token空间
语言建模扩展：利用LLM架构处理音频生成任务
压缩效率：离散表示比连续特征更紧凑
生成控制：离散token便于操控和编辑

1.2 音频信号基础

采样与量化

原始音频是连续的时间序列，通过采样和量化离散化：

import numpy as np
import librosa
 
# 加载音频
audio, sr = librosa.load("speech.wav", sr=16000)  # 16kHz采样
 
# 参数设置
duration = len(audio) / sr  # 音频时长（秒）
n_samples = len(audio)      # 样本点数量
bit_depth = 16               # 位深度（16-bit PCM）

时频表示

音频分析通常在频域进行：

表示方法	描述	适用场景
STFT	短时傅里叶变换	基础时频分析
Mel Spectrogram	梅尔刻度频谱	语音识别/合成
MFCC	梅尔频率倒谱系数	传统语音特征
CQT	常Q变换	音乐分析

# 计算Mel频谱
mel_spec = librosa.feature.melspectrogram(
    y=audio, sr=sr, n_mels=80, hop_length=160
)
log_mel = librosa.power_to_db(mel_spec)  # 转换为dB尺度

2. 矢量量化基础

2.1 VQ (Vector Quantization)

矢量量化是音频tokenization的核心技术。给定一个向量空间和码本 $C = {c_{1}, c_{2}, ..., c_{K}}$ ，VQ将输入向量映射到最近的码字：

q (z) = ar g c_{k} \in C min ∥ z - c_{k} ∥_{2}

VQ-VAE框架

class VectorQuantizer(nn.Module):
    def __init__(self, codebook_size, dim):
        super().__init__()
        self.codebook_size = codebook_size
        self.dim = dim
        # 码本: (codebook_size, dim)
        self.embedding = nn.Embedding(codebook_size, dim)
        self.embedding.weight.data.uniform_(-1/codebook_size, 1/codebook_size)
    
    def forward(self, z):
        # z: (B, T, D)
        z_flatten = z.view(-1, self.dim)  # (B*T, D)
        
        # 计算距离并找到最近码字
        d = torch.sum(z_flatten ** 2, dim=1, keepdim=True) + \
            torch.sum(self.embedding.weight ** 2, dim=1) - \
            2 * torch.matmul(z_flatten, self.embedding.weight.t())
        
        indices = torch.argmin(d, dim=1)  # (B*T,)
        quantized = self.embedding(indices).view_as(z)  # (B, T, D)
        
        return quantized, indices

2.2 训练策略

Straight-Through Estimator (STE)

由于argmin操作不可导，使用STE将梯度直接从输出复制到输入：

class VQLayer(nn.Module):
    def __init__(self, codebook_size, dim, commitment_cost=0.25):
        super().__init__()
        self.codebook = nn.Embedding(codebook_size, dim)
        self.commitment_cost = commitment_cost
    
    def forward(self, z):
        z_flatten = z.view(-1, self.dim)
        
        # 找到最近码字
        d = torch.cdist(z_flatten, self.codebook.weight)
        indices = torch.argmin(d, dim=1)
        
        # STE: 直接复制码字
        quantized = self.codebook(indices).view_as(z)
        
        # 损失: 重构损失 + 码本损失 + 承诺损失
        e_latent_loss = F.mse_loss(z.detach(), quantized)
        q_latent_loss = F.mse_loss(z, quantized.detach())
        commit_loss = self.commitment_cost * e_latent_loss
        
        loss = q_latent_loss + commit_loss
        
        # 保持z的值用于后续计算
        z_q = z + (quantized - z).detach()
        
        return z_q, loss, indices

3. 残差矢量量化 (RVQ)

3.1 RVQ原理

单层VQ受限于码本大小。残差矢量量化（RVQ）通过多层量化逐步细化表示：

z = q_{1} (z) + q_{2} (r_{1}) + q_{3} (r_{2}) + ... + q_{K} (r_{K - 1})

其中 $r_{i} = z - \sum_{j = 1}^{i} q_{j} (\cdot)$ 是第 $i$ 层的残差。

class ResidualVectorQuantizer(nn.Module):
    def __init__(self, codebook_size, n_quantizers, dim):
        super().__init__()
        self.n_quantizers = n_quantizers
        self.codebooks = nn.ModuleList([
            VectorQuantizer(codebook_size, dim) 
            for _ in range(n_quantizers)
        ])
    
    def forward(self, z, n_q=None):
        n_q = n_q or self.n_quantizers
        
        quantized = 0
        residual = z
        indices_list = []
        
        for i in range(n_q):
            # 量化残差
            q, _, indices = self.codebooks[i](residual)
            quantized = quantized + q
            residual = residual - q.detach()
            indices_list.append(indices)
        
        return quantized, indices_list

3.2 RVQ的优势

指数级码本大小： $K$ 层，每层 $N$ 个码字，总计 $N^{K}$ 表示能力
高效存储：每层独立更新
灵活比特率：可动态选择量化层数

# 示例: 8层，每层1024个码字
# 总表示能力: 1024^8 ≈ 2^80
rvq = ResidualVectorQuantizer(codebook_size=1024, n_quantizers=8, dim=256)

4. 音频编解码器架构

4.1 Encodec

Meta提出的端到端音频编解码器，使用RVQ量化：

架构设计

音频输入 → CNN Encoder → 量化器(RVQ) → CNN Decoder → 音频输出
                         ↓
                    离散token序列

class Encodec(nn.Module):
    def __init__(self, dim=128, n_res_block=4, codebook_size=1024, n_quantizers=8):
        super().__init__()
        
        # 卷积编码器
        self.encoder = nn.Sequential(
            nn.Conv1d(1, dim, 4, 2, 1),  # 下采样2倍
            *[
                ResidualBlock(dim) for _ in range(n_res_block)
            ],
            nn.Conv1d(dim, dim, 4, 2, 1),  # 下采样2倍
        )
        
        # RVQ量化器
        self.quantizer = ResidualVectorQuantizer(
            codebook_size=codebook_size,
            n_quantizers=n_quantizers,
            dim=dim
        )
        
        # 卷积解码器
        self.decoder = nn.Sequential(
            nn.ConvTranspose1d(dim, dim, 4, 2, 1),
            *[
                ResidualBlock(dim) for _ in range(n_res_block)
            ],
            nn.ConvTranspose1d(dim, 1, 4, 2, 1),
        )
    
    def forward(self, x):
        # x: (B, 1, T)
        h = self.encoder(x)
        z_q, indices = self.quantizer(h)
        y = self.decoder(z_q)
        return y, indices

关键特性

多尺度量化：不同层的量化层用于不同目的
带宽控制：通过选择量化层数控制比特率
音频质量：24kbps可达CD质量

4.2 SoundStream

字节跳动提出的音频编解码器，与Encodec类似但有改进：

class SoundStream(nn.Module):
    def __init__(self):
        super().__init__()
        
        # 编码器：更深的残差网络
        self.encoder = nn.Sequential(
            nn.Conv1d(1, 32, 7, 2, 3),
            *[_ResidualBlock(32, 64) for _ in range(3)],
            nn.Conv1d(64, 64, 7, 2, 3),
            *[_ResidualBlock(64, 128) for _ in range(3)],
            nn.Conv1d(128, 128, 7, 2, 3),
        )
        
        # 量化器：使用门控卷积
        self.quantizer = GatedResidualQuantizer(
            codebook_size=1024,
            n_quantizers=12,
            dim=128
        )
        
        # 解码器
        self.decoder = nn.Sequential(
            *[_ResidualBlock(128, 128) for _ in range(3)],
            nn.ConvTranspose1d(128, 64, 7, 2, 3),
            *[_ResidualBlock(64, 64) for _ in range(3)],
            nn.ConvTranspose1d(64, 32, 7, 2, 3),
            nn.ConvTranspose1d(32, 1, 7, 2, 3),
        )

5. 音频LM的Tokenization策略

5.1 离散音频Token

不同模型采用不同的tokenization策略：

模型	Token类型	码本大小	层数	比特率
AudioLM	Coarse + Fine	1024	3层	9kbps
VALL-E	Acoustic	1024	8层	-
MusicGen	Enco + Semantic	2048+1024	4层	3kbps
Encodec	RVQ	1024	8层	1.5-24kbps

5.2 多级Tokenization

MusicGen采用语义+声学双层tokenization：

class MusicGenTokenizer:
    """
    MusicGen的tokenization流程：
    1. 语义token：使用Encodec获取离散表示
    2. 声学token：进一步量化高频成分
    """
    
    def __init__(self):
        self.semantic_model = load_model("encodec_32khz")
        self.acoustic_rvq = ResidualVectorQuantizer(
            codebook_size=2048,
            n_quantizers=4
        )
    
    def tokenize(self, audio):
        # 语义token: 粗粒度表示
        semantic_tokens = self.semantic_model.encode(audio)
        
        # 声学token: 细粒度补充
        acoustic_features = self.acoustic_rvq(semantic_tokens)
        
        return {
            'semantic': semantic_tokens,
            'acoustic': acoustic_tokens
        }

6. 与文本Tokenization的对比

特性	文本Tokenization	音频Tokenization
输入类型	离散字符/子词	连续波形/频谱
词汇表大小	32K-128K (BPE)	1K-4K (VQ码本)
表示粒度	词/子词	帧/段落
信息密度	高压缩	低压缩
可逆性	通常不可逆	可通过声码器重建

关键区别

信息保留：音频token必须保留足够的声学信息以重建高质量音频
时序依赖：音频帧之间高度相关，需要考虑时序建模
多尺度表示：同时需要语义和声学信息

7. 实现与实践

7.1 使用HuggingFace音频tokenizer

from transformers import AutoProcessor, EncodecModel
 
# 加载预训练Encodec
model = EncodecModel.from_pretrained("facebook/encodec_32khz")
processor = AutoProcessor.from_pretrained("facebook/encodec_32khz")
 
# 编码
audio_input = processor(
    raw_audio=audio_array,
    sampling_rate=32000,
    return_tensors="pt"
)
encoder_outputs = model.encode(audio_input["input_values"])
 
# 获取离散token
codes = encoder_outputs.audio_codes  # (batch, n_quantizers, seq_len)

7.2 音频重建质量评估

def evaluate_audio_codec(original, reconstructed, sr=16000):
    """评估音频编解码器质量"""
    metrics = {}
    
    # 信噪比
    metrics['SNR'] = 10 * np.log10(
        np.sum(original**2) / np.sum((original - reconstructed)**2)
    )
    
    # 频谱差异
    orig_spec = librosa.stft(original)
    recon_spec = librosa.stft(reconstructed)
    metrics['SpectralDist'] = np.mean(np.abs(orig_spec - recon_spec))
    
    # 语音质量感知评估 (需要pesq库)
    # metrics['PESQ'] = pesq(original, reconstructed, sr)
    
    return metrics

Metaphor

探索

音频Tokenization与离散表示

音频Tokenization与离散表示

1. 概述

1.1 为什么需要音频Tokenization？

1.2 音频信号基础

采样与量化

时频表示

2. 矢量量化基础

2.1 VQ (Vector Quantization)

VQ-VAE框架

2.2 训练策略

Straight-Through Estimator (STE)

3. 残差矢量量化 (RVQ)

3.1 RVQ原理

3.2 RVQ的优势

4. 音频编解码器架构

4.1 Encodec

架构设计

关键特性

4.2 SoundStream

5. 音频LM的Tokenization策略

5.1 离散音频Token

5.2 多级Tokenization

6. 与文本Tokenization的对比

关键区别

7. 实现与实践

7.1 使用HuggingFace音频tokenizer

7.2 音频重建质量评估

8. 总结与展望

核心要点

未来方向

参考资料

关系图谱

目录

反向链接