JEPA as a Neural Tokenizer: Learning Robust Speech Representations with Density Adaptive Attention

Paper: arXiv:2512.07168 Code: gioannides/Density-Adaptive-JEPA Code reference: main @ ac9d866a (2025-12-07)

1. Motivation (研究动机)

这篇论文要解决的问题不是“再做一个音频 codec”，而是把 speech representation 学习和 waveform reconstruction 从同一个目标里拆开。传统 neural audio codec 往往让 encoder 直接服务于重建损失：只要能让波形、频谱或感知质量更像原音频，encoder 就会被奖励。但这种训练把两个目标混在一起：一方面我们希望 representation 记住语言内容、说话人、韵律和较长时间结构；另一方面 decoder 又需要非常细的局部相位、谐波和瞬态信息来恢复 waveform。论文认为，这种耦合会让 encoder 过度关注低层声学细节，而不是形成更适合 sequence modeling / language modeling 的语义型 speech token。

作者引入 JEPA 的动机正是“只在 latent space 预测，不直接重建 waveform”。Stage 1 中，online encoder 从完整 waveform 得到 latent feature，feature-space mask 把部分时间位置替换成 learned mask token，predictor 只根据 visible context 去预测 target encoder 的 masked latent；target encoder 由 EMA 更新并 stop-gradient。这样 Stage 1 的目标是 masked representation prediction，而不是让 decoder 立刻还原每个采样点。Stage 2 才把学到的表示拿去做 FSQ tokenization 和 HiFi-GAN reconstruction。这个分工使得 encoder 可以先学习鲁棒、低帧率、可被语言模型消费的 speech representation，再由 decoder 解决“把 token 变回声音”的问题。

研究价值在于 token rate 和 representation 语义性的组合。论文希望得到一种 reversible、compressed、LM-friendly 的 speech token：Stage 1 的 progressive downsampling 把 frame rate 降到 $2.5\mathrm{Hz}$，Stage 2 的 mixed-radix packing 把 128 个 FSQ dimensions 按 $G=7$ 打包为 19 groups/frame，最终是 $47.5$ tokens/sec。相比 SoundStream / EnCodec 这类 $75\mathrm{Hz}$ frame-rate codec，作者试图证明 JEPA-style latent prediction 可以成为一种 speech neural tokenizer：既能为 decoder 保留足够信息，又不会像 waveform codec 那样产生过密 token 序列，便于后续 decoder-only Transformer 做 next-token prediction over speech tokens。

这篇论文也把 DAAM（Density Adaptive Attention Mechanism）放到 JEPA encoder 里，解决 speech 序列中“哪些时间片真的有信息”的问题。普通 self-attention 问的是“这个 timestep 与哪些 timestep 相似”；DAAM 问的是“这个 timestep 的统计分布是否 unusual / informative”。它用 Gaussian mixture 对时间维统计量建模，再用 density-derived gate 去调制 feature。对于 speech，音素边界、能量变化、共振峰迁移和节奏变化往往是局部但语义上重要的事件；DAAM 的设计目标就是让 encoder 在 masked prediction 时更关注这些统计 salience，而不是均匀处理所有时间位置。

2. Idea (核心思想)

核心 insight：把 JEPA 看成“speech tokenizer 的语义预训练器”——先用 latent masked prediction 学会低帧率、语义化的连续表示，再用 FSQ + mixed-radix packing 把连续表示变成可逆离散 token，最后用 HiFi-GAN 负责从 token 回到 waveform。真正新颖的地方不是单独使用 JEPA、FSQ 或 HiFi-GAN，而是把它们按“representation first, reconstruction second”的顺序连接，并用 DAAM 在 JEPA encoder 内做 density-adaptive temporal feature selection。

与 VQ-VAE / SoundStream / EnCodec 的关键差异是目标分配。VQ-style codec 通常学习 codebook 并直接用重建/感知损失驱动 encoder-decoder；本文 Stage 1 的 encoder 不学习 waveform reconstruction，而是预测 EMA target latent，因此它更像 I-JEPA/BYOL 的 self-supervised representation learner。与 WavLM-Large 这类通用 speech SSL baseline 相比，本文额外强调“tokenizer”属性：Stage 2 必须给出可逆、低 token rate、可直接喂给语言模型的 discrete sequence，而不仅是 downstream embedding。

DAAM 的创新点可以用一句话概括：它不显式计算 $T\times T$ 的 pairwise attention，而是沿 temporal axis 估计每个 feature 的 Gaussian mixture density，再把 density gate 乘回 feature。这样模型以较低复杂度获得“统计显著性”信号；与普通 attention 的相似度检索不同，它更像一种 adaptive salience detector。

3. Method (方法)

3.1 Overall framework：两阶段 neural tokenizer

Figure 1 解读：这张图是 Stage 1 的核心框架。绿色 online encoder 接收完整 waveform，随后在 feature space 里用 learned mask token 替换 masked timesteps；黄色 predictor 根据 masked context 预测所有位置的 latent；紫色 target encoder 也接收完整 waveform，但只通过 EMA 更新，不直接接收梯度。MSE loss 只在 mask 集合 $\mathcal{M}$ 上计算，目标是 $\mathrm{sg}({\mathbf{z}}_{\text{target}})$。因此，模型不会被 waveform reconstruction 牵着走，而是被迫在 latent space 中补全被遮蔽的 speech structure。

整体流程分两段：

1. Stage 1 / JEPA+DAAM pretraining：输入 waveform $\mathbf{x}$，online encoder 得到 ${\mathbf{z}}_{\text{context}}$；block mask 产生 $\mathbf{m}\in \{0,1{\}}^{B\times T}$；masked context 进入 predictor；target encoder 由 EMA 提供稳定 target。训练目标是 masked MSE。
2. Stage 2 / FSQ + HiFi-GAN reconstruction：Stage 1 的 encoder 输出 ${\mathbf{z}}_e$，FSQ 将连续 feature quantize 成 ${\mathbf{z}}_q$ 和 integer indices；mixed-radix packing 把多个 indices 合成更少的 discrete tokens；HiFi-GAN decoder 把 dequantized features 上采样回 waveform。

直觉上，Stage 1 解决“token 应该表示什么”，Stage 2 解决“token 能否复原声音”。如果把两者合在一个端到端 reconstruction loss 里，encoder 可能优先记住短时 waveform 细节；如果只有 Stage 1 而没有 Stage 2，表示不一定可逆，也无法成为 audio codec / speech LM 的 tokenizer。因此两阶段结构是本文的主要设计取舍。

3.2 JEPA masking strategy：在 feature space 做 block prediction

论文使用 temporal block masking，而不是随机遮蔽单点。对 batch 内每个样本，先初始化 $\mathbf{m}=1$，目标 mask 数量为 $n_{\text{mask}}=\lfloor \rho T\rfloor$；然后不断采样 span length $\ell \sim \mathrm{Uniform}(s_{\min },s_{\max })$ 和 start position，把连续区间 $[t_{\text{start}},t_{\text{end}})$ 设为 0，直到达到目标 mask 数量。论文超参数是 $\rho =0.5$，$s_{\min }=2$ frames，$s_{\max }=T/4$ frames；released code 在 training loop 中调用 create_jepa_mask(..., min_span=2, max_span=max(4, T_z // 4))，即小序列时至少允许 4-frame span。

这种 masking 的意义是让 predictor 不能只依赖邻近局部插值，而要用较长 context 补全被遮蔽的语音片段。由于 encoder 的总 stride 是 $9600$ samples，在 $24\mathrm{kHz}$ 下 latent frame rate 约 $2.5\mathrm{Hz}$；一个 masked span 对应的原始音频时间已经相当长，因此 JEPA loss 更偏向 syllable / phrase-level 的结构，而不是 sample-level 的重建。

def make_jepa_mask(batch_size, seq_len, mask_ratio=0.5, min_span=2, max_span=None, device="cuda"):
    max_span = max(4, seq_len // 4) if max_span is None else max_span
    mask = torch.ones(batch_size, seq_len, device=device)
    for b in range(batch_size):
        target = int(seq_len * mask_ratio)
        masked = 0
        while masked < target:
            span_len = random.randint(min_span, max_span)
            start = random.randint(0, max(1, seq_len - span_len))
            end = min(start + span_len, seq_len)
            mask[b, start:end] = 0
            masked += end - start
    return mask

3.3 DAAM：Gaussian mixture density gate for temporal feature modulation

DAAM 模块输入 $\mathbf{x}\in {\mathbb{R}}^{B\times C\times T}$，沿时间维为每个 batch/channel 估计统计量： $\mu =\frac{1}{T}{\sum }_{t=1}^T{x}_{:,:,t},{\sigma }^2=\frac{1}{T}{\sum }_{t=1}^T(x_{:,:,t}-\mu {)}^2.$ 对 $K$ 个 Gaussian components，维护 learnable mean offsets $\mathbf{\delta }=[{\delta }_1,\dots ,{\delta }_K]$ 和 log-scale parameters $\mathbf{\nu }=[{\nu }_1,\dots ,{\nu }_K]$。正尺度用 softplus 保证： ${\tilde{\sigma }}_k=\mathrm{softplus}({\nu }_k)+ϵ=\log (1+\exp ({\nu }_k))+ϵ,$ 其中论文写 $ϵ=10^{-3}$。每个 timestep 在第 $k$ 个 component 下的标准化偏差和 log-density 为： $z_{k,t}=\frac{x_{:,:,t}-(\mu +{\delta }_k)}{\sigma \cdot {\tilde{\sigma }}_k+ϵ},$ $\log p_k(x_t)=-\frac{1}{2}{z}_{k,t}^2-\log {\tilde{\sigma }}_k-\frac{1}{2}\log (2\pi ).$ 混合 density gate 用 log-sum-exp 聚合： $\log \mathbf{G}(x_t)=\mathrm{logsumexp}(\{\log p_1(x_t),\dots ,\log p_K(x_t)\})-\log K,$ $\mathbf{G}(x_t)=\exp (\log \mathbf{G}(x_t)),{\mathbf{y}}_t={\mathbf{x}}_t\odot \mathbf{G}(x_t).$ released code 的 GaussianAdaptiveAttention 与论文公式基本一致：mean_offsets 初始化为 0，log_sigma 初始化为 $\log (0.5)$，variance clamp 到 $10^{-6}$，num_gaussians=4；GAttnGateG 先用 $1\times 1$ convolution 投影成单通道 attention signal，再用 GaussianAdaptiveAttention(..., return_attention_details=True) 得到 gate，最后执行 $\mathbf{y}=\mathbf{x}\cdot (1+\alpha \cdot \text{gate})$，其中 $\alpha$ 初始化为 $0.05$。

class GaussianDensityGate(nn.Module):
    def __init__(self, num_gaussians=4):
        super().__init__()
        self.mean_offsets = nn.Parameter(torch.zeros(num_gaussians))
        self.log_sigma = nn.Parameter(torch.full((num_gaussians,), math.log(0.5)))
 
    def forward(self, x):
        xf = x.float()                         # [B, 1, T] or [B, C, T]
        mean = xf.mean(dim=2, keepdim=True)
        std = xf.var(dim=2, keepdim=True, unbiased=False).clamp_min(1e-6).sqrt()
        sigma = F.softplus(self.log_sigma) + 1e-3
        log_terms = []
        for k in range(len(sigma)):
            z = (xf - (mean + self.mean_offsets[k])) / (std * sigma[k] + 1e-8)
            log_terms.append(-0.5 * z.square() - torch.log(sigma[k]) - 0.5 * math.log(2 * math.pi))
        log_gate = torch.logsumexp(torch.stack(log_terms, dim=-1), dim=-1) - math.log(len(sigma))
        return torch.exp(log_gate).to(x.dtype)
 
class GAttnGateG(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.to_attn = nn.Conv1d(channels, 1, 1)
        self.gaa = GaussianDensityGate(num_gaussians=4)
        self.alpha = nn.Parameter(torch.tensor(0.05))
 
    def forward(self, x):
        gate = self.gaa(self.to_attn(x))
        return x * (1.0 + self.alpha * gate), gate

3.4 JEPA encoder and predictor architecture

Figure 2 解读：online encoder 先用 Conv1D/downsampling path 把 waveform 压到低帧率 latent，再通过 8 个 Conformer blocks 建模较长上下文。每个 encoder block 包含 strided Conv1D、SnakeBeta activation、residual blocks 和 Gaussian Adaptive Attention gating；target encoder 共享同样结构，但通过 EMA 更新而不是 backpropagation。论文给出的 channel progression 是 $64\to 128\to 256\to 384\to 512\to 512$，stride product 是 $8\times 8\times 5\times 5\times 6=9600$。

encoder 的输出形状是 $[B,512,T_z]$，其中 $T_z$ 约等于原始 waveform 长度除以 9600。Conformer 部分使用 8 layers、16 heads；每层包含 feed-forward、multi-head self-attention、depthwise convolution 和第二个 feed-forward。released code 的 CLI defaults 与论文结构一致：--channels 64,128,256,384,512,512、--strides 8,8,5,5,6、--n_res_blocks 8、--n_conformer 8、--heads 16。

Figure 3 解读：predictor 接收的是 ${\mathbf{z}}_{\text{masked}}$，而不是 raw waveform。它先用 expansion Conv1D 把 channel 维加宽，再经过两个 Conformer blocks 和中间 Conv1D refine features，最后 projection Conv1D 回到原始维度，输出所有 timesteps 的 ${\mathbf{z}}_{\text{pred}}$。loss 只看 masked positions，因此 visible positions 的 prediction 只是中间计算，不直接贡献 JEPA objective。

Stage 1 的 masked prediction loss 是： ${\mathcal{L}}_{\text{JEPA}}=\frac{1}{N_{\text{mask}}\cdot C}{\sum }_{t\in \mathcal{M}}{\Vert {\mathbf{z}}_{\text{pred}}^{(t)}-\mathrm{sg}({\mathbf{z}}_{\text{target}}^{(t)})\Vert }^2.$ EMA target update 为： ${\mathbf{\theta }}_{\text{target}}\leftarrow \tau {\mathbf{\theta }}_{\text{target}}+(1-\tau ){\mathbf{\theta }}_{\text{online}},$ 其中 $\tau =0.996$。released code 的 training loop 在每步 engine.step() 后调用 engine.module.update_target_encoder()；它还监控 z_pred.std(dim=[0, 2])，当 mean standard deviation 小于 $0.01$ 时打印 collapse warning，但该监控不反传。

def jepa_train_step(engine, wav, args):
    with torch.no_grad():
        t_z = jepa_time_len_from_wav(wav.shape[-1], args.strides)
        mask = create_jepa_mask(
            batch_size=wav.shape[0], seq_len=t_z,
            mask_ratio=args.mask_ratio, min_span=2, max_span=max(4, t_z // 4),
            device=wav.device,
        )
    z_context, z_pred, mask, z_target = engine.module(wav, mask)
    weight = (1 - mask).unsqueeze(1).to(z_pred.dtype)
    diff2 = (z_pred - z_target.detach()).square()
    loss = (diff2 * weight).sum() / (weight.sum() * z_pred.shape[1]).clamp_min(1e-8)
    engine.backward(loss)
    engine.step()
    engine.module.update_target_encoder()
    mean_std = z_pred.detach().std(dim=[0, 2]).mean()
    return {"loss": loss, "collapse_std": mean_std}

3.5 FSQ quantization and mixed-radix token packing

Stage 2 中，encoder feature ${\mathbf{z}}_e\in {\mathbb{R}}^{B\times C\times T}$ 先经过 $\tanh$： ${\mathbf{z}}_e^{\prime }=\tanh ({\mathbf{z}}_e).$ 对第 $d$ 个 scalar dimension，FSQ 定义固定 quantization levels： $B_d=\left\{\frac{2i-L_d+1}{L_d}:i\in \{0,1,\dots ,L_d-1\}\right\},$ 并选择最近边界： $q_d(x)=\arg {\min }_{b\in B_d}|x-b|.$ 反向传播使用 straight-through estimator： $\frac{\partial \mathcal{L}}{\partial {\mathbf{z}}_e}=\frac{\partial \mathcal{L}}{\partial {\mathbf{z}}_q}.$ 论文/代码配置为 levels=[4,4,4,4]、code_dim=128、temperature=1.0。注意：论文文字说每个 dimension radix 4，$D=128$；代码实现 _fsq_dim_radices(D, levels) 会把 4 个 levels 各复制 $D/4=32$ 次，得到 128 个 radix。FSQ 本身没有 learnable codebook，也没有 commitment loss；released code 的 aux_loss 明确是 0。

Mixed-radix packing 把每 $G=7$ 个 indices 合成一个 token： $\text{token}={\sum }_{k=1}^G{i}_k{\prod }_{j=k+1}^G{r}_j.$ Horner-style iterative computation 写作： $\text{token}=i_1\cdot r_2\cdots r_G+\cdots +i_{G-1}{r}_G+i_G,$ 并在代码中 right-to-left 执行： $\text{token}=i_k+\text{token}\cdot r_k.$ 因为 $D=128$，$G=7$，所以 group 数是 $\lceil 128/7\rceil =19$，需要 padding $19\times 7-128=5$ 个 radix-1 dimensions。frame rate 是 $f=\frac{\text{sample\_rate}}{\text{hop}}=\frac{24000}{9600}=2.5\mathrm{Hz},$ 最终 token rate： $\text{tps}=2.5\times 19=\mathrm{47.5.}$

def fsq_quantize_and_pack(z_e, levels=(4, 4, 4, 4), group_size=7):
    z = torch.tanh(z_e).permute(0, 2, 1).contiguous()  # [B, T, D]
    bounds = [torch.linspace(-(L - 1) / L, (L - 1) / L, L, device=z.device) for L in levels]
    chunks, indices = [], []
    dims_per_level = z.shape[-1] // len(levels)
    for i, b in enumerate(bounds):
        s, e = i * dims_per_level, (i + 1) * dims_per_level
        dist = (z[..., s:e].unsqueeze(-1) - b.view(1, 1, 1, -1)).abs()
        idx = torch.argmin(dist, dim=-1)
        chunks.append(b[idx])
        indices.append(idx)
    z_q = torch.cat(chunks, dim=-1).permute(0, 2, 1).contiguous()
    z_q = z_e + (z_q - z_e).detach()
 
    idx = torch.cat(indices, dim=-1)          # [B, T, 128]
    radices = torch.tensor([4] * 128, device=idx.device)
    pad = math.ceil(idx.shape[-1] / group_size) * group_size - idx.shape[-1]
    idx = torch.cat([idx, torch.zeros(*idx.shape[:2], pad, dtype=idx.dtype, device=idx.device)], dim=-1)
    radices = torch.cat([radices, torch.ones(pad, dtype=radices.dtype, device=idx.device)])
    tokens = []
    for g in range(idx.shape[-1] // group_size):
        tok = torch.zeros(idx.shape[0], idx.shape[1], dtype=torch.long, device=idx.device)
        for k in range(group_size - 1, -1, -1):
            tok = idx[:, :, g * group_size + k] + tok * radices[g * group_size + k]
        tokens.append(tok)
    return z_q, torch.stack(tokens, dim=-1)

3.6 HiFi-GAN decoder and Stage 2 losses

Figure 5 解读：Stage 2 decoder 将 quantized features ${\mathbf{z}}_q$ 先通过 bottleneck Conv1D，再经过 5 个 ConvTranspose1D upsampling blocks，channel progression 是 $512\to 384\to 256\to 128\to 64$，stride 是 $6,5,5,8,8$，总 stride 同样是 9600。每个 block 里有 MRF residual blocks，kernel sizes 包括 3、7、11、15、23、32，用来覆盖不同时间尺度；SnakeBeta activation 为周期性 audio signal 提供 inductive bias。

Stage 2 的总损失为： ${\mathcal{L}}_{\text{total}}={\mathcal{L}}_{\text{rec}}+{\lambda }_{\text{stft}}{\mathcal{L}}_{\text{stft}}+{\lambda }_{\text{gan}}{\mathcal{L}}_{\text{gan}}.$ 其中 waveform L1 reconstruction loss： ${\mathcal{L}}_{\text{rec}}=\frac{1}{T_{\text{wav}}}{\sum }_{t=1}^{T_{\text{wav}}}|{\hat{x}}_t-x_t|.$ Multi-resolution STFT loss： ${\mathcal{L}}_{\text{stft}}={\sum }_{m=1}^M\left({\mathcal{L}}_{\text{sc}}^{(m)}+{\mathcal{L}}_{\text{mag}}^{(m)}\right),$ ${\mathcal{L}}_{\text{sc}}^{(m)}=\frac{{\Vert |S_m(\hat{x})|-|S_m(x)|\Vert }_F}{{\Vert |S_m(x)|\Vert }_F},$ ${\mathcal{L}}_{\text{mag}}^{(m)}=\frac{1}{N_m}{\Vert \log |S_m(\hat{x})|-\log |S_m(x)|\Vert }_1.$ GAN 部分使用 MPD/MSD discriminators： ${\mathcal{L}}_{\text{gen}}={\sum }_{d\in \{\text{MPD},\text{MSD}\}}\mathbb{E}[(D_d(\hat{x})-1{)}^2],$ ${\mathcal{L}}_{\text{feat}}={\sum }_{d\in \{\text{MPD},\text{MSD}\}}{\sum }_{l=1}^{L_d}\frac{1}{N_l}{\Vert D_d^{(l)}(x)-D_d^{(l)}(\hat{x})\Vert }_1,$ ${\mathcal{L}}_{\text{gan}}={\mathcal{L}}_{\text{gen}}+{\mathcal{L}}_{\text{feat}},$ ${\mathcal{L}}_{\text{disc}}={\sum }_{d\in \{\text{MPD},\text{MSD}\}}\left(\mathbb{E}[(D_d(x)-1{)}^2]+\mathbb{E}[D_d(\hat{x}{)}^2]\right).$ 论文给出的 loss weights 是 ${\lambda }_{\text{stft}}=2.0$、${\lambda }_{\text{gan}}=0.1$；discriminator warmup 为 5000 steps，之后每步更新。released code 的 train_decoder_with_frozen_encoder 中 generator loss 正是 rec_loss + spectral_weight * stft_loss + gan_weight * gen_loss，默认 spectral_weight=2.0、gan_weight=0.1、disc_start_step=5000、disc_interval=1；discriminator loss 在 global_step > disc_start_step 后更新。

def decoder_train_step(engine_g, mpd, msd, mr_stft, wav, args, global_step, opt_d):
    rec, indices, aux_loss, z_e = engine_g(wav)
    rec_loss = F.l1_loss(rec, wav)
    stft_loss = mr_stft(rec, wav)
    y_df_r, y_df_g, f_df_r, f_df_g = mpd(wav, rec)
    y_ds_r, y_ds_g, f_ds_r, f_ds_g = msd(wav, rec)
    gen_loss = (generator_loss(y_df_g) + generator_loss(y_ds_g)
                + feature_loss(f_df_r, f_df_g) + feature_loss(f_ds_r, f_ds_g))
    total_g = rec_loss + args.spectral_weight * stft_loss + args.gan_weight * gen_loss
    engine_g.backward(total_g)
    engine_g.step()
 
    if global_step > args.disc_start_step and global_step % args.disc_interval == 0:
        opt_d.zero_grad()
        with torch.no_grad():
            rec_det = rec.detach()
        y_df_r, y_df_g, _, _ = mpd(wav, rec_det)
        y_ds_r, y_ds_g, _, _ = msd(wav, rec_det)
        dloss = discriminator_loss(y_df_r, y_df_g) + discriminator_loss(y_ds_r, y_ds_g)
        dloss.backward()
        opt_d.step()
    return total_g

3.7 论文表述与 released code 实现差异

论文公式与 released code 实现差异：论文 §3 标题写的是 “Fine-Tuning Encoder + FSQ Quantization + HiFi-GAN Decoder”，并在 Stage 2 objective 中说 Stage 2 optimizes FSQ quantizer、HiFi-GAN decoder 和 JEPA encoder；但 released code train_fsqvae_jepa.py 在 Stage 2 明确把 encoder parameter group 设为 lr=0.0，日志也打印 “JEPA Encoder (frozen via lr=0)”。因此，若按 commit ac9d866a 复现，Stage 2 实际是加载 JEPA encoder 后冻结 encoder，只训练 FSQ + decoder/discriminators，而不是论文文字意义上的 encoder fine-tuning。

论文表述与 released code 实现差异：论文正文在 encoder 架构段落中先说 DAAM gating applied in encoder blocks and no DAAM after Conformer blocks，又紧接着写 “After each Conformer block, features pass through GAttnGateG modules”。released code 支持前一种解释：EncoderBlock 在 strided Conv/residual 后使用 GAttnGateG；ConformerBlock 本身没有 DAAM gate。笔记中的方法描述按 released code 和 Figure 2 caption 处理，即 DAAM 位于 encoder downsampling blocks。

代码/文档可复现性差异：GitHub repo 只有 README.md、ds_ckpt_to_pt.py、train_fsqvae_jepa.py 三个文件；README 的命令却使用 train_jepa_fsqvae_hifigan.py 和 ds_config.json，这两个文件在 commit ac9d866a 仓库根目录中不存在。README launch template 还写 Stage 1 --max_steps 200000、Stage 2 --max_steps 800000，而论文 Table 4 报告的是 24K/29K。笔记在实验设置中同时列出论文报告值与 released code launch/defaults，避免把 README template 当作论文实际训练配置。

3.8 Code-to-paper mapping

Code reference: main @ ac9d866a (2025-12-07) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
JEPA block mask	train_fsqvae_jepa.py	create_jepa_mask; training loop calls it with mask_ratio, min_span=2, max_span=max(4,T_z//4)
DAAM / Gaussian density gate	train_fsqvae_jepa.py	GaussianAdaptiveAttention, GAttnGateG
Encoder downsampling + Conformer	train_fsqvae_jepa.py	EncoderBlock, ConformerBlock, JEPAEncoder.encode
EMA target encoder	train_fsqvae_jepa.py	JEPAEncoder._target_encode, JEPAEncoder.update_target_encoder
JEPA masked prediction loss	train_fsqvae_jepa.py	train_jepa_encoder loop: diff2, masked denominator, collapse std logging
FSQ quantization	train_fsqvae_jepa.py	FiniteScalarQuantizer.quantize, FiniteScalarQuantizer.forward
Mixed-radix token packing	train_fsqvae_jepa.py	_fsq_dim_radices, fsq_pack_indices, fsq_token_stats_from_indices
HiFi-GAN decoder	train_fsqvae_jepa.py	HiFiDecoderBlock, MRFBlock, ResBlock, WaveformJEPAFSQVAE.decode
Stage 2 reconstruction/GAN training	train_fsqvae_jepa.py	train_decoder_with_frozen_encoder, MRSTFTLoss, MultiPeriodDiscriminator, MultiScaleDiscriminator
DeepSpeed checkpoint conversion	ds_ckpt_to_pt.py	consolidated checkpoint export utility

4. Experimental Setup (实验设置)

Dataset and preprocessing

论文使用 LibriLight，这是大规模 unlabeled English speech corpus。训练 split 报告为约 9000 hours，覆盖两个阶段；validation 使用 held-out speakers。论文没有给出 utterance/sample count，因此可复现记录应写为“约 9000 小时，sample count 未详细说明”。音频采样率统一为 24 kHz，最大音频长度 15 s；预处理包括 resample to 24 kHz、转 mono（多声道平均），除此之外不做额外 preprocessing，normalization 由模型内部处理。

Baselines and comparisons

论文在 Evaluation Metrics 段明确列出的 representation baselines 是：

Baseline / Variant	作用
JEPA baseline	同样的 JEPA encoder，但没有 DAAM gating，用于 Figure 4 的 loss curve 对比
WavLM-Large	预训练 speech SSL model，用作外部参考 baseline；论文没有报告完整数值表
JEPA+DAAM	本文方法，JEPA encoder with density-adaptive gating

此外，Table 2 用于 token/frame-rate efficiency 对比的 neural codecs 包括 U-Codec、Mimi、DualCodec、SoundStream、EnCodec 和 DAC。这里比较的是 frame rate / token density，而不是同一 benchmark 上的 MOS、WER 或 PESQ。

Evaluation metrics

论文强调这是 limited compute budget 下的 preliminary finding，因此没有完整主结果 benchmark 表。实际使用/报告的指标包括：

Metric	定义与用途
Stage 1 masked prediction MSE	Figure 4 中 JEPA+DAAM 与 no-DAAM baseline 的训练损失；只在 masked latent positions 上计算
Token rate / frame rate	衡量 tokenizer 对 speech LM 是否友好；本文核心数值是 $2.5\mathrm{Hz}$ frame rate 和 $47.5$ tokens/sec
Parameter count	Table 3 报告 Stage 1、Stage 2、inference model 参数规模
Training efficiency	Table 4 报告 trainable params、training steps、batch size、learning rate
Qualitative evaluation	论文文字称 report qualitative evaluations；没有给出完整可复算的 WER/PESQ/STOI/SI-SDR 数值表

Training config：论文报告值 vs released code launch/defaults

论文 Table 4 报告：Stage 1 trainable parameters 121.7M (50.7%)、training steps 24K、batch size 32、learning rate $1.5\times 10^{-4}$；Stage 2 trainable parameters 69.3M (22.4%)、training steps 29K、batch size 8、learning rate $1.5\times 10^{-4}$。论文 §3 还给出 Stage 1 optimizer 为 AdamW，${\beta }_1=0.8$、${\beta }_2=0.99$、weight decay $10^{-3}$、max audio 15 s @ 24 kHz；Stage 2 decoder/discriminator 使用 AdamW，decoder LR $1.5\times 10^{-4}$，discriminator LR $0.75\times 10^{-4}$，weight decay $10^{-3}$，${\lambda }_{\text{stft}}=2.0$，${\lambda }_{\text{gan}}=0.1$，discriminator warmup 5000 steps。

released code/config evidence 必须分开看：README.md 的 Stage 1 launch command 使用 --batch_size 32 --max_steps 200000 --lr 1.5e-4 --mask_ratio 0.5；Stage 2 launch command 使用 --batch_size 8 --max_steps 800000 --disc_start_step 5000 --spectral_weight 2.0 --gan_weight 0.1 --lr 1.5e-4。

train_fsqvae_jepa.py argparse defaults 是 sample_rate=24000、fsq_levels=4,4,4,4、code_dim=128、channels=64,128,256,384,512,512、strides=8,8,5,5,6、n_res_blocks=8、n_conformer=8、heads=16、batch_size=4、max_steps=800000、lr=1.5e-4、max_seconds=15.0、spectral_weight=2.0、gan_weight=0.1、disc_start_step=5000。

由于 repo 没有 paper-specific config 文件，不能把 argparse defaults 或 README template 当成论文实际 24K/29K 运行；这里只能记录二者差异。

硬件方面，论文 §4.3 报告 2× NVIDIA A100 80GB，mixed precision 为 FP16 forward/backward、critical ops FP32，gradient accumulation 1 step，global batch size Stage 1 为 64、Stage 2 为 16。inference pipeline 是 waveform $\to$ JEPA encoder $\to$ latent features $\to$ FSQ quantization $\to$ tokens $\to$ dequantization $\to$ quantized features $\to$ HiFi-GAN waveform。

5. Experimental Results (实验结果)

5.1 Stage 1 loss and DAAM ablation

Figure 4 解读：蓝线 JEPA+DAAM 的 masked prediction MSE 收敛更快，最终约 $0.09$；橙线 JEPA without DAAM 最终约 $0.17$。两者架构相同，差异是 DAAM gating。这个结果支撑了论文的核心论点：DAAM 不只是额外参数，而是帮助 JEPA encoder 在 self-supervised masked prediction 中更有效地找到 acoustically salient temporal regions。需要注意的是，Figure 4 是训练 loss curve，不等价于 downstream ASR / MOS / reconstruction quality benchmark。

5.2 Tokenization efficiency

Approach	Tokens/sec	Reversible	Notes
No packing (128 dims)	320	Yes	Each FSQ dim is a token
Mixed-radix (ours, $G=7$)	47.5	Yes	Pack 7 dims/token
VQ codebook	Variable	Yes	Requires learned codebook

这个表说明 mixed-radix packing 是 token rate 的关键：如果 128 个 FSQ dimensions 每个都作为 token，token rate 会是 $2.5\times 128=320$ tokens/sec；按 $G=7$ 打包后只有 19 groups/frame，所以降为 $47.5$ tokens/sec。该压缩仍然可逆，因为 mixed-radix token 可以用 modular arithmetic 还原成各维 FSQ indices。

Model	Frame Rate	Notes
Ours (JEPA+FSQ)	2.5 Hz	Mixed-radix packing (19 groups/frame)
U-Codec	5 Hz	Ultra-low for LLM-TTS
Mimi	12.5 Hz	Semantic distillation
DualCodec	12.5—25 Hz	Dual-stream architecture
SoundStream (24 kHz)	75 Hz	13.3 ms frames
EnCodec (24 kHz)	75 Hz	75 steps/sec @ 24 kHz
DAC (44.1 kHz)	86 Hz	Stride 512 @ 44.1 kHz

Table 2 的结论是本文 frame rate 比常见 neural codec 低很多，尤其比 SoundStream/EnCodec 的 $75\mathrm{Hz}$ 和 DAC 的 $86\mathrm{Hz}$ 更适合 language model 长序列训练。但这张表只比较 temporal density；它没有证明同等 bitrate 下的重建质量或语义识别性能更好。

5.3 Parameter efficiency and training efficiency

Component	Parameters	Notes
Online encoder	121.7M	Trainable
Target encoder (EMA)	118.5M	Momentum update
Predictor network	3.2M	Trainable
Stage 1 total	240.2M	121.7M trainable
JEPA encoder	240.2M	Paper says fine-tuned; released code freezes via lr=0
FSQ quantizer	~0.01M	Trainable
HiFi-GAN decoder	69.2M	Trainable
Stage 2 total	309.5M	69.3M trainable
Encoder only	121.7M	Online encoder only
FSQ + decoder	69.3M	Inference component
Inference total	191.0M	Single-pass model

Stage 1 训练时同时存在 online encoder、target encoder 和 predictor，但真正通过 gradient 更新的是 online encoder + predictor；target encoder 是 EMA state。inference 时不需要 target encoder 和 predictor，只保留 online encoder、FSQ 和 decoder，因此 inference total 是 191.0M，而不是 Stage 1/Stage 2 训练图里的全部参数。

Metric	Stage 1 (JEPA)	Stage 2 (decoder)
Trainable parameters	121.7M (50.7%)	69.3M (22.4%)
Training steps	24K	29K
Batch size	32	8
Learning rate	$1.5\times 10^{-4}$	$1.5\times 10^{-4}$

论文用这张表强调两阶段训练的 compute split：Stage 1 的表示学习主要训练 encoder；Stage 2 的 reconstruction 学习主要训练 decoder/FSQ。结合 released code 的差异看，如果按公开代码运行，Stage 2 encoder 实际 frozen via lr=0，因此这张表里的 69.3M trainable 与代码更一致；但论文标题“Fine-Tuning Encoder”容易误导复现者以为 encoder 在 Stage 2 继续更新。

5.4 Main findings, limitations, and caveats

主要实验结论可以概括为三点。第一，DAAM 的 ablation 在 Stage 1 loss 上有效：JEPA+DAAM 从训练曲线看收敛更快、最终 MSE 更低，说明 Gaussian mixture salience gate 对 masked latent prediction 有帮助。第二，mixed-radix FSQ packing 把 token rate 从不打包的 320 tokens/sec 降到 47.5 tokens/sec，同时保持可逆；这直接服务于 speech LM 的序列长度约束。第三，模型在 inference 只需要 191.0M 参数，不需要 EMA target encoder 和 predictor，因此训练结构复杂但部署路径相对简单。

作者明确列出的限制包括：固定 block masking strategy 可能不能适应不同语速或 linguistic boundary，未来需要 adaptive masking；实验目前限制在 English LibriLight，tonal languages 和 morphologically rich languages 的泛化仍未验证；pretraining 数据规模相对大型 SSL 系统仍较小，因此结论属于 emerging capabilities；cross-modal JEPA（audio-visual 或 audio-text joint embedding prediction）是未来方向。

还需要额外注意两个阅读/复现 caveats。第一，论文没有给出完整 downstream benchmark 数值（例如 WER、speaker verification、PESQ、STOI、SI-SDR 或 MOS），所以“competitive with codecs”目前主要由 token/frame-rate、loss curve、参数效率和定性评估支撑，而不是完整 head-to-head quality table。第二，公开仓库很小且命令文档与文件名/训练步数存在不一致；如果要复现论文数值，需要作者补充具体 launch config、数据列表和 DeepSpeed config，否则只能复现方法骨架而不能严格复现 24K/29K 训练记录。