3D-JEPA: A Joint Embedding Predictive Architecture for 3D Self-Supervised Representation Learning

TL;DR

3D-JEPA addresses the biases of invariance-based approaches and the inefficiency of generative 3D self-supervised methods by introducing a non-generative Joint Embedding Predictive Architecture for point clouds. It uses a multi-block sampling strategy to extract an informative context block and multiple target blocks from the same scene, paired with a context-aware decoder that injects context across decoder layers to learn semantic representations in embedding space. The training objective distills target-block representations from a teacher model via cosine similarity, avoiding pixel- or token-level reconstruction and enabling efficient pretraining. Across diverse downstream tasks, including real-world and synthetic classification and part segmentation, 3D-JEPA achieves strong accuracy with fewer pretraining epochs, demonstrating both effectiveness and efficiency.

Abstract

Invariance-based and generative methods have shown a conspicuous performance for 3D self-supervised representation learning (SSRL). However, the former relies on hand-crafted data augmentations that introduce bias not universally applicable to all downstream tasks, and the latter indiscriminately reconstructs masked regions, resulting in irrelevant details being saved in the representation space. To solve the problem above, we introduce 3D-JEPA, a novel non-generative 3D SSRL framework. Specifically, we propose a multi-block sampling strategy that produces a sufficiently informative context block and several representative target blocks. We present the context-aware decoder to enhance the reconstruction of the target blocks. Concretely, the context information is fed to the decoder continuously, facilitating the encoder in learning semantic modeling rather than memorizing the context information related to target blocks. Overall, 3D-JEPA predicts the representation of target blocks from a context block using the encoder and context-aware decoder architecture. Various downstream tasks on different datasets demonstrate 3D-JEPA's effectiveness and efficiency, achieving higher accuracy with fewer pretraining epochs, e.g., 88.65% accuracy on PB_T50_RS with 150 pretraining epochs.

Figures (7)

• Figure 1: OBJ_BG Fine-tuing Classification Accuracy. By predicting target blocks from a single context block without any data augmentations, the 3D-JEPA learns strong point cloud representation with less computing.
• Figure 2: Concept comparison of Invariance-based, Generative methods and our 3D-JEPA paradigms. (a) Invariance-based methods aim to generate similar embeddings ${{f}_{i}}$ and ${{f}_{j}}$ for compatible input pairs ${{x}_{i}}$ and ${{x}_{j}}$. (b) After generating the masked point cloud $x_{i}^{m}$ by random masking, Generative methods aim to output embeddings ${{\hat{x}}_{i}}$ to predict the original data ${{x}_{i}}$ as much as possible. (c) After generating the target block $x_{i}^{t}$ and context block $x_{i}^{c}$ by muti-block masking, 3D-JEPA aims to output the embeddings $\hat{f}_{i}^{c}$ to predict the embeddings $f_{i}^{t}$ of $x_{i}^{t}$.
• Figure 3: The Pipeline of 3D-JEPA. Given the input point cloud, the context block will be encoded as sequential tokens after the multi-block sampling. The context representation are then fed to every layer of the decoder to predict the representation of target blocks supervised by the outputs of the teacher model via cosine loss.
• Figure 4: Visualization of multi-block sampling. Given the point cloud, we sample 4 target blocks via FPS with a lower scale. Next, we randomly sample a context block with a larger scale and remove any overlapping target blocks. In this way, the target blocks have global semantic information, and the context block is informative.
• Figure 5: Context-aware Block. In the decoder block, we implement cross-attention layers after the self-attention.