State Representation Learning Using an Unbalanced Atlas
Authors: Li Meng
Morten Goodwin
Anis Yazidi
Paal E. Engelstad
Venue: ICLR 2024 poster (2024)
Keywords: Self-supervised learning
State representation learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We introduce a new self-supervised learning paradigm using an unbalanced atlas to represent a manifold and design a state representation learning method based on the paradigm.

Abstract: The manifold hypothesis posits that high-dimensional data often lies on a lower-dimensional manifold and that utilizing this manifold as the target space yields more efficient representations. While numerous traditional manifold-based techniques exist for dimensionality reduction, their application in self-supervised learning has witnessed slow progress. The recent MSimCLR method combines manifold encoding with SimCLR but requires extremely low target encoding dimensions to outperform SimCLR, limiting its applicability. This paper introduces a novel learning paradigm using an unbalanced atlas (UA), capable of surpassing state-of-the-art self-supervised learning approaches. We investigated and engineered the DeepInfomax with an unbalanced atlas (DIM-UA) method by adapting the Spatiotemporal DeepInfomax (ST-DIM) framework to align with our proposed UA paradigm. The efficacy of DIM-UA is demonstrated through training and evaluation on the Atari Annotated RAM Interface (AtariARI) benchmark, a modified version of the Atari 2600 framework that produces annotated image samples for representation learning. The UA paradigm improves existing algorithms significantly as the number of target encoding dimensions grows. For instance, the mean F1 score averaged over categories of DIM-UA is~75% compared to ~70% of ST-DIM when using 16384 hidden units.

Published as a conference paper at ICLR 2024
STATE REPRESENTATION LEARNING USING AN UN- BALANCED ATLAS
Li Meng
Morten Goodwin University of Agder
Anis Yazidi
University of Oslo Oslo, Norway
Oslo Metropolitan University
Kristiansand, Norway
Oslo, Norway
1i.menge its ts. uio.no no
morten. www.xin@ui.i......
anisy@os dos no
Paal Engelstad
University of Oslo
Oslo, Norway
paal, real.enge engel stadlitattt staeeet vio.no
ABSTRACT
The manifold hypothesis posits that high-dimensional data often lies on a lower- dimensional manifold and that utilizing this manifold as the target space yields more efficient representations. While numerous traditional manifold-based tech- niques exist for dimensionality reduction, their application in self-supervised learning has witnessed slow progress. The recent MSimCLR method combines manifold encoding with SimCLR but requires extremely low target encoding di- mensions to outperform SimCLR, limiting its applicability. This paper introduces novel learning paradigm using an unbalanced atlas (UA), capable of surpassing state-of-the sfate-of-the-at -tthh-rrt self-supervised learning approaches We investigated and engi- neered the DeepInfomax with unbalanced atlas (DIM-UA) method by adapting the Spatiotemporal DeepInfomax (ST-DIM) framework to align with our proposed UA paradigm. The efficacy of DIM-UA is demonstrated through training and eval- uation on the Atari Annotated RAM Interface (AtariARI) benchmark, a modified version of the Atari 2600 framework that produces annotated image samples for representation learning. The UA paradigm improves existing algorithms signif- icantly as the number target encoding dimensions grows. For instance, the mean Fl score averaged over categories of DIM-UA is 75% compared to 70% ST-DIM when using 16384 hidden units.
INTRODUCTION
Self-supervised learning (SSL) is a field in machine learning (ML) that aims to learn useful feature representations from unlabelled input data. SSL includes mainly contrastive methods (Oord al., 2018 Chen et al. 2020: He et al.) (020) and generative models (Kingma & Welling [0133] Gregor al. Oh et al.J 2015). Generative models rely on using generative decoding and recon- struction loss, whereas typical contrastive methods do not involve decoder but apply contrastive similarity metrics to hidden embeddings instead (Liu al. 202T).
State representation learning (SRL) (Anand et al. 2019 Jonschkowski & Brock 2015] Lesort et al.)
2018) focuses on learning representations from data typically collected a reinforcement learning (RL) environment. A collection of images can be sampled through an agent interacting with the environment according to a specified behavior policy. Such images are interesting as study subjects due to their innate temporal/spatial correlations. Moreover, RL can also benefit from self- self-supervised learning just as computer vision (CV) and natural language processing (NLP) do, and successful pretraining of neural network (NN) models may lead improvements in downstream RL tasks.
manifold can be learned by finding an atlas that accurately describes the local structure in each chart (Pitelis et 2013). In SSL, using an atlas can be viewed generalization of both dimen-
Coode available tttps:////ittub.coooo wwwiithuubb.... httpp:///gttttb..om///enggg11//.

---

Published as a conference paper at ICLR 2024
sionality reduction and clustering (Korman 2018
2 b). Namely, it generalizes the case where only one chart exists and where the charts do not overlap in an atlas. In MSimCLR (Korman 2021b), NNs can encode an atlas of a manifold by having chart embeddings and membership probabilities.
Uniform prior
Unbalanced atlas
50
40
30
20
10
200
600
800
1030
Figure 1:
The entropy of the output vector recorded epoch-wise when pretrained on the CIFARIO dataset for total of 1000 epochs, utilizing 8 charts and dimensionality of 256.
One primary issue of MSimCLR is its reliance on uniform prior, which allocates inputs into each chart embedding uniformly. We postulate that although this uniform prior may more effectively represent the data distribution when d is exceedingly small, it concurrently introduces higher pre- diction uncertainty. Simultaneously, also suffers from a problem akin to that faced by bootstrapped methods in RL. It has been noted that multiple NN heads inside model, in the absence of addi- tional noise, tend to output similar results after being trained a large number of epochs (Osband & Van Roy 2015 Osband et al. 2016 Ecoffet et al. 2019 Meng et al. 2022).
To rectify the aforementioned problems, this study introduces a novel SSL paradigm that leverages an unbalanced atlas (UA). In this context, UA denotes the absence of uniform prior distribution, with the membership probability distribution deliberately trained to deviate significantly from uni- formity. As illustrated in Fig. [ it is evident that the entropy of the output vector during pretraining when using UA is markedly lower than that with a uniform prior, which suggests a heightened degree of confidence in its predictions.
Our contribution summarized as follows: (1) We modify the SRL algorithm ST-DIM (Anand al. 2019) with our UA paradigm and introduce new algorithm called DIM-UA. This furthers the research into the integration of RL and SSL with novel manifold-based learning paradigm. DIM-UA achieves the tfte-offtheeart performance on samples collected from 19 Atari games of the AtariARI benchmark. (2) We also provide detailed ablations and additional experiments on
CIFARIO to examine different underlying effects of possible design choices. (3) We demonstrate that our UA paradigm is capable of effectively representing a manifold with a large number (e.g., 2256) of hidden dimensions, whereas previous research (Korman) 202 102labb only showed promise with a small number (e.g.. <8) of hidden dimensions. The UA paradigm thereby showcases its capability to build larger models, transcending the constraints imposed by model backbones.
2 RELATED WORK
Dimensionality reduction with manifolds It is common for nonlinear dimensionality reduction (NLDR) algorithms to approach their goals based on the manifold hypothesis. For example, the manifold structure of an isometric embedding can be discovered by solving for eigenvectors of the matrix of graph distances (Tenenbaum et al. 2000). A sparse matrix can also be used instead with locally linear embedding (Roweis & Saul 2000 Correspondence between samples in different data sets can be recovered through the shared representations of the manifold (Ham et al.) (003). Manifold regularization provides out-of-sample extension compared to graph-based approaches
2

---

Published as a conference paper at ICLR 2024
[Belkin et al. 2006). Manifold sculpting simulates surface tension progressively in local neighbor- hoods to discover manifolds (Gashler et al. 2007).
Self-supervised learning There are relevant works on generative models, such variational au- toencoders (VAEs) (Kingma & Welling 2013) and adversarial autcencoders (AAEs) (Makhzani al. 2013). Meanwhile, contrastive methods have shown promise in the field of SSL. Contrastive Predictive Coding (CPC) learns predictive representations based on the usefulness of the information predicting future samples (Oord al., (018). SimCLR provides a simple yet effective framework
using data augmentations (Chen et al. 2020). Momentum Contrast (MoCo) utilizes dynamic dic- tionary, which can be much larger than the mini- mini-batch size (He ct al. 2020j. The recent trend within research in contrastive learning has been removing the need for negative pairs. BYOL utilizes momentum encoder to prevent the model from collapsing due to a lack of negative pairs (Grill et al, 2020). SimSiam further shows that stop-gradient operation alone is sufficient (Chen & He 2021 Barlow Twins, on the other hand, achieves so by minimizing the redundancy of vector components outputted by two identical networks that take distorted versions of inputs (Zbontar al. 2021
Supe-supe learning with manifolds Representing Non-Euclidean data in NN models is key topic geometric deep learning (Bronstein et al.] 2017). Learning manifolds using NNs was explored in (Korman 2018), in which AAES were used to learn atlas as latent parameters. Constant-curvatuon Riemannian manifolds (CCMs) of different curvatures can be learned similarly using AAEs (Grattarola et al. 2019). Mixture models of VAEs can be used to express the charts and their inverses to solve inverse problems (Alberti et al. 2023). combination of autoencoders and Barlow Twins capture both the linear and nonlinear solution manifolds (Kadeethum al., 2022).
3 METHOD
U
Z
Up
Wa
Vap
Vpp
Figure 2: A manifold embedded in a higher dimension. Two domains are denoted by U, and U/ in Z. th and tha are the corresponding charts that map them lower dimensional Euclidean space. An atlas is then a collection of these charts that together cover the entire manifold.
Our method extends the work from Korman (021b) and builds manifold representation using multiple output embeddings and membership probabilities of those embeddings. First, an illustration
of a manifold Z is in Fig. 2 Instead of directly learning encoder of Z, we can learn encoder functions of tho (Ua) and (Up) together with score function. For given input, we can specify which encoder
according to the output of the score function.
We formally model a distribution of input data by a manifold as follows: x is the input from input space X, 2 is the latent space, is embedding function: X' - Z. I is the identity mapping, d the number of dimensions for each chart output embedding, n the number of charts, and A denotes 21,2, n). by; Z Bd is the inverse mapping of coordinate map: Id + Z. whereas q
(q), ....... gn) Z - [0, 11" is the chart membership function. The output of our model is then
given by Eq.
3

---

Published as a conference paper at ICLR 2024
Output(z)= - (((
(I)
At inference time, the one-hot encoding of ((x)) is used instead (Eq. 2
Output(<> II(x)((()())))) where argmax, T5 (J(())))
(2)
3.1
UNBALANCED ATLAS
Like other SSL methods with manifolds (Korman bb2 a UA uses maximal mean discrepancy
(MMD) objective (Gretton et al., 2012 Tolstikhin al. 2017 7), which is defined by Eq.
(MD,A(P,... P2) - k(s, k(s..)dP((s)
(3)
Here, & is reproducing kernel, Hs is the reproducing kernel Hilbert space of real-valued functions mapping S to Ik, and P1. P2 are distributions on the space.
In paradigm, the input is designed to be represented in charts with higher membership prob- abilities. Thus, we take an MMD loss that moves the conditional membership distribution far away from the uniform distribution. We use the kernel ky A N R. (i. dij. and dig i else 0, and thus have Eq.
1 11
Cx(()) -M. (1(z), UN) -E Σ (qi())-
(4)
Here, UN denotes the uniform distribution on N, is the embedding of f(x).
Unlike MSimCLR, we do not use an MMD objective to make the prior distribution to be uniform, but take another approach to improve the model stability and head diversity when d is not trivial. In Fig. 2 Ua Ug has transitive maps their respective chart coordinates with domains restricted to tha (Ua n Uz) and (Ua n Ug). which are Va. = the o V/a and 19sa = the o V5 What interests us the most is this intersection between Ua and Ug. More precisely, since the overlapping representations in each head of the model become dominant negative factor when d grows larger, we aim at modelling a manifold with dilated prediction targets in pretraining to avoid convergent head embeddings and collapsing solutions. We use the average values of chart outputs model Minkowski sum (Mamatov & Nuritdinov 2020 Wang et al. 2020), which serves key purpose in our paradigm.
While convergent head embeddings should be avoided, the learning process should not break the convergence entirely. Proposition implies that the Minkowski sum the output embeddings con- tains the Minkowski sum of all mappings of intersections, which means that using dilated prediction targets by taking the Minkowski sum does not omit any mapped intersected embedding.
Un be a collection of open subsets of 2 whose union is all of and n Ug not empty. For each E 21, n), there homeomorphism 4/s U, -> V, d=1
Proposition 1. Let U =
n
to an open set Vi C RRd. We have the Minkowski sum V + Vj = {a + V..e Vy). Then
ET Vil M C,
Proof. For any vector a E E vial ( Un). there exists a, E vil n U,) such that a = E Gz- s=1
j=1
Because via( n A U,) c V,, we also have E V. i E (1,2, n). Then Σn aj E Σn Va
f=1
En Vi and thus Ens un En V..
4

---

Published as a conference paper at ICLR 2024
Proposition 2 further states that the average Minkowski sum of the output embeddings together with the Minkowski sum of the average mappings intersections can be used instead and still keeps Proposition 1 true, under the assumption that each mapping of the intersection is convex. More generally, Proposition holds true with scalar multiplications when convexity assumed. However, it should be noted that convexity is not guaranteed here. In Eq. [ we approach this assumption by using an identity mapping Z instead a linear mapping from Korman (021b), More about the convexity assumption is addressed in Appendix B
Proposition 2.
Let UU, Un be collection of open subsets of Z whose union is all of Z, and n Ug is not empty. For each E {1, n., there is homeomorphism t/; Uq - V; to an i=1
open set V, C Rd. The multiplication of set V and a scalar X is defined to be XV =\= a aVV).
We take the Minkowski sum. If each s/s{ n U₃) is convex, then E v n U,) C N Σ V,
j=1
Proof. Follows Proposition and the property of scalar multiplication, E wi n Uy) c
3=1
Vg. Since scalar multiplication is preserved for convex sets, we have Di
4/4 U₁)
Ve
===
3.2 DIM-UA
We experiment with UA paradigm using the SRL algorithm ST-DIM (Anand al. 2019). and propose DIM-UA. ST-DIM develops on Deep InfoMax (HiM) et al. 2018) that uses infoNCE (Oord et al. 2018) as the mutual information estimator between patches. Its objective consists of two components. One is the global-local objective (CGL) and the other one is the local-local objective (LLL), defined by Eq. and Eq. respecctivly
M
N
exp(sm (Fe, l(+1))
Las.
log Σ
(5)
Bin Bi.-EXnont exp(_mm, ( Fea))
m= n=1
exp(kmm, (2t, Ie+1)) Σn. Exnett exp(Immm (Ie, Ita))
CLL
log
Here, Iy and T++1 are temporally adjacent observations, whereas Xnext is the set of next observa- tions and 2is is randomly sampled from the minibatch. M and N are the height and width of local feature representations.
We denote the encoder as the output (global feature) vector of input x, Output(z;), and the local feature vector of I, at point (m, n.) as fm,n((1)) W, and Wa, are linear layers that will be discarded in probing. Then, we have the score function Im.n (Ie, I(+1) = Out W (Wff:r))"" of ST-DIM.
and
hm. (z,, T++1) = (
For DIM-UA, need to redefine the score function of Loi by Eq. because the UA paradigm utilizes dilated prediction targets during pretraining, where 1h (x:)) is the output of x, from the é-th head following encoder f for each in N.
9m, (x, (141) =
Σ
(7)
n
5

---

Published as a conference paper at ICLR 2024
According to Eq. 4 have the MMD objective Lg defined Eq. 8 membership probability the i-th head for each in N when the input is It-
where (;(((x))) is the
LQ 1 ((a(((()))) - -12² -
(a:([(zz41) - })
(8)
Thereby, the UA objective (Eq. 9 is sum of above objectives, where is hyper-parameter.
LUA LGL + CLL ThQ
(9)
EXPERIMENTAL DETAILS
The performance of DIM-UA and other SRL methods is evaluated on 19 games of the AtariARI benchmark. There are five categories of state variables in AtariARI (Anand et al. 2019). which are agent localization (Agent Loc.), small object localization (Small Loc.), other localization (Other Loc.), miscellaneous (Misc.), and score/clock/lives/dissp (Score/.../Dssplay)
We follow the customary SSL pipeline and record the probe accuracy and Fl scores on the down- stream linear probing tasks. The encoder is first pretrained with SSL,, and then is used to predict the ground truth of an image with an additional linear classifier. Notably, the weights of the encoder trained only during the pretraining and are fixed in the probing tasks. The data for pretraining and probing are collected by an RL agent running a certain number of steps using a random policy since it was found that the samples collected by random policy could be more favorable than those collected by policy gradient policies for SSL methods (Anand et al. 2019).
Previous SSL methods in Anand et al. (2019) have used single output head with 256 hidden units. One of the major interests in our experiment is to discover the effect of choosing different values for the number of dimensions d and the number of charts n. Therefore, of hidden units, while keeping the model architecture, to observe the performance using a single output head without UA and of using multiple heads with UA. To make fair comparison, compare the performance when the total number of hidden units are equal, i.e., when x for single output head and n d for multiple output heads are equal in our AtariARI experiment. In contrast, we also modify SimCLR using our UA paradigm and follow the customs from MSimCLR to compare the performance of different methods with d being equal (Krrman) 202Tb) additional experiments CIFAR10.
scale up the number
The experiments are conducted on single Nvidia GeForce RTX 2080 Ti and 8-core CPU, using PyTorch-1.7 (Paszke al., 2019). An illustration of the model backbone, hyper-parameters, and pseudocode of the algorithm are accompanied in Appendix
5 RESULTS
In this section, we show the empirical results of our experiments and compare DIM-UA with other SSL methods verify the efficacy of our UA paradigm. Mcanwhile, we pay special attention the performance of models when choosing different values for and d.
For straightforward comparison, we first observe the probe Fl scores together with standard devi- ations of each game averaged across categories in Table [ "ST-DIM*" denotes ST-DIM with one output head of 16384 units, while DIM- UA uses four output heads with 4096 units in each head. We compare them to various methods of using a single output head of 256 hidden units here, which are taken from Anand et al. (2019). Each table entry an average of independent pretraining/probing runs using images sampled from different seeds. The probe accuracy scores also included in Appendix
Using 16384 units ("ST-DIM*") does not necessarily guarantee better performance than using 256 units (ST-DIM). In out of 19 games, "ST-DIM"" has lower scores than ST-DIM. In particular, the model collapses due overfitting when using too many units to represent the global features Freeway. As a result, "ST-DIM"" only gains an Fl score of 0.3 with standard deviation 0.355 on Freeway. The mean F1 score of "ST-DIM"" is only 0.7 compared 0.72 of ST-DIM. On the other

---

Published as a conference paper at ICLR 2024
Table 1: Probe F1 scores of each game averaged across categories
Game
VAE
CPC
ST-DIM
ST-DIM*
DIM-UA
Asteroids
0.36
0.42
0.49
0.48 + 0.005
0.5 = 0.007
Bowling
0.50
0.90
0.96
0.96 0.021
0.96 - 0.018
Boxing
0.20
0.29
0.58
0.61 = 0.008
0.64 0.007
Breakout
0.57
0.74
0.88
0.88 0.02
Demon Attack
0.26
0.57
0.69
0.71 + 0.01
0.74 0.012
Freeway
0.01
0.47
0.81
0.355
0.86 + 0.02
Frostbite
0.51
0.76
0.75
0.73
0.75 = 0.004
Hero
0.69
0.90
0.93
0
0.94 = 0.004
Montezuma Revenge
0.38
0.75
0.78
0.84 = 0.014
Ms Pacman
0.56
0.65
0.72
0.74
0.76 0.011
Pitfall
0.35
0,46
0.60
0.69 + 0.031
0.73 + 0.029
Pong
0.09
0.71
0.81
0.78 1 0.015
0.85 0.004
Private Eye
0.71
0.81
0.91
0.91 + 0.009
0.93 + 0.009
Qbert
0.49
0.65
0.73
0.026
0.79 = 0.02
Seaquest
0.56
0.66
0.67
0.69 = 0.007
Space Invaders
0.52
0.54
0.57
0.59
0.62 0.013
Tennis
0.29
0.60
0.60
0.64 0.025
Venture
0.38
0.51
0.58
0.57 0.014
0.58 + 0.01
Video Pinball
0.45
0.58
0.61
0.6 + 0.031
0.62 0.023
Mean
0.41
0.63
0.72
0.7 1 0.033
0.033
0.75 1 0.013
ST-DIM
DIM-UA
0.74
0.75
Soore
Accuracy
0.72
0.79
0.7
0.71
0.68
0.89
512
1024
2048
4086
8192
16984
512
1024
2048
4086
8192
1638A
Unit
Unt
Figure 3: The mean Fl and accuracy scores of 19 games when the total number of hidden units varies. The number of heads for DIM-UA is set here.
hand, DIM-UA achieves higher scores and more stable performance. The F1 scores of DIM-UA equal or higher than those of both "ST-DIM*" and ST-DIM in every game. The mean Fl score of DIM-UA is 0.75, the highest among all methods in Table
5.1 ABLATIONS
In this subsection, our goal to observe the behavior of models with different settings while the total number of units (n x d) on the horizontal axis of figures varies. Fig. eexamines the difference between ST-DIM and DIM-UA. Fig. 4 examines the effects changing the number of heads for DIM-UA. In addition, method uses the paradigm from MSimCLR, denoted by "+MMD". The other is similar to ours, which minimizes the loss of Eq. 9 but without modifying the score function as in Eq. (namely, DIM-UA without using dilated prediction targets), denoted by "-UA". The 6 games mentioned here Asteroids, Breakout, Montezuma Revenge, Private Eye, Seaquest, and Video Pinball.
compare DIM-UA with two methods designed on ST-DIM in Fig, 5 One

---

Published as a conference paper at ICLR 2024
hoads
heads
heads
0.75
0.78
0.723
0.735
F1 Score
Accuracy
0.895
0.71
0.668
0.685
0.84
0.88
512
1024
2048
4008
8102
16384
512
1024
2048
4030
6182
16394
Unit
Figure 4: The mean F1 and accuracy scores of DIM-UA on 6 games when the number of output heads 4, or
0.75
0.76
0.73
0.74
F1 5oore
Accurscy
0.71
0.72
0.09
0.67
0.68
512
1024
2048
4096
8192
16384
512
1024
2048
4086
8192
10384
Unit
Uhe
Figure 5: The mean F1 and accuracy scores on 6 games with different adaptations. All methods use 4 output heads.
In Fig. 3 ST-DIM performs better than DIM-UA when the number of hidden units is small. Their scores become close to each other when the number units is around 2048. DIM-UA continues to improve as the total number of units grows, whereas the performance of ST-DIM drops at the same time. is expected for DIM-UA to have lower Fl and accuracy scores when the encoding dimensions are low since the diversity among output heads demands more epochs of training. How- ever, the efficacy of our UA paradigm is clearly demonstrated, as it allows the model to extend its capability by continuously expanding the encoding dimensions.
Since we expect DIM-UA to converge slower because of the diversity among output heads, we expect this become more obvious as the number of heads increases. We can verify that in Fig. where the model with two output heads has the highest F1 and accuracy scores when the total number of hidden units is below 2048. On the other hand, it obtains the lowest Fl and accuracy score when the total number of units grows to 16384. Meanwhile, the model with eight output heads gives the worst results when the number of units small but shows no sign of plateau, even with very high encoding dimensions. Increasing while keeping d the same in our UA paradigm helps with the manifold representation but also lowers the performance if dl not large enough.
In Fig, 5 it is not surprising that "+MMD" obtains the worst results in spite of the number of units, since MSimCLR was only found out to be helpful when the number of units is extremely small (e.g.,
2, 4). "-UA" obtains better results than DIM-UA UA when the number of units is 512 but gets overrun by DIM-UA when the number becomes even larger. This empirically demonstrates that the dilated prediction targets in our UA paradigm are critical to achieve effective manifold representations.
8

---

Published as a conference paper at ICLR 2024
Table 2: Linear evaluation accuracy on CIFARIO
Method
Head
Dimension
256
512
1024
SimCLR
0.881
0.002
0.883 + 0.002
0.881 = 0.003
MSimCLR
2
0.877 + 0.002
0.878
0.001
0.866 = 0.003
MSimCLR
4
0.873 + 0.001
0.873 = 0.001
0.861 = 0.002
MSimCLR
8
0.864 0.001
0.859
0.005
0.857 + 0.002
SimCLR-U UA 2
0.882
0.001
0.884
0.001
0.885 + 0.001
SimCLR-U
4
0,885 + 0.001
0.884
1 <0.001
0.88
+ 0.001
SimCLR-UA
8
0.882 1 <0.001
0.886
+ 0.002
0.876 + 0.005
5.2 ADDITIONAL. EXPERIMENTS ON CIFAR10
We modify SimCLR using the UA paradigm (SimCLR-AA and perform additional experiments on CIFARIO, following the parameter settings and evaluation protocol from Korman (202Tb); Chen et al. (2020). SimCLR-UA uses multiple heads with dilated prediction targets instead in pretraining, and adds Lg in Eq. to the contrastive loss of SimCLR. Here, ResNet50 is used as the backbone, which is significantly larger than the default backbone in ST-DIM. We also slightly modify our model here based on an ablation study on CIFARIO. Please check Appendix [ for more details.
In Table each entry is average accuracy score obtained from three independent pretraining and evaluation runs, together with the standard deviation. SimCLR obtains an accuracy score of 88.3% when it uses 512 hidden units. In the mean time, SimCLR-UA achieves the highest accuracy score of 88.6% among three methods when uses eight heads with 512 units in each. The second best score is also achieved by SimCLR-UA. which is 88.5% when using four heads with 256 units in each head, or using two heads with 1024 units in each. We acknowledge that our improvement over SimCLR is small under this experimental setup. Nonetheless, there is significant increase in accuracy when comparing SimCLR-UA to MSimCLR, especially in the case where the number of heads is larger. For instance, the highest evaluation accuracy score of MSimCLR is 87.8%, 87.3% and 86.4% respectively, when using two, four or eight heads. In contrast, SimCLR-U obtains the highest accuracy score when using eight heads. This supports our hypothesis that UA can be universal paradigm to effectively create manifold representations in SSL.
6 DISCUSSION
We have demonstrated that our UA paradigm helps improve the performance of both ST-DIM and SimCLR when encoding dimensions are high. Furthermore, we argue that training NNs with multi- ple output heads is inherently slower and more demanding than training with single output head, which has restrained the study in its domain. It is evident that our paradigm can overcome this headwind by generating effective manifold representations. Moreover, UA paradigm has gained a significant amount of improvement compared to the most related stateoof-th-aan art manifold repre- sentation paradigm MSimCLR in the experiments on AtariARI and CIFARIO.
Notably, the UA paradigm also exhibits the potential of modeling a manifold using further higher dimensions while increasing the number of output heads (Fig. 4.. It can be an important contribu- tion because this means the performance of the model scales with the size of output heads. Using 16384 hidden units in total not very efficient economically when the entire model is small, but the additional overhead introduced by doing so can be relatively insignificant when the model itself is large. particular, this trade-off may also be worthwhile in challenging downstream tasks where the smallest increase in probe accuracy can make difference.
Our work has illustrated that SSL methods with manifolds have great potential, and more topics can be researched in this area. The relationship between the number of hidden units and the number of output heads in an NN model demands more study (see Appendix B for more discussion on this). The convexity assumption is crucial in representing the manifold. Future research may focus representing a manifold using an unbalanced atlas more efficiently, e.g., designing new objectives and convexity constraints.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was performed on the [ML node] resource, owned by the University of Oslo, and operated by the Department for Research Computing at USIT, the University of Oslo IT-department.
REFERENCES
Giovanni S Alberti, Johannes Hertrich, Matteo Santacesaria, and Silvia Sciutto. Manifold learning by mixture models of vaes for inverse problems. arXiv preprint arXiv:2303. 2023.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. Advances neural information pro- cessing systems, 32, 2019.
Mikhail Belkin, Partha Niyogi, and Vikas Sindhwani. Manifold regularization: A geometric frame- work for learning from labeled and unlabeled examples. Journal of machine learning research, (11), 2006.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Piere Vandergheynst. Geomet- ric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18-42, 2017.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning. pp. 1597-1607. PMIR, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 15750-15758, 2021.
Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O Stanley, and Jeff Clune. Go-explore: new approach for hard-exploration problems. arXiv preprint arXiv:1901.1099 2019.
Michael Gashler, Dan Ventura, and Tony Martinez. Iterative non-linear dimensionality reduction with manifold sculpting. Advances neural information processing systems, 20, 2007.
Daniele Grattarola, Lorenzo Livi, and Cesare Alippi. Adversarial autoencoders with constant-
curvature latent manifolds. Applied Soft Computing, 81:105511, 2019.
Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Rezende, and Daan Wierstra. Draw: A recurrent neural network for image generation. In International conference on machine learning, pp. 1462- 1471. PMLR, 2015.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach self-supervised learning. Advances in neural information processing systems, 33:21271-21288. 2020.
Hun Ham, Daniel D Lee, and Lawrence K Saul. Learning high dimensional correspondences from low dimensional manifolds. International Conference Machine Learning Workshop, 2003.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9773 2020.
Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint arXiv:1808 0,,66600,, 2018.
Devon Hjelm, Alex Fedorov, Samuel
Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Au- tonomous Robots, 39(3):407-428, 2015.
10

---

Published as a conference paper at ICLR 2024
Teeratorn Kadcethum, Francesco Ballarin, Daniel O' O'malley, Youngsoo Choi, Nikolaos Bouklas, and Hongkyu Yoon. Reduced order modeling for flow and transport problems with barlow twins self-supervised learning. Scientific Reports, 12(1):20654, 2022.
Diederik Kingma and Max Welling.
Auto-encoding variational bayes.
arXiv preprint
arXiv:1312.61 14, 2013.
Eric O Korman. Autoencoding topology. arXiv preprint arXiv: 2018.
arXiv preprint Eric O Korman. Self-supervised representation learning on manifolds. In ICLR 2021 Workshop on Geometrical and Topological Representation Learning, 2021b.
Eric O Korman. Atlas based representation and metric learning on manifolds. arXiv:2106.07062 2021a.
Timothée Lesort, Natalia Diaz-Rodríguez, Jean-Franois Goudou, and David Filliat. State represen- tation learning for control: An overview. Neural Networks, 108:379-392, 2018.
Xiao Liu, Fanjin Zhang, Zhenyu Hou, Li Mian, Zhaoyu Wang. Jing Zhang, and Jie Tang. Self- supervised learning: Generative or contrastive. IEEE Transactions on Knowledge and Data En- gineering, 35(1):857-877 2021.
Alireza Makhzani. Jonathon Shlens, Navdeep Jaitly, lan Goodfellow, and Brendan Frey, Adversarial autoencoders. arXiv preprint arXiv:151 2015.
Mashrabjon Mamatov and Jalolxon Nuritdinov. Some properties the sum and geometric differ- ences of minkowski. Journal of Applied Mathematics and Physics, 8(10):2241-222 2020.
Li Meng, Morten Goodwin, Anis Yazidi, and Paal Engelstad. Improving the diversity of boot- strapped dqn by replacing priors with noise. IEEE Transactions Games, 2022.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard L Lewis, and Satinder Singh. video prediction using deep networks in atari games. Advances in neural information processing systems, 28, 2015.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint 1r7.v:1807.03748, 2018
Ian Osband and Benjamin Van Roy. Bootstrapped thompson sampling and deep exploration. arXiv preprint /507.vv:500..00 00300, 2015.
Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29:4026 2934026-0034, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performanee deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Nikolaos Pitelis, Chris Russell, and Lourdes Agapito. Learning a manifold as an atlas. In Pro- ceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 1642-1649, 2013.
Sam T Roweis and Lawrence K Saul. Nonlinear dimensionality reduction by locally linear embed- ding. science, 290(5500):2333-23333 2000.
Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science, 29005500):2319-2333, 2000.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. arXiv preprint 1rXiv: 2017.

---

Published as a conference paper at ICLR 2024
Xiangfeng Wang, Junping Zhang, and Wenxing Zhang. The distance between convex sets with minkowski sum structure: application to collision detection. Computational Optimization and Applications, 77:465-490, 2020.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning, pp. 12310- 12320. PMLR, 2021.
A DETAILS OF DIM-UA
Table 3 provides values of some crucial hyper-parameters experimented on AtariARI that are kept the same across different methods. In addition, Eq. is set to 0.1 for DIM-UA.
Fig. 6 illustrates the standard backbone in the AtariARI experiment. The output values of the last convolutional layer in the backbone are taken as the feature map to get the local feature vector fm.n (2x) of input 2i at location (m,,1) For the original ST-DIM, fully connected layer of 256 units immediately follows the backbone. For DIM-UA, a projection head (1)) and membership probability head (g) branch from the backbone.
Table 3: The values of hyper-parameters on AtariARI
Value
Image size
160 x 210
Minibatch size
64
Learning rate
3e-4
Epochs
100
Pretraining steps
80000
Probe training steps
35000
Probe testing steps
10000
flattened, 3456 units
ReLU kernel 3, stride 64 channels
ReLU
kernel 4, stride 2, 128 channels
ReLU
kernel 4, stride 2, 64 channels
ReLU
kernel 8, stride 4, 32 channels
Atari image 160x210
Figure 6: An illustration of the standard backbone used by ST-DIM.
Pytorch-style pseudocode of the DIM-UA algorithm is provided in Algorithm I
On a side note, the output of an unbalanced atlas at inference time relies on a single output head, since Eq. 4 moves the membership probability far away from the uniform distribution. As result, the rest of the output heads does not play role at inference time. This different from MSimCLR, which partitions inputs into each head by simultaneously forcing uniform prior and low entropy conditional distributions. The role of those remaining output heads in an unbalanced atlas is compa- rable to the moving average network of BYOL, which produces prediction targets to help stabilize
12

---

Published as a conference paper at ICLR 2024
Algorithm 1: Pytorch-style pseudocode for DIM-UA.
encoder network
fm: I but only up to the last conv layer
pj: projection head
mp: membership probability head
cl, c2: classifier layers only used in pretraining
B: batch size
N: number of heads
D: number of hidden units
C: local feature map channels
Н, W: local feature map height and width
mean: mean function along a specified dimension
matmul: matrix multiplication cross_entropy: cross entropy loss mmd: mmd loss
size: get the size along specified dimension
range: get the range vector of an integer
t: transpose
for St, It-1 in loader: # load samples of St, F++1
get feature maps
On Mt, Mt+1 mean(pi(((z.)), 1), (x4), (=(41) # BXD, BXHXWXC get the membership probabilities
Mt 9+1 mp(f(xz)). mp(f(x++1)) BXN get the feature map size
Sb, 8m, 8, = size(y,, 0), size(y,, 1), size 2) B, H, W initialize the loss va values
loss, 0, loss: = 0
mnd loss
loss = -0.05 * (mmd(q.) + mmd(4+1))
spatial-temporal
loss
for m in range(sm):
for n in range(s..::
global-locae loss
logits, mat) Mtt1l m, n, :l.t()) BXB
loss, + = cross entropy (logits.. range(s:)) $ cross entropy loss
local ocal loss
logits: = matmul(c2(xaf: m, n, :)), M+t1l: n, :).t()) # BXB
loss; + cross_entrus enssentrropylloittaa range(s:)) # cross entropy loss
lossg = (am * sn), loss: / (&m 8x) # averaged by map size
loss + lossy + loss; # total loss
opt optimization step
loss. backwad))))d)) optimizer.step()
the bootstrap step. The UA paradigm accomplishes similar goal by using dilated prediction targets of output heads instead.
The probe accuracy scores shown in Table 4 which overall similar to the Fl scores in Table The accuracy scores of DIM-UA are equal or higher than those of both "ST-DIM and ST-DIM every game. The mean accuracy of DIM-UA is 0.76, the highest among all methods.
13

---

Published as a conference paper at ICLR 2024
Table 4: Probe accuracy scores of each game averaged across categories
Game
VAE
CPC
ST-DIM
ST-DIM*
DIM-UA
Asteroids
0.41
0.48
0.52
0.51 0.005
0.53 0.007
Bowling
0.56
0.90
0.96
0.96 0.021
0.96 = 0.017
Boxing
0.23
0.32
0.59
0
0.64 = 0.007
Breakout
0.61
0.75
0.89
0.89
0.91 = 0.015
Demon Attack
0.31
0.58
0.70
0
0.74 + 0.011
Freeway
0.07
0.49
0.82
0.33 0.34
0.86 0.017
Frostbite
0.54
0.76
0.75
0.704
0.75 = 0,004
Hero
0.72
0.90
0.93
0
0.94 0.004
Montezuma Revenge
0,41
0.76
0.78
081+000
0.84 - 0.014
Ms Pacman
0.60
0.67
0.73
0.75 0.016
0.77 0.01
Pitfall
0.35
0.49
0.61
0.7
0.74 + 0.027
Pong
0.19
0.73
0.82
0
0.85 0.004
Private Eye
0.72
0.81
0.91
0.91 :: 0.01
0.93 + 0.009
Qbert
0.53
0.66
0.74
0.025
0.8 0.019
Seaquest
0.61
0.69
0.69
0.006
0.7 = 0.006
Space Invaders
0.57
0.57
0.59
0.009
0.63 0.014
Tennis
0.37
0.61
0.61
0.58 0.066
0.65 0.024
Venture
0,43
0.52
0.59
0.58 0.012
0.59 + 0,009
Video Pinball
0.47
0.59
0.61
0.6
0.63 0.022
Mean
0.46
0.65
0.73
0.71 + 0.031
0.76 + 0.013
B MORE ABOUT CIFAR10 EXPERIMENT
The convexity assumption crucial in order effectively model a manifold (c.g., scalar multipli- cation is not preserved when set is non-convex). However, the universal approximation theorem implies that weights of multiple linear layers can approximate any non-convex functions. Thus, whether multiple linear layers should be used here or not could be an interesting ablation topic. Moreover, clamping can be introduced to define open sets our method. An ablation study of SimCLR-UA is performed on CIFARIO, using 4 heads with 512 units in each head. The results are shown in Table 5 where "FC1" denotes the linear layer immediately following the ResNet50 backbone and "FC2" denotes the projection layers following coordinate mappings. The range of clamping is set to 10, 10). As the results suggest, the combination of clamping and "FC2" yields the best accuracy and is hence used to obtain results in Table 2
Referring to the performance of SimCLR-UA in Table the accuracy reaches the highest when the number of heads eight with 512 units in each head, but when the number of heads is four, the optimal number of units is 256. It also appears that a small number of hidden units can be sufficient. This finding is different from what observed in the AtariARI experiment. where using eight heads and 2048 units in each head is not sufficient to guarantee the optimal (Fig. 4 This may be attributed the image size and the number of ground truth labels in an image, and more challenging tasks may demand better representations. However, we do think there should be limited number of heads needed, related to the intrinsic dimension (ID) of data. Techniques to find ID can potentially be used to decide the optimal number of heads.
Table Ablation study on CIFARI0
FCI
Clamp
FC2
Accuracy
0.857
0.876
0.883
0.872
0.875
0.884
0.872
14

---

Published as a conference paper at ICLR 2024
After observing the results in Table it appears that the performance of SimCLR-UA could be further enhanced. Whilst comparing MSimCLR with SimCLR- UA, we maintained most hyper- parameters identical for both. However, SimCLR-UA may attain higher performance given differ- ent set of aapee-parametens using uniform prior. Ideally, it may be advantageous for the entropy remain high during the initial stages and to decrease gradually. We suggest that the hyper-parameter т, which regulates the Lo loss, could be set smaller or set to 0 initially and gradually increased time.
Fig. D the initial entropy of using UA is notably lower than when
We conduct an additional small-scale experiment to validate this hypothesis. In this context, we use R the backbone instead of ResNet50, and the training duration set 100 epochs as opposed to 1000. The model incorporates 8 heads, with each head containing 512 hidden units. If is linearly scaled, would increment linearly, from zero up to its final value over the pretraining epochs. The outcome of this experiment is detailed in Table 6
The table clearly illustrates that implementing linear scaling or utilizing smaller values can gen- uinely enhance the performance of SimCLR-UA. For values of 0.1 and 0.2, adopting linear- scaling scheme is instrumental for optimizing the performance. However, for small T values of 0.05 and 0.02, such scheme is not needed. Thus, the performance of SimCLR-UA, as presented in Table [ could potentially be boosted further, since uses relatively large 7 value of 0.1 without any linear scaling.
Table 6: Changing 7 in SimCLR-UA
Linear scaling
Accuracy
0.2
0.791
0.2
0.797
0.1
0.791
0.1
0.8
0.05
0.799
0.05
0.796
0.02
0.802
0.02
0.785
15


Reviews:
[{"review_id": "6FAj3Iv5tn", "reviewer": "Reviewer_hci6", "summary": "This paper developed a state representation learning method leveraging an unbalanced atlas (UA). The authors have modified the ST-DIM algorithm to align with the proposed UA paradigm. Although the main contribution is not stated intuitively, empirical evaluations on 19 games of the AtariARI benchmark suggested an improved performance compared with three established baseline methods (many existing self-supervised methods are omitted for comparison). Furthermore, the authors performed a comprehensive ablation study for the design choices of the proposed method.", "questions": "1. 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.\n\n2. 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.\n\n3. In the 3rd paragraph of Section 4, *d* and *n* are used without proper definition.\n\n4. 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).\n\n5. 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.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "+ The experiments are conducted across 19 games of the AtariARI benchmark, covering a variety of vision tasks.\n\n+ There are comprehensive ablation studies for the technical components of the proposed method.", "weaknesses": "- The clarity of the introduction could be enhanced by providing a more explicit context for the specialized terminology introduced (see Q1-3).\n\n- The comparison would benefit from the inclusion of key baseline models which are currently absent (see Q4).\n\n- 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.\n\n- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bd2tvXazp2", "reviewer": "Reviewer_ES5Q", "summary": "The paper addresses the problem of learning a low-dimensional manifold from high-dimensional data in the context of state representation learning. A new approach based on self-supervised learning is proposed in order to learn an unbalanced atlas representation. The proposed approach is called DeepInfomax with an unbalanced atlas (DIM-UA) and is evaluated on 19 Atari games of the AtariARI benchmark. The evaluation shows good performance in this benchmark.", "questions": "- What is the motivation for redefining the score function L_GL? \n\n- Equation 9 has an hyperparameter for L_Q but not for L_GL or L_LL. Why?\n\n- The improvements reported in table 2 seem relatively small. Can you comment on this further?\n\n- Why is the linear evaluation accuracy on CIFAR10 a suitable evaluation metric? The improvements shown here (table 3) are very marginal at best.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper addresses an important topic in the area of state representation learning. The proposed algorithm seems novel and yields good performance in the tested benchmarks.", "weaknesses": "I am struggling to understand the rationale behind the DIM-UA algorithm. What is the motivation for redefining the score function L_GL? Equation 9 has a hyperparameter for L_Q but not for L_GL or L_LL. Why? The writing of the paper is not clear in several points as some parts are difficult to follow. For example, Figure 1 is not quite clear to me and the caption is not very explicit. Additionally, the results could be presented in a more concise fashion (esp. Table 1 and 2 - showing the best results in bold would increase the readability).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H0a99LJDab", "reviewer": "Reviewer_Ww1x", "summary": "Authors propose to use multiple heads at the end of an encoder for\ncontrastive learning, instead of one. These heads are considered to\nmodel different charts in an atlas, mapping the data manifold to the\nembedding space. For a given sample, a score is computed to determine,\nprobabilistically, which chart should be used to encode\nit. A theoretical discussion is presented to support outputting a\nweighted average of charts based on scoring function, which relies on\nMinkowski sum of open sets to which charts map. The scoring function\nis forced to be different than a uniform distribution, thus the\nunbalanced nature of the mapping. Experiments are conducted to\nunderstand whether using multiple heads and weighted averaging at the\noutput leads to better representations than using a single head with\nthe same number of latent dimensions.", "questions": "1. can authors discuss further why the CIFAR experiments do not show a\n   similar improvement?\n2. can authors please improve the notation consistency in the\n   presentation?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. More expressive models for contrastive learning are very relevant\n   and interesting for the community. Here authors point out that when\n   the embedding dimension is very high, naive contrastive learning\n   may not use the embedding space very well. Instead of using a\n   single projection head, authors proposal to use multiple projection\n   heads seem to lead to better results according to Tables 1\n   and 2. This is a simple yet - seems to be - an effective\n   modification.  \n2. Results of the ablation study shown in Figure 2 are very\n   convincing. This simple approaches surely uses the dimensions much\n   more efficiently, and provides the expected gains in accuracy.\n3. The difference between +MMD and the proposed version, which I\n   assume is -MMD, is striking.\n4. This reviewer appreciates the experiments with CIFAR.", "weaknesses": "1. Technical contribution is not at a very high level, but the\n   contribution is focused and pertinent.\n2. In the ablation study, the model \"-UA\" is not clearly specified. If\n   authors do not use the modifications of 7, that means the -MMD loss\n   is also void. What does that yield? Are authors using a single\n   projection head in this case?\n3. CIFAR experiments show that the gains are much lower in these\n   experiments compared to the ones obtained in ATARI games. A\n   discussion towards this end is not provided but it would be very\n   valuable for the readers.\n4. Notation in the presentation of the method seems a bit\n   inconsistent. I recommend authors to improve the consistency in the\n   notation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)