Paper: Dense Representation Learning for a Joint-Embedding Predictive Architecture (ICLR-2024)

Reviewer comment (Presentation):
Some symbols and equations make the paper a bit hard to follow and might be redundant. For example, the symbols in equation 2-5 — adding symbols of features in Figure 2 might make it easier to read and understand. Furthermore, providing structure figures like those of transformer blocks, indicating what the $Q$, $K$ and $V$ features of $h_{\theta}$ and $h_{\hat{\theta}}$ would also make it quicker to understand.

Author response:
**Q1. Clarity for Equations and Figure 2**

We acknowledge your suggestion to improve the clarity of symbols and equations in the paper, particularly around equations 2-5. To address this, we updated Figure 2 to include symbols of features, enhancing readability and understanding. Additionally, we provided a more detailed structure of local aggregation head, clearly indicating the Q, K, V of features $h_\theta$ and $ h_{\hat{\theta}}$. This should make the methodology quicker and easier to grasp for readers.

Reviewer comment (Writing):
The subscript of $s_I$ and $x_t$ are weird in $s^{LAT}_{i} = h_{\hat{\theta}}({x_j}_{j\in P_i}, x_t)$ of Eq. 4. What is the relationship between subscript $t$ and $i$? Can we understand it as $s^{LAT}_{t} = h_{\hat{\theta}}({x_j}_{j\in P_i} , x_t)$?

Author response:
**Q2. Clarity for Equation 4**

Your observation about the confusing subscript in equation 4 is valid. The subscript t and i in \( s^{LAT}_i = h_{\hat{\theta}}(x_{j_{j\in P_i}}, x_t) \) indeed require clarification. We will correct and clarify this in the final draft. The subscript t refers to the target patch, while i refers to the index of the patch in the local aggregation process. We will ensure this distinction is more clearly articulated in the final draft.

Reviewer comment (Theory):
The motivation of the Local Aggregation Target module might require further explanation. Why does adding such a module improve performance? The modification proposed in this paper is simple, thus, more explanation, analysis, and insights would be beneficial for the community to understand it better.

Author response:
**Q3. Clarity for Local Aggregation Target**

We understand your request for a more detailed explanation of the LAT module's motivation and its contribution to performance improvement. The LAT module is designed based on the premise that semantically similar patches should have similar targets. By incorporating LAT, our model can achieve better semantics and discriminative capabilities, as not every patch is treated with the same target. We will expand on this explanation in the manuscript, providing more analysis and insights to help the community understand the benefits of this approach.

Reviewer comment (Writing):
The reviewer does not quite understand the structure of $h_{\theta}$. In Eq. 4, it seems that $s_x$ and $x_c$ are different embeddings and Table 7 studies $h_{\theta}$ using cross-attention. However, from Figure 2, the definition of $s_x$ around Eq. 2 ('... given the output of the context encoder, $s_x$, ...') and the definition of $x_c$ around Eq. 4 ('where $s_x$ denotes context embeddings and $x_t$, $x_c$ denote the averaged embeddings from all patches in the target encoder and only unmasked patches in the context encoder, respectively') suggest that $s_x$ and $x_c$ are the same thing.

Author response:
**Q4. Clarity for Local Aggregation Head structure**

Regarding the structure of $h_\theta$ and the apparent confusion between $s_x$ and $x_c$, we will make sure to clarify these in our revised paper. $s_x$ and $x_c$ indeed refer to different embeddings. $s_x$ denotes the context embeddings, while $x_c$ represents the averaged embeddings from unmasked patches in the context encoder. This distinction is crucial for understanding the operation of our model, and we will provide a clearer explanation to eliminate any confusion in the final draft.

Reviewer comment (Novelty):
The results are not surprising, dense self-supervised learning has been explored extensively in the literature, and the same conclusion and gains on segmentation tasks have already been observed. DenseCL [1] uses a contrastive loss function to match local vectors. VICRegL [2] combines local and global aggregation and uses a similar nearest neighbour search for semantically similar patches. ODIN [3] computes segmentation masks online and aggregates the patches corresponding to the same regions. These papers are not mentioned or cited, I would therefore recommend to do a more extensive literature review, and highlight the novelty within the dense SSL literature.

Author response:
**Q1. Our novelty and discussion to dense SSL literature**

We acknowledge your suggestion to include references such as DenseCL, VICRegL, and ODIN in our literature review. While these works are indeed relevant, we emphasize that Dense-JEPA's core novelty lies in enhancing the latent JEPA architecture for learning local/dense representations, a different focus from the mentioned works. Our approach is not merely about dense representation learning but about innovatively applying it to the existing latent JEPA framework. For example, our idea is independent to hand-crafted augmentations and we believe that ours can be incorporated with further augmentation-based dense SSL ideas. Nevertheless, we will add regarding discussion in final draft to position our work appropriately within the field.

Regarding why our work is important and have significance, we belive that mask reconstruction on latent space will have gain much attention on its efficacy and efficiency, so study of training objective on latent space is crucial. Therefore, our investigation of the latent objective is invaluable and could inspire future research direction to explore better SSL objective, unlike the previous stereotype of masked reconstruction objective on pixel space.

Reviewer comment (Experiments):
There is no comparison with DINO on most of the tasks. The DINO family of methods [4,5] has shown that dense loss functions might actually not be required to learn very good dense features, and offers the best performance on local tasks as of today. For example, on DAVIS-2017, DINO reports 61.8 J&M with a ViT-S/16, while the best number reported in the paper is 58.3 with a much larger ViT-L backbone. On classification on ImageNet, a DINOv2 ViT-L model is over 86 points. The gains over I-JEPA are interesting, but the overall performance is very far from the state-of-the-art. MAE is not really a good baseline as it performs very poorly on frozen tasks, which are standard to evaluate SSL methods.

Author response:
**Q2. Comparision with DINO family methods**

Thank you for highlighting the comparison between our Dense-JEPA method and the DINO family of methods. We acknowledge the impressive achievements of DINO and DINOv2 in the domain of self-supervised learning. However, we would like to emphasize a few key aspects of our work to clarify our contributions and the significance of our results.

1. Contributions on Numbers

| Method           | Pretrain data | Seg (mIoU) | Det (AP_b) | Det (AP_m) |
|------------------|---------------|------------|------------|------------|
| DINO             | ImageNet-1K   | 46.8       | 50.1       | 43.4       |
| MAE              | ImageNet-1K   | 48.1       | 50.3       | 44.9       |
| DenseJEPA (ours) | ImageNet-1K   | **49.0**   | **50.9**   | **45.6**   |

Our results shown in the Table above, particularly in segmentation and detection tasks using a ViT-B/16 backbone, demonstrate significant improvements over DINO. While DINO's performance is commendable, our method shows a marked advantage in these specific areas. This is a crucial aspect of our contribution, as it highlights the effectiveness of Dense-JEPA in handling diverse and complex visual tasks. 
In the context of DAVIS-2017, where linear evaluation on frozen tasks is critical, our method also shows promising results. While DINO reports higher J&M scores with a smaller ViT-S/16 backbone, our Dense-JEPA method achieves competitive scores with a larger ViT-L backbone in Table 2, suggesting potential for further optimization and tuning to enhance performance.


2. Comparison Fairness with DINOv2

It is important to note that DINOv2 was trained on a significantly larger dataset compared to the dataset used for training Dense-JEPA. This difference in training conditions can substantially affect the performance outcomes. Therefore, a direct comparison may not entirely reflect the intrinsic strengths of our method. Our approach still shows competitive performance under more constrained data conditions, which is a testament to its robustness and efficiency.
Compared to DINOv2 using ViT-L/16 using different pre-training resolutions in Figure 6 in the original paper, our Dense-JEPA with ViT-B/16 backbone achieves competitive results (49.0 mIoU vs 46.2 mIoU). However, it should be noted that DINOv2 performed linear evaluation settings on the ADE20K dataset by freezing representations at different resolutions. Although the comparison is not totally fair, these results further demonstrate the competitive fine-tuning performance on segmentation by learning local semantics during self-supervised pre-training. We believe that our new framework can be scaled up to more pre-training data to achieve better downstream performances.

3. Fine-Tuning Comparison with DINO and MAE

In terms of full fine-tuning capabilities, particularly using a ViT-B/16 model, our method exhibits superior performance compared to both DINO and MAE. This is a crucial metric, as fine-tuning is often essential in real-world applications of self-supervised learning models. Our method's adaptability and effectiveness in fine-tuning scenarios underscore its practical utility beyond standard benchmark tasks.

| Method           | Pretrain data | Backbone | Top-1 Accuracy |
|------------------|---------------|----------|----------------|
| DINO             | ImageNet-1K   | ViT-B/16 | 82.8           |
| MAE              | ImageNet-1K   | ViT-B/16 | 83.6           |
| DenseJEPA (ours) | ImageNet-1K   | ViT-B/16 | **84.6**       |
| MAE              | ImageNet-1K   | ViT-L/16 | 85.9           |
| DenseJEPA (ours) | ImageNet-1K   | ViT-L/16 | **86.6**       |

Overall, while we recognize the excellence of the DINO family in certain areas, our Dense-JEPA method demonstrates its unique strengths and contributions, especially in segmentation, detection, and fine-tuning contexts. We believe these aspects significantly contribute to the field of self-supervised learning and merit recognition.

Reviewer comment (Experiments):
The gains over I-JEPA are probably not worth the complexity. The paper is missing a study on the tradeoff in terms of compute compared to I-JEPA. In terms of memory, running time and data efficiency.

Author response:
**Q3. Computational complexity of Dense-JEPA**

We understand your concern about the computational complexity of Dense-JEPA compared to I-JEPA. In our revised manuscript, we include a detailed analysis of total pre-training epochs, max memory usage, and running time per training step in Section B, demonstrating that Dense-JEPA maintains similar efficiency to I-JEPA while achieving superior performance. Besides, our method is much efficient and achieves short training time than MAE.

| Method            | Pre-train Epochs | Max Memory  per GPU | Running Time  per Step |
|-------------------|------------------|---------------------|------------------------|
| I-JEPA            | 600              | 21.9G               | 606.4 ms               |
| Dense-JEPA (ours) | 600              | 21.9G               | 608.2 ms               |

Reviewer comment (Presentation):
Figure 2 is very unclear. There is not much additional information compared to the similar figure in the I-JEPA paper, and the explanation for “Local aggregation head” and “Masked Semantic Neighboring” is only in the text. I would recommend clarifying the Figure and emphasizing the difference with I-JEPA.

Author response:
**Q4. Clarify of Figure 2**

We acknowledge the need to improve Figure 2 for better clarity and differentiation from I-JEPA. The revised figure includes a clearer distinction and detailed explanation between “Local Aggregation Head” and “Masked Semantic Neighboring” components, emphasizing the distinctions from I-JEPA.

Reviewer comment (Presentation):
The visualization of Figure 1 is not convincing. I-JEPA patches already contain highly semantic information and Dense-JEPA clearly brings a bias towards neighboring patches, which is not necessarily a good thing as we would like to let the system pick patches that are far away but semantically very similar.

Author response:
**Q5. Qualitative anlaysis on Dense-JEPA**

In order to qualitatively demonstrate the effectiveness of the proposed Dense-JEPA, we further compared learned attention maps from the target encoder using pretrained I-JEPA and our Dense-JEPA, and attention maps from the cross-attention layer in the proposed Local Aggregation Target (LAT) modules in Appendix section C in our revised manuscript. Our Dense-JEPA achieves much better attention maps, compared to I-JEPA. 
Furthermore, we provide more examples on learned cosine similarity map to further demonstrate that our Dense-JEPA performs effectively by encouraging the model to learn dense representations. Regarding Figure 1, we will enhance the visualization in the final draft to more convincingly demonstrate how Dense-JEPA improves upon I-JEPA's semantic information capture, especially in the context of neighboring patches.

Reviewer comment (Experiments):
Why do you use cross-attention to do the aggregation? Have you tried max-pooling or other types of pooling? Do you consider that it is an essential component of the method?

Author response:
**Q6. Role of cross-attention in Local aggregation head**

The choice of cross-attention for aggregation in Dense-JEPA is deliberate. We experimented with other pooling methods, such as average pooling and max pooling, but found them less effective. Cross-attention allows for more nuanced aggregation of features, which is crucial for our model's performance. We included additional results in Table B.1 in the Appendix to support this claim.

Reviewer comment (Theory):
The transformer blocks in the encoder are performing self-attention between every pair of tokens, which might already have the effect of grouping patches using a distance in the embedding space. Have you thought about that and do you think that your method theoretically brings something?

Author response:
**Q7. Rationale of Dense-JEPA**

Your question about the role of transformer blocks in our encoder is insightful. While these blocks do perform self-attention among token pairs, our method enhances this process by focusing on semantically relevant patches. This targeted approach is theoretically significant as it guides the transformer to learn more meaningful local semantics. We provided additional evidence in the revised manuscript to demonstrate the superiority of our targeted patch representations over generic self-attention mechanisms, as shown in Figure C1-C4 in the Appendix.

Reviewer comment (Novelty):
The novelty can be improved. It seems that Dense-JEPA is just an improvement on I-JEPA, especially Figure 2 is so similar to figure 2 in I-JEPA. The main points proposed by the paper are not well presented.

Author response:
You pointed out that Dense-JEPA's novelty is questionable, particularly in reference to Figure 2's similarity to I-JEPA's Figure 2. We acknowledge this observation and updated Figure 2 to better illustrate the distinctions between Dense-JEPA and I-JEPA.

Regarding why our work is important and have significance, we believe that mask reconstruction on latent space will have gain much attention on its efficacy and efficiency, so study of training objective on latent space is crucial. Therefore, our investigation of the latent objective is invaluable and could inspire future research direction to explore better SSL objective, unlike the previous stereotype of masked reconstruction objective on pixel space. 
To clarify, Dense-JEPA's novel contribution lies in its unique approach to enhancing target representations. Unlike I-JEPA, Dense-JEPA employs a more nuanced strategy in selecting neighboring patches based on semantic similarity. This approach leads to a more accurate and granular understanding of local semantics, significantly contributing to the field of dense representation learning.

Reviewer comment (Writing):
The descriptions of MSN and LAT are ambiguous. From section 3.2, MSN is just using the presentation extracted by target encoder to find the most similar neighbor patches, which is prepared for LAT to aggregate from these similar patches, so MSN could not be regarded as an individual module.

Author response:
We agree that the descriptions of the MSN (Masked Semantic Neighboring) and LAT (Local Aggregation Target) modules could be more precise. MSN is indeed preparatory for LAT, but it is crucial for establishing fine-grained semantic relations among patches. On the other hand, LAT focuses on a coarser level of aggregation.

In the revised paper, we provided a clearer distinction between these two modules in Figure 2. MSN is not just a preparatory step but a critical component in refining the target representation by identifying the most semantically relevant neighbors. LAT then aggregates these for a comprehensive representation. We will revise further for better clarity in the final draft.

Reviewer comment (Experiments):
Without LAT, does MSN actually influence the training procedure (Table 5, the third row)?

Author response:
Regarding your question about the role of MSN without LAT (Table 5, second row), MSN independently contributes to the training procedure by identifying the most semantically relevant patches. This process is crucial for Dense-JEPA's effectiveness, as it ensures that even in the absence of LAT, the model doesn't revert to a coarse, generalized representation but maintains a degree of fine-grained semantic understanding (e.g., 53.7->55.2 scores in Table 5). Adding MSN to the vanilla baseline highly increases the results of the video segmentation performance regarding all metrics, which validates the benefit of masked semantic neighboring in finding semantically similar neighboring patches for masked patches (e.g., 53.7->55.2 scores in Table 5).

Without MSN, Dense-JEPA would lack the nuanced understanding of local semantics, leading to a more generalized and potentially less accurate representation (e.g., 57.1->54.6 scores in Table 5).