Point-JEPA: A Joint Embedding Predictive Architecture for Self-Supervised Learning on Point Cloud
Ayumu Saito Jiju Poovvancheri

[
    Cédric Pilatte
    Accepted XXX, Received YYY; in original form ZZZ
====================================================







Recent advancements in self-supervised learning in the point cloud domain have demonstrated significant potential.
However, these methods often suffer from drawbacks, including lengthy pre-training time, the necessity of reconstruction in the input space, or the necessity of additional modalities.
In order to address these issues, we introduce Point-JEPA, a joint embedding predictive architecture designed specifically for point cloud data.
To this end, we introduce a sequencer that orders point cloud tokens to efficiently compute and utilize tokens' proximity based on their indices during target and context selection. 
The sequencer also allows shared computations of the tokens' proximity between context and target selection, further improving the efficiency.
Experimentally, our method achieves competitive results with state-of-the-art methods while avoiding the reconstruction in the input space or additional modality.

    


§ INTRODUCTION


Sef-supervised learning is a paradigm that allows the model to learn a meaningful representation from unlabeled data.
This allows the model to utilize a vast amount of unlabeled data and learn a strong representation, leading to advancements in natural language processing <cit.> and 2D computer vision <cit.>.

Only recently, we have seen the successful applications of self-supervised learning in the point cloud domain <cit.>, achieving state-of-the-art results on downstream tasks.
However,  our initial investigation found that they require a significant amount of pre-training time as shown in <Ref>.
This slow pre-training process can pose constraints in scaling to a larger dataset or deeper and more complex models, hindering the key advantage of sel-supervised learning; its capacity to learn a strong representation from a vast amount of data.

The recently proposed Joint-Embedding Predictive Architecture <cit.> and successful implementations of the architecture for pre-training a model <cit.> show JEPA's ability to learn strong off-the-shelf semantic representations.
The idea behind JEPA is to learn a representation by predicting the embedding of the input signal, called target, from another compatible input signal, called context, with the help of a predictor network.
This allows learning in the representation space instead of the input space, leading to efficient learning.

















Inspired by I-JEPA <cit.>, we aim to apply Joint-Embedding Predictive Architecture in the point cloud domain.
While this application introduces a promising direction for self-supervised learning in the point cloud domain, it also presents a unique challenge.
Unlike an image, a point cloud is a set of points with no specific order.
Regardless of the permutations applied to the set, the set presents represents the same object.
This implies that the model consuming N points needs to be invariant to N! permutations of the order of the input data <cit.>. 
This unordered nature of the point cloud data makes the context and target selection of the data difficult especially if we aim to select spatially contiguous patches similar to I-JEPA <cit.>.
In order to overcome this challenge while utilizing the full potential of Joint-Embedding Predictive Architecture for efficiency, we introduce Point-JEPA.
To this end, we introduce a sequencer that orders the input sequence such that elements keep the spatial proximity when they are adjacent in the data.
This allows the shared computation for proximity between the target and context, allowing efficiency while achieving a state of the performance with linear evaluation on ModelNet40 <cit.> as shown in <Ref>.
Our method also achieves promising results in other downstream tasks showing the strong and transferable learned representation.












































































































































































































































































































































































§ POINT-JEPA ARCHITECTURE

Our aim is to bring Joint-Embedding Predictive Architecture to the point cloud domain while measuring the efficient implementation.
The overall framework, as shown in <Ref>, first converts the point cloud to a set of tokens, then the sequencer orders the tokens based on the spatial proximity, and Joint-Embedding Predictive Architecture is applied to the ordered tokens.
We utilize a small PointNet <cit.> architecture for encoding the grouped points and standard Transformer <cit.> architecture for the context and target encoder as well as the predictor.
It is important to note that our JEPA architecture operates on the token instead of patches in order to share the point encoder network between context and target encoder for efficiency similar to Point2Vec <cit.>.  



  
Point Cloud Tokenizer.
Similar to previous studies utilizing standard Transformer architecture on point cloud objects <cit.>, we adopt a point cloud tokenizer that embeds groups of points.
Given a point cloud object, c center points are first sampled using the farthest point sampling <cit.>. 
Then using the k-nearest neighbors algorithm, we sample k closest points to each of c center points. 
These point patches are then normalized by subtracting the center point coordinate from the coordinates of the points in the patches.
This allows the separation between local structural information and the positional information of the patches. 
In order to embed the local point patches, we utilize a small PointNet <cit.> architecture.
This ensures that the patch embedding, or token, remains invariant to any permutations of data feeding order of points within the patch.
Specifically, this PointNet contains two sets of a shared MLP  and a max-pooling layer. 
First, a shared MLP  maps each point into a feature vector. Then, we apply max-pooling to these vectors and concatenate the result back to the original feature vector.
Another shared MLP then processes these concatenated vectors and a max-pooling layer is applied to produce tokens.




  
Sequencer.
Because of the previously observed advantages of spatial proximity in patches for targets and context, referred to as a block in I-JEPA <cit.>, we aim to sample tokens that are spatially close to each other.
As previously mentioned, point cloud data is permutation invariant to data feeding order, which implies that even if the indices of tokens are sequential, they might not be spatially adjacent.
Moreover, our method involves selecting spatially contiguous M blocks of encoded embedding vectors as the target while ensuring that the context does not include the tokens corresponding to these embedding vectors (details in the next paragraph).
In order to overcome these challenges, we introduce a sequencer that is applied after
the tokenization of the points.
This sequencer orders tokens based on their associated
center points. 
The process begins with a chosen center point and its associated token.
In each subsequent step, the center point closest to the one previously chosen and
its associated token are selected.
This is iterated until the sequencer visits all of the
center points.
The resulting arrangement of tokens is in a sequence where contiguous elements are also spatially contiguous in most cases.
This allows the shared computation of spatial proximity between context and target selection.
At the same time, this also allows simpler implementation for context and target selection.
It is worth noting, however, that selecting two adjacent token indices does not always guarantee spatial proximity; there might be a gap between them. While this is true, the experiment results show that this iterative ordering is effective enough in our JEPA architecture.





  
Context and Target.
Targets in Point-JEPA can be considered patch-level representations of the point cloud object, which the predictor aims to predict.
As illustrated in <Ref>, the target encoder initially encodes the tokens conventionally, and we randomly select M possibly overlapping target blocks, which are sets of adjacent embedding vectors.
We denote these sets as target blocks.
Context, on the other hand, is the representation of the point cloud object which is passed to the predictor to facilitate the reconstruction of target blocks. 
In order to avoid trivial learning, we remove the corresponding tokens of the target blocks during the context selection step.
Out of this subset of tokens, we select a set of tokens that are adjacent in their indices and denote it as a context block.
Because of this, context may be formed from multiple sets of spatially contiguous patches while targets often consist of one block of contiguous patches as shown in <Ref>.



  
Predictor and Loss
The task of predictor given targets and context is analogous to the task of
supervised prediction. 
Given a context as input along with a certain condition, it aims to predict the target representations.
Here, the condition involves the mask tokens, which are created from shared learned parameters, as well as positional encoding created from center points associated with the targets.
Because the predictor’s task is to predict the representation produced by the target encoder, the loss can be defined to minimize the disagreement between the predictions and targets.
To this end, we utilize Smooth L1 loss to measure the dissimilarity between target embedding and predicted embedding.
The parameters of the target encoder, as well as the point tokenizer, are updated via gradient-based optimization, while the parameters of the target encoder are updated via the exponential moving average of the context encoder.










§ EXPERIMENTS

In this section, we first describe the details of self-supervised pre-training.
Following that, we compare the performance of the learned representation to some of the top-performing self-supervised learning methods in the point cloud domain, especially those that utilize the ShapeNet <cit.> dataset in pre-training.
We specifically evaluate the learned representation using linear probing, end-to-end fine-tuning, and a few-shot learning setting.



 §.§ Self-Supervised Pre-training

We pre-train our model on training set of ShapeNet <cit.> following the previous studies utilizing the standard Transformer <cit.> architecture such as Point-MAE <cit.>, Point-M2AE <cit.>, PointGPT <cit.>, and Point2Vec <cit.>.
The dataset consists of 41952 3D point cloud instances created from synthetic 3D meshes from 55 categories.
As previously mentioned, standard Transformer <cit.> architecture is used for the context and target encoder as well as the predictor.  
During pre-training, we set the number of
center points to 64 and the group size to 32. The point tokenization is applied to the input point cloud containing 1024 points per object. We set the depth of the Transformer in the context and target encoder to 12 with the embedding width of 384 and 6 heads. For the predictor, we use the narrower dimension of 192 following I-JEPA <cit.>. 
The depth of the predictor is set to 6, and the number of heads
is set to 6.



 §.§ Downstream Tasks

In this section, we report the performance of the learned representation on several downstream tasks.
Following the previous studies <cit.>, we report the overall accuracy as a percentage.
Unless specified otherwise, we report the mean accuracy and standard deviation from 10 independent runs, each with a different seed, in order to account for variability across independent runs.



  
Linear Probing.
After pre-training on ShapeNet <cit.>, we evaluate the learned representation via linear probing on ModelNet40 <cit.>. 
Specifically, we freeze the learned context encoder and place the SVM classifier on top. To enforce invariance to geometric transformation, we utilize max and mean pooling on the output of the Transformer encoder <cit.>.
Here, we utilize 1024 points for both training and test sets. 
As shown in Table <ref>, our method achieves state-of-the-art accuracy, providing +0.8% performance gain, showing the robustness of the learned representation. 








  
End-to-end Fine-Tuning
We also investigate the performance of the learned representation via end-to-end fine-tuning.
After pre-training, we utilize the context encoder to extract the max and average pooled outputs.
These outputs are then processed by a three-layer MLP for classification tasks.
This class-specific head as well as the context encoder is fine-tuned end-to-end on ModelNet40 <cit.> and ScanObjectNN <cit.>.
ModelNet40 consists of 12311 synthetic 3D objects from 40 distinct categories, while ScanObjectNN contains objects from 15 classes, each containing 2902 unique instances collected by scanning real-world objects.
For ModelNet40, we sub-sample 1024 points per object and sample 64 center points with 32 points in each point patch.
On the other hand, we utilize all 2048 points for the ScanObjNN dataset and sample 128 center points with 32 nearest neighbors for the grouped points.
As shown in Table <ref>, our method achieves competitive results when compared to other state-of-the-art methods.
Especially, in the  variant of the ScanObjNN <cit.> dataset,  which presents a realistic representation of a point cloud that includes both the object and its background, our method achieves an improvement of +1% over the best-performing method.
This shows the learned representation obtained from pre-training with Point-JEPA can easily be transferred to a classification task.








  
Few-Shot Learning
We conduct few-shot learning experiments on Modelnet40 <cit.>. 
Experiments are done in m-way, n-shot setting as shown in Table <ref>.
Specifically, we randomly sample n instances of m classes for training. We select 20 instances of m support classes for evaluation.
Under one setting, we run 10 independent runs under a fixed random seed on 10 different folds of the dataset, and we report the mean and standard deviation of overall accuracy.
As shown in Table <ref>, our method exceeds the performance of the current state-of-the-art in all settings.
Our method yields a +1.1% improvement in the most difficult 10-way 10-shot setting, showing the robustness of the learned representation of Point-JEPA, especially in the low-data regime.





§ CONCLUSION

This work introduced Point-JEPA, a joint embedding predictive architecture applied to point cloud objects.
In order to efficiently select targets and context blocks even under the invariance property of point cloud data, we introduced a sequencer, which orders the center points and their corresponding tokens by iteratively selecting the next closest center point.
This eliminates the necessity of computing spatial proximity between every pair of tokens or embedding vectors when sampling the targets and context.
Our method experimentally achieves state-of-the-art performance in several downstream tasks, showing a strong learned representation.
Specifically, Point-JEPA excels in few-shot learning as well as linear evaluation, making the method highly useful when there is a large amount of unlabeled data and a limited amount of labeled data.
It is also worth noting that Point-JEPA converges much faster during pre-training, offering a more efficient pre-training alternative in the point cloud domain.

    ieeenat_fullname
    



 





§ FURTHER PRE-TRAINING DETAILS




  
Optimization.
We utilize AdamW <cit.> optimizer with cosine learning decay <cit.>.
Starting from learning rate of 10^-5, we increase it to 10^-3 in the first 30 epochs and decay it to 10^-6.
The batch size for pre-training is set to 512, and the beta value for Smooth L1 loss is set to 2, similar to Point2Vec <cit.>.
The target encoder and context encoder initially have identical parameters. 
The context encoder's parameters are updated via backpropagation, while the target encoder's parameters are updated using the exponential moving average of the context encoder parameters, that is 
θ_t+1τθ_t + (1 - τ)θ_t
where τ∈ [0,1] denotes the decay rate, θ denotes the parameter of the context encoder, and θ denotes the parameter of the target encoder.
We gradually increase the decay rate of the exponential moving average from 0.995 to 1.0 during pre-training.



  
Masking and Ordering.
To determine the sequence of tokens, we utilize the iterative ordering of associated center points, as previously mentioned.
We chose the starting point in this sequence with the lowest sum of its coordinates.
This method allows the sequence to start from a point on the outer edge of the object rather than from a point within the object's interior.
This consistency in selecting the initial point is experimentally shown to deliver a slightly better-learned representation than taking the first available index (See <ref> for the details).

For masking, we define a range of ratios with both upper and lower limits similar to I-JEPA <cit.>.
To start with, we clarify that the term block refers to a sequence of tokens and their corresponding encoded embedding vectors that are contiguous in their indices. 
Because of the sequencing process applied before the target and context selection, most contiguous tokens and encoded embedding vectors are also spatially contiguous.
For the target, we randomly select 4 blocks of embedding vectors from within the 0.15 to 0.2 range.
We then remove the corresponding tokens of encoded embedding vectors that have already been chosen as targets for further selection.
Following this, we sample tokens within the ratio range of 0.4 to 0.75 out of the available tokens. 
Because some of the tokens are not available for context selection, we note that context block usually consists of multiple sets of tokens that are spatially contiguous.




§ FURTHER RESULTS

Following previous studies <cit.>, we report the performance of Point-JEPA in part segmentation task.
Here, we utilize the ShapeNetPart <cit.> dataset, consisting of 16881 objects from 16 categories.
We utilize the identical architecture employed in Point2Vec <cit.> for this task.
Specifically, we take the embeddings from 4^th, 8^th, and 12^th Transformer block and take the average of them.
Then, we apply mean and average pooling to this averaged output. 
The max and mean pooled embedding along with a one-hot encoded class label of an object is used as a global feature vector for the object.
The original averaged output is also up-sampled using the PontNet++ <cit.> feature propagation layer to create a feature vector for each point.
Then, each feature vector is concatenated with the global feature vector. A shared MLP is utilized on this concatenated vector to predict the segmentation label for the given point.
Although the Point-JEPA shows competitive results as shown in Table <ref>, its performance is slightly worse than the state-of-the-art methods.
This, coupled with the fact that its learned presentation has shown to be superior in classification tasks suggests that Point-JEPA's learned representation emphasizes global features over local features.






§ ABLATION

In this section, we conducted thorough ablation studies to understand the effect of moving parts of Point-JEPA.
We pre-train Point-JEPA on the ShapeNet <cit.> dataset under various settings and evaluate the learned representation with linear probing on the ModelNet40 <cit.> dataset. 




  
Masking Strategy.
We investigate the impact of the masking type on the performance.
We consider single-block masking and multi-block masking.
For single-block masking strategies, we consider random masking and contiguous masking. 
For random masking, we randomly select the 60% of indices out of all encoded embedding vectors.
Similarly, for contiguous masking, embedding vectors that are spatially contiguous are selected. 
In this setting, all tokens not corresponding to the selected target blocks are used as context (denoted as rest).
On the other hand, in the muli-block masking setting, we sample multiple possibly overlapping spatially contiguous embedding vectors as targets, and we remove the corresponding tokens already utilized for targets during context selection.
In this setting, we set the ratio range of 0.15 to 0.2 for targets, while we set the ratio range of 0.4 to 0.75 for context. 
As shown in Table <ref>, the single-block masking achieves sub-optimal performance regardless of the spatial contiguity of the target embedding.
It shows that our method learns stronger representation by utilizing a smaller amount of targets with a larger frequency.



  
Ratio of Targets.

We change the ratio of the selected embedding vectors for the target selection while keeping the number of target blocks and the ratio of context tokens fixed.
As shown in <Ref>, the performance increases when you increase the ratio to a certain point. 
However, beyond this point, further increasing the ratio results in decreased performance.
This implies that Point-JEPA does not require a large size for the target blocks and benefits from a sufficient amount of available tokens for context selection.






  
Ratio of Context.
In this study, we change the ratio of tokens selected for context encoding while keeping the number of targets and the ratio range for targets fixed.
As shown in Table  <ref>, having a relatively large difference between the lower and upper bound of the ratio can improve performance.
In other words, Point-JEPA learns a better representation when the number of selected context tokens varies more between training iterations.
Additionally, when the upper bound of the ratio is somewhat constrained, we see increased performance.  






  
Number of Target Blocks.
We will also consider the effect of the number of blocks chosen for targets on the performance of the learned representation while we keep the ratio for targets and context fixed.
As shown in Table <ref>, the performance increases as we increase the number of targets.
However, the performance decreases as you increase the number of target blocks after a specific frequency. 
Similarly to what we observed in Section <ref>, we note that our method benefits from having a sufficient amount of tokens available for context encoding.







  
Predictor Depth
We also study the effect of the predictor's depth on the learned representation.
To this end, we vary the predictor depth and observe its effect on the linear evaluation accuracy.
As shown in Table <ref>, Point-JEPA benefits from a deeper predictor.



  
Initial Point for Sequencer

In this section, we study the importance of the initial center point in the sequencer.
We compare a sequencer that takes the initial index of the center point against another sequencer that takes the minimum sum of the coordinates as the initial point.



As previously mentioned, selecting the first center point presented in the sequencer does not guarantee the selection from the edge of the object.
On the other hand, selecting the point with minimum coordinate sum at least guarantees that the starting center point in the sequencer lies near the edge of the object.
As shown in Table <ref>, taking the point near the edge of the object helps with learning stronger representation.




  
Learned Representation
In order to qualitatively analyze the learned representation, we reduce the dimension of the learned representation by utilizing t-SNE <cit.>.
Specifically, we introduce max and mean pooling on the output of the context encoder, similar to the classification setup, and apply t-SNE on the pooled embedding. 
We visualize the learned representation on ModelNet40 <cit.> with no fine-tuning on the dataset. 
In other words, the context encoder has not been trained with this data.
Despite this, our context encoder produces discriminative features as shown in <Ref>, showing the robustness of the learned representation.
















