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

TL;DR

This paper presents a novel self-supervised framework combining JEPA and DAAM to learn robust, hierarchical speech representations that enable efficient tokenization and high-quality waveform reconstruction, advancing neural audio coding.

Contribution

It introduces a density-adaptive attention mechanism into JEPA for adaptive feature selection and hierarchical speech structure discovery at low frame rates.

Findings

Tokens are highly compressed and reversible.

Model achieves competitive audio coding performance.

Hierarchical speech structure is effectively captured.

Abstract

We introduce a two-stage self-supervised framework that combines the Joint-Embedding Predictive Architecture (JEPA) with a Density Adaptive Attention Mechanism (DAAM) for learning robust speech representations. Stage~1 uses JEPA with DAAM to learn semantic audio features via masked prediction in latent space, fully decoupled from waveform reconstruction. Stage~2 leverages these representations for efficient tokenization using Finite Scalar Quantization (FSQ) and a mixed-radix packing scheme, followed by high-fidelity waveform reconstruction with a HiFi-GAN decoder. By integrating Gaussian mixture-based density-adaptive gating into the JEPA encoder, the model performs adaptive temporal feature selection and discovers hierarchical speech structure at a low frame rate of 2.5~Hz. The resulting tokens (47.5 tokens/sec) provide a reversible, highly compressed, and language-model-friendly representation that is competitive with, and often more efficient than, existing neural audio codecs.