PaperlayerView exact revision on arXiv

Source-linked AI summary

A-JEPA: Joint-Embedding Predictive Architecture Can Listen

Zhengcong Fei, Mingyuan Fan, Junshi Huang

arXiv:2311.15830v3cs.SDcs.CVeess.AS

TL;DR

Masked modeling has driven progress in foundational vision models, but applying JEPA-style learning to audio remains an open endeavor. A-JEPA predicts masked audio-spectrogram regions in latent space using curriculum masking and regularized fine-tuning. It achieves state-of-the-art results across audio and speech classification tasks, including a 1.3-point mAP gain on AudioSet-2M over contemporary masking and reconstruction competitors.

  • Problem

    Although JEPA learns image representations by predicting masked regions in latent space, applying this learning principle to sensory data such as audio remains an open endeavor.

    Source: §1 · p. 1

  • Method

    A-JEPA uses a context encoder and curriculum masking to predict masked spectrogram-region representations in latent space, then fine-tunes with regularized patch masking.

    Source: §1 · p. 2

  • Results

    A-JEPA achieves state-of-the-art performance across multiple audio and speech classification tasks and exceeds contemporary masking and reconstruction competitors by +1.3 points in mAP on AudioSet-2M.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Larger pre-training audio datasets continue to improve performance even under fixed computational constraints.

    Source: §1 · p. 2

Abstract

from arXiv · show

This paper presents that the masked-modeling principle driving the success of large foundational vision models can be effectively applied to audio by making predictions in a latent space. We introduce Audio-based Joint-Embedding Predictive Architecture (A-JEPA), a simple extension method for self-supervised learning from the audio spectrum. Following the design of I-JEPA, our A-JEPA encodes visible audio spectrogram patches with a curriculum masking strategy via context encoder, and predicts the representations of regions sampled at well-designed locations. The target representations of those regions are extracted by the exponential moving average of context encoder, \emph{i.e.}, target encoder, on the whole spectrogram. We find it beneficial to transfer random block masking into time-frequency aware masking in a curriculum manner, considering the complexity of highly correlated in local time and frequency in audio spectrograms. To enhance contextual semantic understanding and robustness, we fine-tune the encoder with a regularized masking on target datasets, instead of input dropping or zero. Empirically, when built with Vision Transformers structure, we find A-JEPA to be highly scalable and sets new state-of-the-art performance on multiple audio and speech classification tasks, outperforming other recent models that use externally supervised pre-training.

1 Introduction

A-JEPA extends joint-embedding prediction to audio spectrograms, using curriculum masking and latent-space targets. It reports state-of-the-art results across multiple audio and speech classification tasks, including a 1.3-point mAP gain on AudioSet-2M over contemporary masking and reconstruction competitors.

  • JEPA predicts masked image regions in a high-level representation space, unlike masked-image methods that predict pixels or low-level visual tokens; applying this principle to audio remains an open endeavor.The paper motivates latent-space prediction as a way to focus on semantic features rather than unnecessary pixel-level details.Source: §1 · p. 1
  • A-JEPA uses curriculum masking that progresses from random blocks to time-frequency-aware patterns, then fine-tunes with regularized patch masking rather than dropping patches or setting them to zero.The context encoder predicts masked-region representations in latent space using guidance from an exponential-moving-average target encoder.Source: Figure 1 · p. 1, §1 · p. 2
  • A-JEPA achieves state-of-the-art performance across multiple audio and speech classification tasks and exceeds contemporary masking and reconstruction competitors by +1.3 points in mAP on AudioSet-2M.The reported comparison is for a self-supervised audio-only model.Source: §1 · p. 2
  • The paper presents A-JEPA as a unified, scalable framework for self-supervised audio representations and reports continued performance gains as pre-training audio data grows under fixed computational constraints.Source: §1 · p. 2

2 Related Works

Prior work combines joint-embedding and reconstruction objectives, while I-JEPA predicts target-block representations from a context block. Audio representation learning spans different input signals and objectives; A-JEPA extends latent-space prediction to spectrograms.

  • I-JEPA [Assran et al., 2023] predicts representations of target blocks from one context block; A-JEPA extends the JEPA framework to audio spectrograms.Source: §2 · p. 2
  • Joint-embedding methods have been combined with patch-level reconstruction losses, while view-invariance methods can favor global image representations; local loss terms have been proposed to improve other vision tasks.The passage also relates local objectives to contrastive predictive coding.Source: §2 · p. 2
  • Audio-only self-supervised methods vary by input signal and objective, including waveform, frame-level-feature, and spectrogram-patch inputs with contrastive or predictive and reconstructive learning.A-JEPA focuses on audio-only self-supervised pre-training from scratch and predicts in latent space rather than at a low level.Source: §2 · pp. 2–3

3 Approach

A-JEPA predicts target audio-spectrogram representations from visible context in latent space. It combines curriculum masking during pre-training with regularized masking during fine-tuning.

3.2 Application to Audio

  • A-JEPA encodes visible spectrogram patches and predicts target-region representations in latent space, using curriculum masking for pre-training and regularized masking for fine-tuning.The encoder processes non-masked patches, while the target encoder tracks the context encoder through an exponential moving average; the objective averages L2 distances between predicted and target features.Source: §3.2 · pp. 3–4
  • Audio masking moves from random blocks toward time-frequency-aware masks because block masking can leave nearby audio context that makes prediction easier.The schedule randomly selects between the two strategies at each training step, gradually favoring time-frequency-aware masking; Figure 2 illustrates both strategies.Source: §3.2 · p. 3, Figure 2 · p. 3
  • Regularized masking excludes selected patch tokens from attention-weight calculations during fine-tuning, requiring representations to use information from remaining neighbors.The vanilla attention calculation is retained at testing.Source: §3.2 · p. 4, Figure 4 · p. 4

4 Experiments

Experiments evaluate A-JEPA on audio and speech classification tasks and analyze benchmark performance, model and masking choices, training scale, and predictor representations. The results show strong comparisons with prior systems alongside benefits from larger models, curriculum masking, and more pre-training data.

4.1 Experimental settings.

  • Evaluation covers AudioSet (AS-2M and AS-20K), ESC-50, Speech Commands (SPC-1 and SPC-2), and VoxCeleb speaker identification.For AS-20K, AudioSet is used for pre-training and the balanced 20K set for fine-tuning.Source: §4 · p. 4, §4.1 · p. 5

4.2 Comparison with the State-of-the-art

  • 38.4 mAP on AudioSet-20K; A-JEPA leads in-domain self-supervised models across tasks and exceeds prior ImageNet-supervised systems.The comparison reports A-JEPA as highest among in-domain self-supervised models on all tasks.Source: §4.2 · pp. 5–6, Table 1 · p. 5

4.3 Model Analysis

  • Larger encoders tend to perform better but require more computation and memory, with ViT-L's advantage particularly pronounced on AS-20K.The reported ViT-S–ViT-B performance gap can shrink from 4.7 to 2.5 mAP when fine-tuning uses more in-domain data.Source: Table 2 · p. 5, §4.3 · p. 6
  • Curriculum masking outperforms random block masking in pre-training, while regularized patch masking at lower ratios improves downstream fine-tuning performance.The inverse masking order produces a larger performance drop; task-agnostic pre-training favors easy-to-hard curriculum masking with a high ratio.Source: §4.3 · p. 6, Table 3 · p. 6, Table 4 · p. 6
  • A 16-layer predictor outperforms shallower variants, while a 512-dimensional predictor balances compute and performance without significant gains from greater width.The width comparison concerns predictor embedding dimension.Source: §4.3 · p. 6, Table 5 · p. 6, Table 6 · p. 6
  • Performance consistently improves as pre-training data increases, while longer training helps only until performance plateaus after epoch 24.The data-size analysis randomly samples subsets from the original dataset.Source: §4.3 · p. 6, Figure 6 · p. 7

4.4 Predictor Visualizations

  • Across generated samples conditioned on positional mask tokens, shared characteristics reflect information in the predictor representation, which captures positional uncertainty and high-level audio components.The visualization decodes predictor outputs from a frozen pre-trained encoder and predictor.Source: §4.4 · pp. 6–7, Figure 5 · p. 7

5 Conclusion

The conclusion reports that A-JEPA achieves superior performance across multiple audio and speech classification tasks. It highlights curriculum masking and regularized fine-tuning, and identifies multimodal self-supervised learning as future work.

5 Conclusion

  • A-JEPA achieves superior performance on multiple audio and speech classification tasks.The authors describe this as a favorable outcome from extending JEPA to audio spectrograms.Source: §5 · p. 7
  • Longer pre-training benefits downstream task performance on the ViT-B AudioSet ablation.Source: Figure 7 · p. 7
  • Time-frequency aware masking that progresses from easy to hard can enhance learned audio representations.The conclusion contrasts this strategy with a straightforward application of JEPA.Source: §5 · p. 7
  • Regularized masking during fine-tuning, rather than directly dropping inputs or setting them to zero, contributes to accuracy improvement.Source: §5 · p. 7
  • The authors identify multimodal self-supervised learning with joint latent embeddings, such as video and text, as future work.Source: §5 · p. 7

A Implementation Details

The implementation uses a vanilla ViT-B encoder and Transformer decoder, with audio converted to non-overlapping spectrogram patches. Pre-training and fine-tuning use specified masking, sampling, and training schedules.

A Implementation Details

  • A vanilla 12-layer ViT-B serves as the encoder, with a 16-layer vanilla Transformer decoder.Source: A Implementation Details · p. 12
  • Audio is converted to 128 Mel-frequency bands, then divided into non-overlapping patches using 16×16 convolutional kernels with matching strides.For a 10-second AudioSet recording, the spectrogram has dimensions 1×1024×128.Source: A Implementation Details · p. 12
  • Pre-training uses AudioSet-2M for 24 epochs, while fine-tuning applies a default 10% regularized patch-masking ratio.AudioSet-2M supervised fine-tuning uses weighted sampling; AudioSet-20K fine-tuning runs for 60 epochs without it.Source: A Implementation Details · p. 12

B Pytorch Code for Time-frequency Aware Masking

The implementation samples block masks, optionally restricts them to acceptable regions, and uses complement masks to form time-frequency aware masking regions.

B Pytorch Code for Time-frequency Aware Masking

Loading 2311.15830v3…