Paper: Masked Completion via Structured Diffusion with White-Box Transformers (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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.

|Model (mask ratio)|CRATE-MAE-Base (25%)|CRATE-MAE-Base (50%)|CRATE-MAE-Base (75%)|CRATE-MAE-Base (90%)|*MAE-Small (25%)*|*MAE-Small (50%)*|*MAE-Small (75%)*|*MAE-Small (90%)*|
|-|-|-|-|-|-|-|-|-|
Finetuning|89.87|92.83|96.4|96.58|95.79|96.8|97.61|97.34|
Linear Probing|58.41|62.72|71.12|64.15|54.22|68.99|79.43|75.13|

Reviewer comment (Novelty):
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).

Author response:
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.

Reviewer comment (Evaluation):
Omitting linear probing results restricts the paper from showcasing the practicality and caliber of the representations derived using CRATE-MAE.

Author response:
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.

|Model (mask ratio)|CRATE-MAE-Base (25%)|CRATE-MAE-Base (50%)|CRATE-MAE-Base (75%)|CRATE-MAE-Base (90%)|*MAE-Small (25%)*|*MAE-Small (50%)*|*MAE-Small (75%)*|*MAE-Small (90%)*|
|-|-|-|-|-|-|-|-|-|
Finetuning|89.87|92.83|96.4|96.58|95.79|96.8|97.61|97.34|
Linear Probing|58.41|62.72|71.12|64.15|54.22|68.99|79.43|75.13|

Reviewer comment (Experiments):
The study's dependence on a confined dataset for classification raises concerns about the breadth of its applicability and potential generalization to diverse scenarios.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Novelty):
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).

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

||Train Reconstruction Loss|Test Reconstruction Loss|
|-|-|-|
|CRATE-MAE-Base|0.266|0.303|
|MAE-Small|0.250|0.283|

Reviewer comment (Evaluation):
How are the evaluations done in Table 2? Are they using the encoder only? Are they with fully-visible inputs?

Author response:
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).

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.


||CIFAR10 FT|CIFAR10 LP|CIFAR100 FT|CIFAR100 LP|Oxford Flowers FT|Oxford Flowers LP|Oxford Pets FT|Oxford Pets LP|
|-|-|-|-|-|-|-|-|-|
|CRATE-MAE-Base|96.4|71.12|80.24|48.64|70.29|59.23|76.18|38.38|
|MAE-Small|97.61|79.43|82.96|54.36|84.22|68.45|81.68|51.21|

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
Will the layer normalization influence the rate deduction process, or it's just for making the training process more stable or other reasons?

Author response:
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.