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

1.1 Overview

• The two-stage framework pairs JEPA and DAAM for speech representations, separating masked latent prediction from later FSQ tokenization and HiFi-GAN reconstruction.The authors motivate the separation as addressing semantic representation learning and perceptual reconstruction as distinct goals.Source: §1.1 · p. 2, §1.2 · p. 3

1.2 Motivation: Why JEPA for Speech?

• DAAM learns to select informative temporal regions and features, while JEPA separates semantic representation learning from waveform reconstruction.The stated benefits include reduced Stage 2 training cost through encoder fine-tuning and reuse of the encoder for downstream tasks.Source: §1.2 · p. 3

2.1 JEPA Masking Strategy

• Block masking creates contiguous masked spans rather than random individual positions, forcing the model to learn longer-range temporal dependencies and semantic content.Minimum spans are 2 frames and maximum spans are T/4 frames, adaptive to sequence length; at 2.5 Hz, spans vary with sequence length.Source: §2.1 · p. 4

2.2 Density Adaptive Attention for Temporal Feature Modulation

• DAAM uses learned Gaussian-mixture distributions to identify statistically salient temporal regions, unlike standard self-attention’s pairwise comparisons between positions.It operates along time, computing a gate from a learned one-channel temporal projection that scales the full feature tensor.Source: §2.2 · p. 4, §2.2.1 · pp. 4–5

2.2.1 Mathematical Formulation

• The DAAM formulation aggregates Gaussian-component log densities into a gate that modulates features through element-wise scaling.The implementation uses four Gaussian components and positive scales computed with softplus.Source: §2.2.1 · p. 5

2.3 JEPA Encoder Architecture

• The JEPA encoder uses online and EMA target pathways, with strided convolutions producing a 512-channel latent representation at approximately 2.5 Hz.The architecture also uses eight Conformer layers with 16 attention heads; DAAM gating is applied in the encoder blocks, not after the Conformer blocks in the described implementation.Source: §2.3 · p. 6, §2.3.1 · p. 6, Figure 2 · p. 7

2.4 JEPA Predictor Network

• The predictor processes visible-context features through two Conformer blocks and predicts features at all temporal positions, including masked regions.Its architecture expands channels, refines features between the Conformer blocks, then projects back to the original dimensionality.Source: §2.4 · p. 6, Figure 3 · p. 8

2.5 Stage 1 Training Objective

• The JEPA objective is pure self-supervised prediction in latent space, with squared-difference loss computed only on masked regions and normalized by the number of masked tokens times channels.After each training step, the target encoder parameters are updated via EMA with momentum coefficient τ = 0.996.Source: §2.5 · p. 6, §2.5.1 · p. 7, §2.5.2 · p. 7

3 Stage 2: Fine-Tuning Encoder + FSQ Quantization + HiFi-GAN Decoder

• Stage 2 fine-tunes the JEPA encoder and introduces quantization and waveform reconstruction, using FSQ with fixed scalar quantization rather than a learned codebook.The decoder uses HiFi-GAN to map quantized representations back to waveform.Source: §3 · p. 8, §3.1 · p. 8, §3.3 · p. 11

3.1 Finite Scalar Quantization (FSQ)

• FSQ is used to create discrete tokens without codebook learning, while the mixed-radix scheme groups indices for compact integer representation.The supplied FSQ formulation passages identify four quantization levels per dimension and a code dimension of 128.Source: §3.1 · p. 8, §3.1.1 · p. 8

3.1.1 FSQ Formulation

• With identical architectures except for DAAM gating, JEPA+DAAM reaches a lower final masked-prediction MSE loss than JEPA without DAAM, about 0.09 versus 0.17.The figure caption also reports faster convergence for JEPA+DAAM.Source: Figure 4 · p. 9

3.2 Mixed-Radix Token Packing

• Mixed-radix packing combines groups of FSQ indices into compact integer tokens using dimension-specific radices and Horner’s method.The reported implementation uses group size G = 7 and a per-token vocabulary of 16,384.Source: §3.2 · p. 9, §3.2.2 · p. 9, §3.2.5 · p. 10

3.2.5 Comparison to Alternatives

• The packing scheme is described as perfectly reversible via modular arithmetic, near-optimal for its radices, and implementable with integer-only operations.Flexible group size trades vocabulary size against token rate; unlike VQ-VAE, the method uses no learned codebook.Source: §3.2.5 · p. 10
• For speech generation, the compact tokens enable direct training of decoder-only Transformers, establishing their stated use as a training representation in that setting.Source: §3.2.6 · p. 11

3.4 Stage 2 Training Objective

• The Stage 2 objective optimizes the FSQ quantizer, HiFi-GAN decoder, and JEPA encoder, and includes an L1 reconstruction loss.The supplied passages also identify multi-resolution STFT and GAN losses as training components.Source: §3.4 · p. 11, §3.4.1 · p. 12

3.3.1 Decoder Architecture

• The HiFi-GAN decoder upsamples quantized features through five ConvTranspose1D blocks with multi-receptive-field residual blocks.Its residual kernels span sizes 3, 7, 11, 15, 23, and 32 to capture multi-scale temporal patterns.Source: §3.3.1 · p. 11, Figure 5 · p. 12

2. Multi-resolution STFT loss [Yamamoto et al., 2020].

• The listed multi-resolution STFT settings use FFT sizes 2048, 1024, 512, 256, and 128 with hop sizes 512, 256, 128, 64, and 32.The passages identify spectral convergence as part of the loss description.Source: §2 · p. 12
• The Stage 2 loss setup includes multi-period and multi-scale discriminators for its GAN loss.The generator loss is described alongside these discriminators.Source: §2 · p. 12
• The discriminator is frozen for a 5,000-step warmup and then updated every step.The listed training configuration uses batch size 8 and weight decay 10^-3.Source: §2 · p. 13

4.1 Dataset

• The dataset is LibriLight, a large-scale unlabeled English speech corpus, with about 9,000 training hours combined across both stages.Validation uses held-out speakers.Source: §4.1 · p. 13

4.2 Data Preprocessing

• Audio is resampled to 24 kHz if needed and converted to mono by averaging channels, with no further preprocessing.Normalization is handled in-model.Source: §4.2 · p. 13

4.3 Distributed Training

• Training uses two NVIDIA A100 GPUs with mixed precision and global batch sizes of 64 for Stage 1 and 16 for Stage 2.Gradient accumulation is one step.Source: §4.3 · p. 14

4.4 Inference Pipeline

• 47.5 tokens/sec is the reported inference token rate with G = 7 packing.The supplied passages give no further inference-pipeline detail.Source: §4.4 · p. 14

5 Model Architecture and Efficiency

• The training description characterizes the setup as self-supervised pretraining followed by supervised fine-tuning.The supplied passage does not specify additional efficiency figures.Source: §5.2 · p. 14
• The inference model has 191M parameters and does not use an EMA network.The passage labels this as inference efficiency.Source: §5.2 · p. 14
• Table 4 is identified as covering training efficiency for the two stages, but its caption provides no numeric comparisons.No table values are supplied in the passage.Source: Table 4 · p. 15

6 Evaluation Metrics

• The evaluation is qualitative, with the work presenting preliminary findings after variants were trained under limited computational budgets.The supplied passage does not report quantitative evaluation results.Source: §6 · p. 15
• The JEPA baseline is an encoder without DAAM gating.Source: §6 · p. 15
• The comparison includes pre-trained WavLM-Large and the JEPA encoder with DAAM gating.Source: §6 · p. 15

7.1 Why DAAM Improves JEPA Representations

• The discussion frames DAAM integration into JEPA as an approach with several advantages.Source: §7.1 · p. 15
• Unlike standard softmax attention, which computes pairwise correlations between positions, DAAM uses Gaussian mixture modeling to identify statistically salient timesteps.Source: §7.1 · p. 15
• DAAM operates on temporal statistics and can capture salient patterns without the quadratic complexity of full self-attention.Source: §7.1 · p. 15

8 Limitations and Future Work

• Fixed block masking may not adapt optimally to varying speech rates or linguistic structure; adaptive masking around acoustic or linguistic boundaries is proposed for future work.Source: §8 · p. 15
• Experiments are limited to English using LibriLight, leaving generalization to tonal and morphologically rich languages open.Source: §8 · p. 15
• Pretraining uses relatively modest data compared with large-scale SSL systems, restricting conclusions to emerging capabilities.Source: §8 · p. 15
• Extending JEPA to audio–visual or audio–text joint embedding prediction is identified as a promising direction for multimodal representations.Source: §8 · p. 15

9 Code Availability

• The complete implementation includes training scripts, model architectures, and data-processing pipelines.Source: §9 · p. 16
• The listed code covers Stage 1 JEPA encoder training with DAAM and Stage 2 decoder training with the encoder.Source: §9 · p. 16
• The implementation also includes FSQ and mixed-radix packing, a HiFi-GAN decoder with optional DAAM gating, and DeepSpeed integration.Source: §9 · p. 16

• Stage 1 trains a DAAM-gated JEPA encoder with masked-region MSE, while Stage 2 uses the representations for reconstruction with L1, multi-resolution STFT, and adversarial GAN losses.Stage 2 also uses FSQ and HiFi-GAN.Source: §10 · p. 16
• Gaussian mixture-based DAAM gating provides adaptive feature selection during self-supervised learning.Source: §10 · p. 16
• 47.5 tokens/sec tokens from mixed-radix FSQ packing are reversible and substantially lower in rate than many existing neural audio codecs.The conclusion describes the scheme as efficient tokenization.Source: §10 · p. 16
• The training paradigm separates pure self-supervised pretraining from reconstruction-focused fine-tuning.Source: §10 · p. 16
• The authors report that probabilistic attention dynamically identifies acoustically salient regions during masked prediction, and that JEPA can serve as a tokenizer for large language models and other sequence models.Source: §10 · p. 16