Paper: State Representation Learning Using an Unbalanced Atlas (ICLR-2024)

Reviewer comment (Experiments):
In the ablation study, the model "-UA" is not clearly specified. If authors do not use the modifications of 7, that means the -MMD loss is also void. What does that yield? Are authors using a single projection head in this case?

Author response:
"-UA" uses the original prediction targets (as in Eq. 1), and the dilated prediction targets (the first half of Eq. 7) are not used. It uses multiple heads and $L_Q$ is still present. We have added more descriptions in the text.

Reviewer comment (Evaluation):
CIFAR experiments show that the gains are much lower in these experiments compared to the ones obtained in ATARI games. A discussion towards this end is not provided but it would be very valuable for the readers.

Author response:
We very much agree that this was a primary concern and we have run more experiments to check the reasons behind. We have added discussion on this and supplemented a small-scale experiment in Table 6 at the end of Appendix B. In general, we have used the same set of hyper-parameters from MSimCLR, but that may not be optimal. We noticed that a small $\tau$ or a linear increasing scheme of $\tau$ increases the performance of SimCLR-UA significantly.

Reviewer comment (Writing):
Notation in the presentation of the method seems a bit inconsistent. I recommend authors to improve the consistency in the notation.

Author response:
We checked the paper carefully and spotted that the primary inconsistency in notations were caused by using $\theta$ to denote the NN encoder when formalizing DIM-UA. We thought about the pros and cons and decided to remove the use of $\theta$. We now use $f$ to denote the encoder and made changes based on that. Hopefully, this can increase the readability.

    We have also added more descriptions to notations far apart from where they were first introduced.

Reviewer comment (Evaluation):
Can authors discuss further why the CIFAR experiments do not show a similar improvement?

Author response:
We very much agree that this was a primary concern and we have run more experiments to check the reasons behind. We have added discussion on this and supplemented a small-scale experiment in Table 6 at the end of Appendix B. In general, we have used the same set of hyper-parameters from MSimCLR, but that may not be optimal. We noticed that a small $\tau$ or a linear increasing scheme of $\tau$ increases the performance of SimCLR-UA significantly.

Reviewer comment (Writing):
Can authors please improve the notation consistency in the presentation?

Author response:
We checked the paper carefully and spotted that the primary inconsistency in notations were caused by using $\theta$ to denote the NN encoder when formalizing DIM-UA. We thought about the pros and cons and decided to remove the use of $\theta$. We now use $f$ to denote the encoder and made changes based on that. Hopefully, this can increase the readability.

    We have also added more descriptions to notations far apart from where they were first introduced.

Reviewer comment (Writing):
I am struggling to understand the rationale behind the DIM-UA algorithm.

Author response:
Thank you very much for your kind review! We were sorry that the notations in DIM-UA were hard to follow. We put some effort on making the notations in DIM-UA more consistent. Hope it helps with the readability.

Reviewer comment (Theory):
What is the motivation for redefining the score function L_GL?

Author response:
We redefined $L_{GL}$ because the UA paradigm uses dilated prediction targets in pretraining, instead of using the normal prediction targets.

Reviewer comment (Experiments):
Equation 9 has a hyperparameter for L_Q but not for L_GL or L_LL. Why?

Author response:
We have a hyper-parameter $\tau$ for $L_Q$ because this hyper-parameter controls the MMD loss on membership probabilities. For $L_{GL}$ or $L_{LL}$, they are directly controlled by the learning rate of the optimizer. We have showed that this hyper-parameter $\tau$ is indeed important. Please check Table 6 at the end of Appendix B to see the effects of changing $\tau$.

Reviewer comment (Writing):
The writing of the paper is not clear in several points as some parts are difficult to follow.

Author response:
Thank you very much for your kind review! We were sorry that the notations in DIM-UA were hard to follow. We put some effort on making the notations in DIM-UA more consistent. Hope it helps with the readability.

Reviewer comment (Presentation):
Figure 1 is not quite clear to me and the caption is not very explicit.

Author response:
We have added more content to the caption of old Figure 1 (now Figure 2).

Reviewer comment (Evaluation):
The improvements reported in Table 2 seem relatively small. Can you comment on this further?

Author response:
The old Table 2 now is labeled as Table 4, which has the accuracy scores. Our DIM-UA method has a mean accuracy of 76%, compared to 73% and 71% of ST-DIM using 256 units and using the same amount of units. First, VAE underperforms too much compared to ST-DIM, which may make a 3% increase seem not that big, but 3% increase in accuracy is quite decent since the difference between DIM-UA and ST-DIM is not even near to the difference between generative methods and contrastive methods.

    Moreover, when comparing ST-DIM to DIM-UA with the same large number of units, DIM-UA wins by 5%, which suggests that a large number of hidden units in ST-DIM does not help with the performance. This is strong empirical evidence that our UA paradigm can create more effective representations.

Reviewer comment (Evaluation):
Why is the linear evaluation accuracy on CIFAR10 a suitable evaluation metric? The improvements shown here (Table 3) are very marginal at best.

Author response:
This is a great point. We also thought that the performance of SimCLR-UA in Table 3 (now Table 2) needed more discussion. Our emphasis in the past few days has put on finding the key reason behind this. We have supplemented a small-scale experiment in Table 6 at the end of Appendix B. Mainly, we argued that the main reason why SimCLR-UA did not outperform SimCLR by a large margin is that it used the same set of hyper-parameters as in MSimCLR, which were not optimal. This, in turn, also suggests that SimCLR-UA beats MSimCLR by a large margin with no doubt.

Reviewer comment (Writing):
The clarity of the introduction could be enhanced by providing a more explicit context for the specialized terminology introduced (see Q1-3).

Author response:
1. We agree that the terminology should be better introduced. We now have more descriptions in the caption of Figure 2.

Reviewer comment (Presentation):
Tables 1 and 2 appear to be redundant, presenting analogous results through different evaluative metrics (F1 score and Accuracy, respectively). Although a comprehensive evaluation is encouraged, putting these two sizable tables back to back in the main paper gives the impression of lacking sufficient materials for the paper. It would be more appropriate to consolidate these findings, perhaps through a combined analysis or in supplementary materials, to avoid repetition and maintain the conciseness of the paper.

Author response:
Thank you very much for your detailed review! We have moved Table 2 to Appendix A (now Table 4). Other part of weakness will be addressed in answering questions.

Reviewer comment (Novelty):
The paper lacks a clear statement of its underlying motivation and significance, which is pivotal for readers to comprehend the value and potential impact of the research (see Q5).

Author response:
5. This is a very good point. The reasons why this work is important should be better explained. Thus, we emphasized on our motivation in Introduction and extended the contribution summary. One important note is that our work can help build large models in deep learning by enlarging the heads (complementary to the typical research that deepens the backbones).

Reviewer comment (Writing):
The introduction used specialized terminology that may not be universally familiar, necessitating additional clarification for a broader audience. Specifically, the first sentence of the third paragraph introduces concepts such as *manifold*, *atlas*, *local structure*, and *chart*, which would benefit from further exposition to contextualize the study and its objectives.

Author response:
Dear reviewer hci6,

Thank you very much for your response! Indeed, an annotated version would be very helpful. Please check our updated pdf. The terms are mainly introduced in the caption of Figure 2, since we thought that those terms can be better understood with a figure illustration. 

We agree that the term of local structure is still not independently introduced. However, please note that this term is not anything special in this context, it is more like a broad concept in maths. For instance, the concept of locality here is not much different compared to when we say that a function is locally smooth, or locally differentiable. It can be better understood to consider that the domain for a chart is local to the whole manifold for an atlas.

Reviewer comment (Writing):
The paper's motivation remains unclear, partly owing to the use of undefined terms. The concept of an *atlas*, and particularly the distinction between *unbalanced* and *balanced* atlases within this framework, needs clarification. The terms *prior distribution* and *membership probability distribution* introduced later also lack clear definitions, impeding the reader's understanding.

Author response:
2. The concept of an atlas is now described in Figure 2. Moreover, we described the notion of a uniform prior and unbalanced atlas in Introduction (above the contribution summary). We used the term 'balanced' in contrast to 'uniform'. While introducing these notions, we also strengthened our motivation towards our method. The membership probability can be viewed as a simple score function (by using Softmax) to attribute a probability score to each chart.

Reviewer comment (Writing):
In the 3rd paragraph of Section 4, *d* and *n* are used without proper definition.

Author response:
3. We have added short descriptions of $n$ and $d$ in Section 4.

Reviewer comment (Experiments):
This paper suggests that pre-training a model using a reinforcement learning task and then fine-tuning it on downstream reinforcement learning tasks is beneficial. However, this point is not fully demonstrated because the authors did not compare the proposed method with other self-supervised learning methods, as reviewed in the introduction, e.g., contrastive models (SimCLR is compared) and generative models (none is compared).

Author response:
4. We started our work with a focus on SRL, and we saw that generative methods did not perform well in SRL tasks (see the much lower F1 scores of VAE compared to CPC in Table 1). It was never our intention to omit generative methods, but we simply thought that contrastive methods were more suitable in SRL tasks. Of course, there are many other tasks where generative methods are more suitable, and it can be interesting future work.

Reviewer comment (Novelty):
From the introduction section, it is not intuitive to me why this study is important. For example, the last paragraph lists technical achievements but does not convey their broader significance. Specifically, (1) fitting the ST-DIM to UA paradigm (why UA paradigm is important?), (2) detailed ablations for better design choices (that's standard, not sure if it counts as a contribution), and (3) representing a manifold with a larger number (why this is important anyway?). Not limited to the introduction section, the authors did not describe the significance of the proposed method and all these ablation studies in the entire paper.

Author response:
Dear reviewer hci6,

Thank you very much for your response! Indeed, an annotated version would be very helpful. Please check our updated pdf. The terms are mainly introduced in the caption of Figure 2, since we thought that those terms can be better understood with a figure illustration. 

We agree that the term of local structure is still not independently introduced. However, please note that this term is not anything special in this context, it is more like a broad concept in maths. For instance, the concept of locality here is not much different compared to when we say that a function is locally smooth, or locally differentiable. It can be better understood to consider that the domain for a chart is local to the whole manifold for an atlas.

For the motivation of enlarging heads, it is because we face bottlenecks when enlarging the backbone. For instance, over-parameterization is common for vision transformers (ViTs) and over-parameterized ViTs might perform worse on ImageNet 1K. When facing such bottlenecks, an alternative to build larger deep learning models would naturally be increasing the number of hidden units in output heads.