Paper: Self-Guided Masked Autoencoders for Domain-Agnostic Self-Supervised Learning (ICLR-2024)

Reviewer comment (Presentation):
I am missing an end-2-end sketch of the proposed architecture. I understand the goal is to keep it somewhat agnostic to task, architecture (self/cross-attention etc), but essential components can still be sketched. Algorithm 1 (which appears by the way too early in the text) to some extent covers this, but not totally.

Author response:
For completeness, we have included an additional full visualization of the architecture in our supplementary materials (Fig 3). We hope this clarifies any remaining questions regarding the end-to-end structure.

Reviewer comment (Presentation):
Am I right to assume (3) - (5) are a mathematical notation of Fig. 1? If so, this is one of those examples where math becomes more of a problem than a solution!

Author response:
While we agree that the figure is easier to understand as it conveys the core mechanism in a more intuitive way, we believe that including a formal definition is beneficial for completeness and clarity.

Reviewer comment (Experiments):
The improvement SMA offers in settings without strong priors (such as protein property prediction, and chemical property prediction) is marginal, which is fine. However, here a simple random masking pattern seems to work pretty well, and that can be an efficient natural choice which is difficult to argue against. (See Tables 1, 2, 3 and 5).

Author response:
Our method improves over the random masking method in every benchmark we have tested without any drawbacks, as the per-step training time is the same. Therefore we believe our method is comprehensively better than the random masking method. Additionally, the amount of improvement over random masking is closely tied to how much redundancy exists in the domain itself, with protein biology likely having the least redundancy in discrete protein tokens.

Reviewer comment (Experiments):
I think it would be useful to have some reflections on the computational complexity (time and/or model space) of SMA. Is it the reason why ImageNet-100 is chosen instead of ImageNet itself?

Author response:
The main purpose of the ImageNet100 experiments is to provide a comparison to hand-designed masks in a controlled setting and not to suggest our method is state-of-the-art for image pretraining. We do acknowledge that the Perceiver architecture (our method does not actually introduce any significant additional compute) on raw images still has a high computational cost. Compared to the original perceiver, our architecture requires approximately 1/10th of the FLOPs, but scaling to the full ImageNet dataset is beyond our computational resources. In fact, the Imagenet100 experiments still incurred the most computational cost out of all domains.

Reviewer comment (Writing):
The paper can benefit from another proof read; e.g. here are some minor suggestions:  
a) On page 2: These architectures demonstrates … => Demonstrate  
b) On page 3: Formally for an unlabeled … => Formally,  
On Page 4: the inputs of a single "group," => “group”,  
On Page 9: also use resize image … ? Doesn’t read well.  
And several others.

Author response:
Thank you for the thorough editing suggestions. We have edited the paper to take all of these into account.

Reviewer comment (Experiments):
One can argue masking highlight correlated areas on the input actually poses an easier job to the prediction model no?

Author response:
Empirically we find our method to have a higher reconstruction loss compared to random masking despite lower masking rates across modalities indicating that our method generates a more challenging masking objective.

Reviewer comment (Theory):
Why would repeated top-k(.) operations be too complex or hard to parallelize for different tokens? And why would (3) help to elevate this bottleneck?

Author response:
The repeated top-k operation is challenging to parallelize because we must ensure each query masks unique tokens meaning the top-k operations need to be done sequentially to exclude the indices that were masked by other queries.

Regarding the complexity, suppose we select m queries from which we compute k inputs to mask for each from a total of n inputs. Without our approximation, the average-case runtime complexity is O(m(n + k)) as we must perform m top-k operations. In comparison, with our method, the average case complexity is only O(n + mk) because we just need to perform a single top-k operation with mk selected values. Since m, the number of queries selected is generally a fraction of n, and it is always the case that mk < n (we cannot mask more than or equal to 100% of inputs), our approximation reduces the complexity of top-k costs from squared to linear with respect to n (the number of inputs).

Empirically, the running time without our proposed approximation on the image domain is 1.13 steps/second. In comparison, the approximation training runs at 3.58 steps/second. Note that with the approximation, the time to compute masks does not significantly impact train step time (less than 2% of the time per step), but without the approximation, the time to compute masks becomes the bottleneck for large input counts. Additionally, empirically we found no difference in performance between the two methods across domains, and our approximation actually improves stability in our testing.

Reviewer comment (Experiments):
In Table 1, why shouldn’t one pick random masking? It seems to function as good as SMA, and the prior art, no? It’s way more efficient as well.

Author response:
Random masking is actually not more efficient by a noticeable amount, as even in the image modality where the difference in runtime is greatest, our algorithm is only 1% slower in training steps/second (3.63 steps/second vs 3.58 steps/second). Secondly, our method surpasses the random masking performance across all benchmarks. We suspect the difference is relatively smaller in the protein biology domain as the genomics data is somewhat preprocessed from its raw form by the benchmark.

Reviewer comment (Experiments):
While the authors report the results of ImageNet100, the results on the full dataset are also expected to ensure a comprehensive evaluation. Additionally, I'm also curious about the pre-trained encoder's performance on segmentation and object detection tasks.

Author response:
> While the authors report the results of ImageNet100, the results on the full dataset are also expected to ensure a comprehensive evaluation.

Unfortunately, scaling to the full ImageNet dataset is beyond our computational resources. In fact, the Imagenet100 experiments still incurred the most computational cost out of all domains due to the high number of input tokens per image (36,864). We even downscaled the original Perceiver architecture to approximately 1/10th the size in terms of FLOPs.

>Additionally, I'm also curious about the pre-trained encoder's performance on segmentation and object detection tasks.

Object detection is challenging given the time constraints since we are not aware of any prior works using Perceiver for object detection and testing and tuning a new architecture is unfortunately not possible during the response period. In the case of semantic segmentation, we are aware of a prior Perceiver work [1] which demonstrate semantic segmentation results on PASCAL VOC. However, their architecture used a resolution of 512 x 512 whereas ours uses one of size 192 x 192. Additionally theirs is much larger in terms of number of latents and hidden state dimension and is pretrained over ImageNet-1k. Therefore it would be challenging to pretrain this new larger architecture during the response period. Lastly, the purpose of the image domain experiments is not to suggest our pretrained model is state-of-the-art, but rather to just compare our guided masking method against a domain-specific masking method (patch masking) which we outperform.

Reviewer comment (Experiments):
For tabular datasets like HIGGS, the results are promising. However, I'd suggest authors extend the work to broader tabular datasets as the performance of deep learning-based models may vary a lot. Additional experiments are not super expensive in this domain but will give a more comprehensive evaluation.

Author response:
> For tabular datasets like HIGGS, the results are promising. However, I'd suggest authors extend the work to broader tabular datasets as the performance of deep learning-based models may vary a lot. Additional experiments are not super expensive in this domain but will give a more comprehensive evaluation.

Previously we used the HIGGS tabular dataset as we were focused on scientific applications (in line with the chemistry and biology domains we used). However, we agree it would make the evaluation more comprehensive with more tabular results. We have benchmarked 1 more tabular dataset, the year prediction regression benchmark, to now have both a classification and regression task. We have rendered the table below and added it to the supplementary materials of our paper as well.

| Method                                 | RMSE (lower is better ↓) |
|----------------------------------------|---------------------------|
| TabNet                                | 8.909                     |
| SNN                                   | 8.895                     |
| AutoInt                               | 8.882                     |
| GrowNet                               | 8.827                     |
| MLP                                   | 8.853                     |
| DCN2                                  | 8.890                     |
| NODE                                  | 8.784                     |
| ResNet                                | 8.846                     |
| CatBoost                              | 8.877                     |
| XGBoost                               | 8.947                     |
| FT-T                                  | 8.855                     |
| Transformer Scratch                   | 8.809                     |
| Transformer Random Masking             | 8.762                     |
| **Transformer Guided Masking (Ours)**  | **8.695**                 |

Reviewer comment (Experiments):
Missing important related work: DABS [1, 2] is a benchmark for domain agnostic self-supervised learning algorithms. This benchmark consists of semiconductor wafers, multispectral satellite imagery, protein biology, bacterial genomics, particle physics, Speech Recordings, Chest Xrays. DABS also has baselines in the form of Generalized masked autoencoding, Capri (Hybrid Masked-Contrastive Algorithm), e-Mix and Shuffled Embedding Detection (ShED). Demonstrating the effectiveness of SMA on this benchmark would strengthen the paper. I understand that the authors have already shown results on particle physics and protein datasets but comparing with these baselines would lead to a more complete results section. The authors can also discuss and compare SMA with these baselines.

Author response:
The main reason we do not use the DABS benchmark is that their definition of domain agnostic slightly differs from ours, as their benchmark allows for domain-specific tokenization, which is one of the main processes we aim to avoid with our method. As a result, we believe using their benchmark would create some confusion around our definition of domain-agnostic especially because there is some overlap with tasks and datasets used (we have updated the paper to acknowledge these similarities with citations).

However, we agree it is important to have other domain-agnostic baselines to compare against, and we have added results comparing against [5], the current state-of-the-art in domain-agnostic representation learning (far outperforming the e-Mix method as an augmentation method and to our knowledge state-of-the-art on DABS). We significantly outperform their method and include the table of results below, along with updating the original ImageNet100 results in our paper.

| Architecture                             | Pretraining Method                                  | Domain Agnostic | w/o augmentation | w/ augmentation |
|-------------------------------------------|------------------------------------------------------|------------------|-------------------|------------------|
| **Perceiver**                            | Patch                                                | ❌              | **79.6**          | 84.3             |
| **ResNet**                                | None                                                 | ✔️             | 72.5              | 78.1             |
| **ResNet**                                | Random Quantization [6] + MoCo-v3                         | ✔️                 | 76.9              | 81.6             |
| **Perceiver**                            | None                                                 | ✔️             | 52.5              | 62.7             |
| **Perceiver**                            | Random                                               | ✔️                 | 77.6              | 85.6             |
| **Perceiver**                            | SMA (ours)                                           | ✔️                 | **78.5**          | **86.1**         |

Despite their method using a ResNet-50 architecture which outperforms our Perceiver Architecture when trained from scratch on ImageNet-100, our method still outperforms this prior art, we add the results from their method to our other ImageNet100 results. Also note that when running their method on our same Perceiver architecture, we found it resulted in embedding collapse, which we theorize is due to the scratch initialization of positional embeddings, and the contrastive loss does not facilitate the learning of these positional embeddings well.  We include a rendering of the updated Imagenet100 table below for quick reference. We apply the method of [5] with the MoCo-v3 framework, as the original paper found MoCo-v3 to work best with their proposed augmentation. Technically the ResNet architecture is not domain-agnostic, but we denote their method as domain-agnostic because their proposed augmentation can be done domain-agnostically.

Additionally, we now compare against [6], which explores tabular pretraining methods in the same setting as our HIGGS Small benchmark. The main difference is they use the ROC-AUC metric for the HIGGS Small subset achieving a best score of 82.1 ROC-AUC. In comparison, our scratch model achieves an ROC-AUC score of 80.78, our random masking model achieves a score of 82.0 ROC-AUC, and our guided masking model achieves a score of 82.8. Therefore, our method remains the best in this setting.

Reviewer comment (Novelty):
Missing important related work: The authors should compare and contrast with related literature [3, 4, 5].

Author response:
We have compared against [5], which has the best previous empirical results. See our response above for comparisons on the image domain.

Reviewer comment (Experiments):
The authors can run some ablations to better understand the proposed method, SMA. For instance, how does the masking ratio impact performance in various domains? It may be interesting to analyze the performance if the masks are computed in second/ third layer (or any kth layer) instead of first layer in all the experiments.

Author response:
We are not entirely certain what is meant by computing masks in the second/third layer. In our architecture, the attention with the inputs is computed only once. While it could be computed again in theory by computing attention with inputs, computing masks later would require recomputing the initial hidden states again with masked inputs. In comparison with our formulation, no recomputation is necessary, which we believe is important as it allows our method to be applied without any overhead computation.

Reviewer comment (Reproducibility):
What are the hyperparameters in SMA? Is masking ratio the only hyperparameter? Can the authors explain why certain domains require a high masking rate compared to others (as mentioned in Table 7 to 11)?

Author response:
Yes, masking ratio is the only additional hyperparameter that is tuned for each domain. Regarding the difference in masking rate across domains, this is still somewhat of an open question in masked modeling, but most works point to increased redundancy between different domains requiring different optimal masking ratios. We provide details on how masking ratio sweeps were done below.

For NLP and vision domains we simply set the masking ratio of our method to the masking ratio of the hand-designed masks. For vision random masking we used the 0.85 value from the Perceiver paper, which performed an experiment with random pixel masking as well and already performed a sweep. For NLP random masking we performed a sweep of values from 0.15 to 0.3 at increments of 0.05. We observed a decrease in performance beyond 0.2 and therefore did not sweep higher.

For the protein and chemistry domains, we swept from 0.15 to 0.3 and used the optimal value. Prior work has also found optimal masking values to be in the same range [1, 2]. For the HIGGS benchmark, we found further improvements for random masking at higher rates therefore, we swept until 0.5, where we eventually noticed worse performance beyond this masking rate.

Reviewer comment (Experiments):
How would the performance differ if a domain-specific tokenizer is used in Chemical property prediction?

Author response:
Our work aims to avoid using any domain-specific tokenizers. In the case of the chemistry domain, the domain-specific tokenizers simply reduce certain element symbols in SMILES strings from many characters into a single token. While the results may vary depending on the domain-specific tokenizer used, from our testing with the ChemBERTa-2 tokenizer we found the difference in performance to be statistically insignificant on downstream tasks. Additionally we found the tokenizer to only merge approximately 2% of the tokens on the pre-training set meaning the difference between their tokenizer and a simple UTF-8 mapping (the one we use) is relatively small.

Reviewer comment (Experiments):
The results on image classification and NLP tasks are only on toy datasets and seem to be too preliminary to convince people from these communities to try the approach. The gains are marginal and only a small set of methods are compared. The final model is far from the state-of-the-art in NLP and vision. I would recommend if possible to be more ambitious and demonstrate results on more large scale tasks such as linear evaluation on ImageNet.

Author response:
We agree with this assessment. The purpose of these experiments was not to suggest our method be used for the NLP or vision modalities, as these two domains have well-developed inductive biases such as causality in NLP or augmentations in vision.  Rather, these experiments showcase how our method compares with hand-designed masks in a controlled setting, given that the other domains we benchmark on do not have hand-designed masks. Our method is most useful in domains without well-developed inductive biases and allows for competitive self-supervised representation learning without domain knowledge.

We note that our purpose is _not_ to advance the state of the art in established domains such as language or images but to develop a domain-agnostic method that readily applies to new domains such as protein biology or particle physics. Currently, scaling to the full ImageNet dataset is beyond our computational resources. In fact, even at ImageNet100, the image experiments are the most costly in terms of compute, and an experiment with ImageNet-1k would require at least 10x the computational resources.

Reviewer comment (Experiments):
Some design choices are not well motivated and should be ablated properly. For example masking ratios, number k of elements in Eq.3.

Author response:
Following prior work we sweep for the optimal masking ratio in each of our domains. The exceptions are in the language and image domains we do not sweep for patch masking or word masking, given these masking ratios have been well studied. The number of elements $k$ is dependent on the masking ratio and is always 2 * # of inputs / # of queries. Generally, we found this to work well across domains, and we try to minimize the need to sweep for hyperparameters with the exception of masking ratio.

Reviewer comment (Experiments):
Do you need to mask queries? Could you simply mask the input randomly? Could you clarify if this corresponds to “Random Masking” in your tables?

Author response:
Yes, “simply masking the input randomly” is the random masking baseline in our tables. We find that we outperform this baseline in every domain with our method not introducing any additional compute. Lastly, we would like to clarify that we are not masking the queries but rather using the query attention to calculate an optimal mask of the raw input values.

Reviewer comment (Novelty):
Originality: the paper did not cite or compare other domain-agnostic SSL methods, either contrastive [3] or masking [4]. Also, the key components, latent query tokens (Jaegle 2021b) and the KeepTopK (Shazeer et al. 2017) are not novel, further weakening the originality of this work.

Author response:
Thank you for the additional related work. Regarding comparisons, we have now added [5] as a point of comparison; this is a more recent improvement over [3] mentioned by reviewer 3 on our image domain benchmark settings. We believe [5] is the best prior work in domain agnostic learning, and it far outperforms [3]. We show the results below:

| Architecture                             | Pretraining Method                                  | Domain Agnostic | w/o augmentation | w/ augmentation |
|-------------------------------------------|------------------------------------------------------|------------------|-------------------|------------------|
| **Perceiver**                            | Patch                                                | ❌              | **79.6**          | 84.3             |
| **ResNet**                                | None                                                 | ✔️             | 72.5              | 78.1             |
| **ResNet**                                | Random Quantization [5] + MoCo-v3                         | ✔️                 | 76.9              | 81.6             |
| **Perceiver**                            | None                                                 | ✔️             | 52.5              | 62.7             |
| **Perceiver**                            | Random                                               | ✔️                 | 77.6              | 85.6             |
| **Perceiver**                            | SMA (ours)                                           | ✔️                 | **78.5**          | **86.1**         |

Despite their method using a ResNet-50 architecture which outperforms our Perceiver Architecture when trained from scratch on ImageNet-100, our method still outperforms this prior art, we add the results from their method to our other ImageNet100 results. Also note that when running their method on our same Perceiver architecture, we found it resulted in embedding collapse, which we theorize is due to the scratch initialization of positional embeddings, and the contrastive loss does not facilitate the learning of these positional embeddings well.  We include a rendering of the updated Imagenet100 table below for quick reference. We apply the method of [5] with the MoCo-v3 framework, as the original paper found MoCo-v3 to work best with their proposed augmentation. Technically the ResNet architecture is not domain-agnostic, but we denote their method as domain-agnostic because their proposed augmentation can be done domain-agnostically.

Regarding [4], we believe this work is more focused on different “domains” within images (real images, sketches, paintings, etc), which is different from our work which focuses on different data modality “domains” (image, text,  For this reason we cannot compare with [4] as their method addresses a different problem.

Reviewer comment (Evaluation):
Quality: the quality of this work is lacking. Empirical performance improvement can be limited, such as results for the MoleculeNet (Table 2), where the proposed method is sometimes worse than the baseline method (BACE and HIV of Uni-Mol-20M), or improvement can seem limited (Lipo).

Author response:
Uni-Mol is not a baseline but rather the prior art for this benchmark. Notably, it uses a separate data representation that gives additional information to the model not available to our text-based representation. Compared to the prior art text-based representation we far outperform their best results across all tasks. Additionally, there are many other prior graph-based works that we did not include for the sake of space.

Reviewer comment (Experiments):
Quality: while a lower performance is common, since there are only limited baselines (TabNet for the HIGGS benchmark and baselines for the HIGGS benchmark are all from 2021 or prior) and this work is quite empirical, the performance difference can seem noticeable. Nevertheless, the reviewer is not familiar with the benchmarks and did not extensively search for new work with higher results, and will hugely benefit from a response from the authors explaining why the baselines are few.

Author response:
The prior works for the HIGGS small benchmark span 9 prior works. While it is true that this is less than popular benchmarks such as ImageNet, architecture research for tabular data is less common.

For the HIGGS self-supervised learning benchmark, TabNet is the only prior point of comparison, as it was originally introduced in that paper. Previously we were not aware of any prior works extensively evaluating pre-training, which is why we also compared under this setting so that we would have at least 1 point of external comparison for self-supervised learning. However, we performed a wider literature search and found [6], which explores pretraining methods for tabular data. They use the ROC-AUC metric for the HIGGS subset achieving a best score of 82.1 ROC-AUC. In comparison, our scratch model achieves an ROC-AUC score of 80.78, our random masking model achieves a score of 82.0 ROC-AUC, and our guided masking model achieves a score of 82.8 ROC-AUC. Therefore, our method remains the best in this setting.

Additionally, we found one more comparison point [7], which achieves 73.0 on the HIGGS benchmark, 0.1 higher than the best we previously found FT-T. Previously we also had 5 additional comparison points that we did not include due to them having worse performance than our other points of comparison. The full table with these entries is shown below and in our revised paper.

| Method                           | Accuracy (%) |
|----------------------------------|--------------|
| TabNet                       | 71.9         |
| SNN                           | 72.2         |
| AutoInt                       | 72.5         |
| GrowNet                       | 72.2         |
| MLP                           | 72.3         |
| DCN2                          | 72.3         |
| NODE                          | 72.6         |
| ResNet                        | 72.7         |
| CatBoost                      | 72.6         |
| XGBoost                       | 72.7         |
| FT-T                          | 72.9         |
| Transformer-PLR [7]             | 73.0         |
| Transformer Scratch               | 72.7         |
| Transformer Random Masking        | 74.0         |
| Transformer Guided Masking (Ours) | **74.8**     |

Reviewer comment (Evaluation):
Page 5, “First, because the top attention tokens chosen for each query have significant overlap, we often do not actually mask the target masking ratio, and this issue worsens as $n$ increases.“ How much overlap did the top attention tokens have? Did the authors quantify them? How does this issue worsen as $n$ increases – a log or linear relationship?

Author response:
We conducted an experiment in the image domain to empirically evaluate this relationship. We resize an image into different resolutions, thus producing inputs with a different number of tokens, and then compute the overlap ratio. The overlap ratio is computed as (# of tokens in top attentions - # unique tokens in top attentions) / # tokens in top attentions. Therefore, a higher value indicates greater overlap. The graph is shown in the appendix of our updated paper in Figure 4; the overlap ratio increases with the number of tokens and typically ranges between 35% and 50%.

Reviewer comment (Experiments):
Page 5, “...while achieving the desired masking ratio and parallelizing well.” Unfortunately, there are no follow-up experiments or discussions on better parallelization. Why is this the case, and how much parallelization improvement does the proposed method bring?

Author response:
The parallelization is due to performing only a single top-k operation compared to needing a top-k operation for every selected query. Compared to the other additional operations (softmax and a sum along the input length dimension), the top-k operation is the bottleneck in running time on modern GPUs. By requiring only 1 top-k operation we can significantly improve running time.

Empirically, the running time without our proposed approximation on the image domain is 1.13 steps/second. In comparison, when using our approximation, the training runs at 3.58 steps/second, over a 3x increase in throughput. Note that with the approximation, the time to compute masks does not significantly impact train step time (less than 2% of the time per step), but without the approximation, the time to compute masks becomes the bottleneck for large input counts. Additionally, we found no difference in performance between the two methods across domains, and our approximation actually improves stability in our testing.

Reviewer comment (Theory):
Page 5, “let $\mathcal{P}$ represent the set of all permutations of indices” and Eq.(1). If $p$ is a permutation, what does $pX^{(i)}$ mean in Eq.(1)? And more importantly, why Eq.(1) defines domain-agnostic? It is not clear as there are no direct citations or proof supporting this claim (the Perceiver paper did not seem to include any specific math statement like this). The reviewer would appreciate more explanation on this part.

Author response:
The Perceiver paper evaluates the domain-agnostic nature of architectures by proposing to shuffle all tokens along with their corresponding positional embeddings. In their experimental results, they denote these results with “shuffled,” where Perceiver is the only architecture that does not experience a performance drop seeing as it is invariant to this form of transformation. 

We expand on this definition but define a learning process as domain-agnostic if a consistent shuffling of all tokens results in the same performance following pretraining. One way to understand this definition would be to consider treating any data format as flattened tabular data, given that most tabular methods (with the exception of the ResNet model we compare with) do not use knowledge of how related two columns are. If we were to shuffle all the columns, the performance of the learning algorithm would remain the same. We essentially treat all data in this manner, thereby allowing our method to be seamlessly applied to new domains without any adaptation or knowledge of the domain.