Masked Completion via Structured Diffusion with White-Box Transformers

Thank you for your valuable comments. We especially appreciate your remarks regarding the “theoretical depth and scientific rigor” of our work, as well as the significance of the problem. We would like to respectfully point out a minor error in your review: instead of using 30% fewer parameters than base MAE (that is, using 70% of the parameters of MAE), we use 30% of the total parameters of MAE. We hope that this further clarifies the value of our empirical evaluations.

Below, we attempt to resolve your remaining comments and concerns.

Evaluation Concerns: The introduction of the CRATE-MAE architecture in the paper falls short in offering a comprehensive quantitative analysis when compared to established benchmarks like MAE or contrastive methods. The results presented in Table 2 seem somewhat restrictive, making it challenging for readers to gauge the model's efficacy in relation to others.

Absence of Linear Probing Results: Omitting linear probing results restricts the paper from showcasing the practicality and caliber of the representations derived using CRATE-MAE.

Thank you for pointing it out. We summarize some recent experimental results, which we will put in the camera-ready version, which evaluate the linear probing and fine-tuning performance of CRATE-MAE-Base (44.6M parameters) versus MAE-Small (47.5M parameters) [1] over different mask ratios. This experimental result shows that while CRATE-MAE-Base slightly underperforms its black-box counterpart, the performance is still quite reasonable.

Local Self-Supervised Learning Comparison: From a broader perspective on self-supervised learning (SSL), this paper could be classified under local SSL, emphasizing layer-specific objectives and training. Although this area might be less traversed, incorporating findings from related works, such as [A], would enhance the paper's credibility. [A] Siddiqui, Shoaib Ahmed, et al. "Blockwise self-supervised learning at scale." arXiv preprint arXiv:2302.01647 (2023).

Thank you for pointing this out. We will add a comparison to these methods in the camera-ready version, since we were not aware of them previously. We note that there is a significant conceptual difference between our work and “local SSL”; instead of training layer-wise, we train our model end-to-end and the architecture (derived through discretized structured denoising-diffusion) ensures that the layer-wise characteristics (e.g. compression/sparsity improvements) are preserved. In contrast to the layer-wise training in [A], where each layer has an interpretable goal but black-box mechanisms, each layer and operator in our white-box model is interpretable by construction, and the interpretations hold throughout training.

Dataset Limitations: The study's dependence on a confined dataset for classification raises concerns about the breadth of its applicability and potential generalization to diverse scenarios.

We pre-train on ImageNet-1K, a standard large 2D image classification dataset commonly used for pre-training large image models. When fine-tuning for classification, we use many datasets such as CIFAR10/CIFAR100, Oxford Flowers, and Oxford Pets. We believe that this resolves any concern about fine-tuned classification performance overfitting to one dataset.

Again, thank you for your helpful comments. Please let us know if you have any more concerns or questions. We look forward to a fruitful discussion period.

[1]: Gandelsman, Y., Sun, Y., Chen, X., & Efros, A. (2022). Test-time training with masked autoencoders. Advances in Neural Information Processing Systems, 35, 29374-29385.

Thank you for your valuable comments. We would like to respectfully reiterate that the methodology we introduce is not a diffusion model per se – when one discusses “diffusion models” they usually refer to generative models trained through score matching [1]. On the other hand, in our work we have introduced a framework for invertible representation-learning models, whose representations are achieved through discretizations of probability flow equations, and trained via masked autoencoding (among many possible tasks). Thus the mechanism is analogous, not identical, to diffusion models.

Below, we attempt to resolve the remaining questions and concerns.

In Sec 2.2, the authors intend to learn representations Z, and hope that the results learned are low-dimensional, sparse, bijective, etc. If the method proposed by the authors is a white box, then these properties of Z should be verifiable through experiments. Therefore, the authors should provide experimental results of representations Z to support their theory.

We measure the low-dimensionality (indeed, the compression) of $Z$ by the rate distortion $R^{c}(Z \mid U_{[K]})$. The network is shown to iteratively optimize this quantity at each layer, as demonstrated in Figure 4 of the paper. In this figure, not only are the final representations shown to be compressed, (almost) every intermediate representation in the network is shown to become progressively more compressed, showing that each layer implements compression, which aligns with our conceptual framework. Figure 4 also demonstrates a similar trend with the sparsity. We ensure that the representation is invertible by using a decoder, and training the encoder and decoder together on the masked autoencoding task. Figure 5 and Table 3 demonstrate invertibility of the representations.

I checked the provided Pytorch code and find that MSSA and ISTA are composed of Linear layer, which implies large GPU Memory consumption. So I'm wondering whether this method can be extended to large images, just like stable diffusion. In addition, can the authors provide results from larger datasets? The images from CIFAR and ImageNet-1k are too small.

We are not sure what you mean; linear layers are present in the vast majority of machine learning systems, such as Stable Diffusion [2]. We resize our images to 224 x 224 pixels before processing in order to standardize the input processing and tokenization, so the dataset would not matter. Increasing this size would make further pre-training much more computationally expensive, and infeasible for us to complete during this rebuttal period. If you think it would be better, we can increase the size of the figures in the camera-ready version, so as to make the fine details easier to see.

How effective is this network at unconstrained image generation? In other words, if the learning target is pure noise, which may not be viewed as an image compression task, would this method still work?

Unconditional image generation is at odds with masked autoencoders, which are empirically used to learn representations instead of perform generation, and mathematically attempt to output $\mathbb{E}[X \mid \mathtt{Mask}(X)]$ by minimizing the training loss. The best that masked autoencoders (even black-box MAE) can do with 100% mask ratio is to always output the constant image $\mathbb{E}[X]$, which would be a blurring of all possible images, and so would not be desirable. Our framework does not have to use masked autoencoding as a pretext task, however, and there are several such tasks (e.g. denoising) which may enable generation. Concrete extensions of our framework to generation are out of scope of this paper.

Again, thank you for your helpful comments. Please let us know if you have any more concerns or questions. We look forward to a fruitful discussion period.

[1]: Song, Y., & Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32.

[2]: Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. (2021). High-resolution image synthesis with latent diffusion models. 2022 IEEE. In CVF Conference on Computer Vision and Pattern Recognition (CVPR) (pp. 10674-10685).

Thank you for your detailed response and the effort put into addressing the concerns raised. After careful consideration of your responses and a re-evaluation of the manuscript, I have decided to maintain my original score.

Table 4 again lacks a comparison of a similarly-sized MAE (how it performs on the same dataset as the mask ratio changes). It is not clear a conclusion from ImageNet classification can be transferred to CIFAR.

Thank you for pointing it out. We compared CRATE-MAE-Base against MAE-Small when pre-trained on different mask ratios, and conducted fine-tuning and linear probing on the training sets of CIFAR10. Below, we show the resulting models’ classification accuracies on the test sets of CIFAR10.

We are unfortunately not sure what you mean by “It is not clear a conclusion from ImageNet classification can be transferred to CIFAR.” The pre-training is on ImageNet, and the fine-tuning in the referenced table is on CIFAR10, so we are not extrapolating classification performance anywhere.

Figure 4: how does it compare with a supervised encoder/decoder trained on ImageNet? I want to know the compression and sparsity behavior is a result of unsupervised learning, or a result of the architecture design. The same applies to Figure 6 and 7.

Thank you for the question. Actually, we observe from [1] that such a compression and sparsification effect holds for supervised classification training of the encoder network. Figures 6 and 7 are replicated in the supervised classification context in [3, Figures 1 and 3], which also demonstrates that black-box networks generally do not have this interpretability unless specifically trained for it [3, Figure 14]. To answer your question directly: [1] and [3] demonstrate that the main effect of the interpretability properties we observe is due to the architecture design, and not due to training procedure. We will briefly clarify this in the camera-ready version.

Again, thank you for your helpful comments. Please let us know if you have any more concerns or questions. We look forward to a fruitful discussion period.

[1]: Yu, Y., Buchanan, S., Pai, D., Chu, T., Wu, Z., Tong, S., ... & Ma, Y. (2023). White-Box Transformers via Sparse Rate Reduction

[2]: Gandelsman, Y., Sun, Y., Chen, X., & Efros, A. (2022). Test-time training with masked autoencoders. Advances in Neural Information Processing Systems, 35, 29374-29385.

[3]: Yu, Y., Chu, T., Tong, S., Wu, Z., Pai, D., Buchanan, S., & Ma, Y. (2023). Emergence of segmentation with minimalistic white-box transformers.

Thank you for your valuable comments. We especially appreciate your remarks regarding the work’s “good job in presentation” and the quality of the various visualizations. Below, we respond to the raised questions and comments.

I see a lot of similarities to the main white-box paper (White-Box Transformers via Sparse Rate Reduction) which is already out there for supervised learning. I want more justifications for the meaningfulness of the current work apart from doing unsupervised learning.

While the main point of the work is to extend the white-box framework in [1] to unsupervised learning, we believe that this extension is in itself meaningful. To begin with, one cannot use the unrolled optimization framework in [1] to understand or construct white-box decoder models; this is because they map from structured feature spaces to the original data distribution, so they cannot optimize any representation learning objective at all. This means that one needs to develop a new conceptual toolkit to extend the white-box approach in [1] to unsupervised learning. In this work, we propose the novel “structured denoising-diffusion” framework towards this end. We recover the unrolled optimization framework from [1] as a special case, but this time the framework is able to handle inversion and conditional generation (e.g., able to construct a white-box decoder network) via time-reversal. Note that this structured denoising-diffusion framework could be of independent interest in building practical and principled models for various tasks (such as unsupervised generation and self-supervised learning) beyond what we explored in this work. What’s more, it relies on certain quantitative connections, derived in the paper (e.g. Theorem 1), between statistical notions of denoising and information-theoretic notions of lossy compression, which would be of independent technical and conceptual interest.

While remarks to the above effect are interspersed through the paper, we will consolidate and expand on them in the camera-ready version.

I am not convinced what's the advantage of models being white-box here, especially whether it has synergies with unsupervised learning. The paper explains a lot about what's done, but I don't see a strong motivation of why it is done (especially since the introduction is less of a story but more of a break-down of context and contributions).

Thank you for the comment on desiring more motivation; we will certainly include more exposition towards this end in the camera-ready version. At a high level, white-box or theoretically principled models serve to make models more understandable or interpretable at a high granularity. This line of work thus attempts to bridge the divide between theory and practice; such a divide exists in both supervised and unsupervised deep learning. For theorists, simpler and principled models are easier to analyze and build meaningful theory off of. For practitioners, white-box networks retain empirical benefits over black-box models. To wit, we demonstrate in the paper that we obtain comparable autoencoding results to the MAE-ViT with only 30% of the parameters. (Below, we will clarify several of your remarks about to what degree the results are “comparable”.) We also observe that attention maps in the encoder are interpretable.

The paper lacks a fair comparison with MAE, especially on downstream tasks. Table 1 lists that MAE-base has more parameters and it could partially explain why Crate-MAE has a higher reconstruction loss. So what would a fair comparison (in terms of model parameters) look like? I think it is very easy to train MAE with a smaller encoder/decoder pair given the open-sourced code.

Thank you for pointing it out. We ran a comparison against MAE-Small [2], which uses a similar number of parameters (47.5M) as CRATE-MAE-Base (44.6M). Since you ask for reconstruction loss, we evaluate the reconstruction loss on the training and testing datasets of ImageNet-1K. It shows that CRATE-MAE slightly underperforms its black-box counterpart, but the performance remains reasonable.

How are the evaluations done in Table 2? Are they using the encoder only? Are they with fully-visible inputs?

The precise fine-tuning methodology is discussed in Appendix D due to space constraints. In short: the fine-tuning networks use the encoder only with an attached classification head, and use fully-visible inputs (i.e., no masking).

(Remaining responses and references posted in a followup.)

Empirically, people think that the attention map $Q$/$K$ plays a different role as the mapping matrix $V$ and they'd better not be set to be the same. However, in this papers' theoretical framework, they can be set to the same parameter $U$. If change the mapping matrix from $U$ to $V \neq U$, will the performance of CRATE-MAE change a lot? If no, what's the main reason why CRATE-MAE has such property?

Thank you for pointing this out! We were not aware that people suggest that the $Q$ and $K$ matrices have distinct conceptual roles from the $V$ matrices, and would be interested in any references(s) which explain(s) this, so that we could incorporate them into our paper. We believe that in fact the $Q$, $K$, and $V$ matrices could be considered conceptually in the same role. Our justification in this paper, as well as other works on white-box transformers [2, 3], yield this interpretation from both empirical and theoretical perspectives, as well as a proposal of Hinton: [4]. White-box transformers are simpler [2] and have different, more interpretable [3] operational characteristics than the ViT. Thus, the suggestion that $V \neq U$, obtained through some subtle post-hoc analysis in the ViT context, may not be straightforwardly carried over to CRATE.

Will the layer normalization influence the rate deduction process, or it's just for making the training process more stable or other reasons?

The approximations made to derive the MSSA and ISTA blocks require the features to have constant-order magnitude, so some normalization is required in order to make the implementation match the underlying theoretical principles. Using layer-normalization as opposed to “true” normalization (i.e., manually normalizing each token feature to the unit sphere) is motivated by improved empirical performance and easier training. Our framework does not attempt to give a principled reason why we should use one over the other. We will incorporate this discussion into the camera-ready version.

Once again, we thank you for your insightful comments. Please let us know if you have any other questions or concerns. We look forward to a fruitful discussion period.

[1]: Gandelsman, Y., Sun, Y., Chen, X., & Efros, A. (2022). Test-time training with masked autoencoders. Advances in Neural Information Processing Systems, 35, 29374-29385.

[2]: Yu, Y., Buchanan, S., Pai, D., Chu, T., Wu, Z., Tong, S., ... & Ma, Y. (2023). White-Box Transformers via Sparse Rate Reduction

[3]: Yu, Y., Chu, T., Tong, S., Wu, Z., Pai, D., Buchanan, S., & Ma, Y. (2023). Emergence of segmentation with minimalistic white-box transformers.

[4]: Hinton, G. (2023). How to represent part-whole hierarchies in a neural network. Neural Computation, 35(3), 413-452.

Thank you for your valuable comments. We especially appreciate your remarks that our work is a “novel attempt to both the theory and empirical community” and that “[t]he visualization results are quite impressive.” Below, we respond to the raised questions and comments.

I think the comparisons between CRATE-MAE-Base amd MAE-Base are not fair. I understand that the empirical evaluations are not optimally engineered and actually the visualization results are every good. However, I still think that the author should compare CRATE-MAE and MAE with (almost) the same amount of parameters and different performance, or alternatively, (almost) the same performance and different amount of parameters, and then to compare these two models.

Thank you for the comment. The CRATE-MAE-Base has similar parameters to MAE-Small [1], which uses a similar number of parameters (47.5M) as CRATE-MAE-Base (44.6M). In the following table, we demonstrate a comparison of linear probing (LP) and finetuning (FT) accuracy on the test data on CIFAR10 between these two models (pretrained on ImageNet-1K), showing that CRATE-MAE-Base is slightly outperformed by its black-box counterpart, but the performance remains reasonable.

What's the choice of LASSO coefficient hyperparameter $\lambda$, and the step size of discretization $\kappa$, $\eta$? Are they chose carefully and is the model sensitive to them?

We follow the hyperparameter choices in [2], that is, $\lambda = 0.1$, $\eta = 0.1$, and $\kappa$ and $\epsilon^{2}$ tuned so that $\kappa p/(N\epsilon^{2}) = 1$. We do not choose hyperparameter values carefully (e.g., through ablation or other robust tuning). Indeed, as you already noted the experiments in this work are intended as a proof of concept to illustrate the promise of our white-box approach and framework, though they already demonstrate some empirical benefits such as parameter saving and interpretability. If you would like a more detailed quantitative analysis, let us know and we can upload it to the camera-ready version, but performing a detailed robustness analysis on ImageNet would be computationally infeasible during the review period.

Which dataset does Fig.4 belongs to? Does the model have the similar patterns as Fig.4 on other datasets evaluated in this paper?

Thank you for pointing out that the paper is missing this information; we will include it in the camera-ready version. Figure 4 demonstrates the compression and sparsity of tokens within an image from ImageNet-1K, averaged over 10,000 random samples with error bars displayed. All models in this work were pre-trained on ImageNet-1K, and this task is evaluated before fine-tuning on any datasets. Thus we are not sure if you suggest that we look at models that are pre-trained from scratch on the other datasets using masked autoencoding, or look at models that are pre-trained on ImageNet-1K and fine-tuned on other datasets. We have tried both on the CIFAR10 dataset; we generally observe the same trends as in the paper Figure 4 for both settings.

• Here are the results for the model pre-trained on CIFAR10: https://imgur.com/a/KdTzZDU
• Here are the results for the model pre-trained on ImageNet-1K and fine-tuned on CIFAR10: https://imgur.com/a/GZ2iYbK

(Remaining responses and references posted in a followup.)

Thank you for the reply, and for raising the score!

We totally agree with your point – we believe that the white-box MSSA block is an efficient and useful operator (in part) because of its derivation from first principles as compressing the representations towards the desired statistical model. In particular, such derivations from first principles could potentially inform useful architecture modifications, even if not all architectures currently have such interpretations.

We will add this discussion on parameter-sharing (including some discussion of related work, such as [5], which considers sharing $Q$/$K$ parameters), and an ablation on the sparsity parameter $\lambda$ and learning rate $\eta$, to the camera-ready version.

Thank you again! Please let us know if you have other questions or comments during the discussion period.

Thank you for your detailed response. Your additional experiments and illustration address most of my concerns and I will increase my score.

As for the sensitive analysis, the ablation study for parameters like $\lambda$ may be helpful but not necessary since as you mentioned, the main focus of this study is to construct the theoretical framework of a white-box transformer.

For the roles of $Q/K$ and $V$, I’m very sorry that my claim may be confusing. Some previous works like reformer[5] may claim that shared-QK parameters will not significantly influence the performance of model, but people may heuristically think that the attention matrix $W_Q$/$W_K$ and the value matrix $W_V$ will project the input to different spaces. Actually, I can’t find other reference papers for this common choice, and the materials you provide address my concern, I think it may be beneficial to add the discussion to the paper. Maybe the MSSA block and the update process (2.4~2.6) which is different from the heuristically designed transformer result in the effectiveness of using the single matrix $U$. And I think this is really an interesting point of this approach

[5] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ICLR, 2020.