Exploring Compositionality in Vision Transformers using Wavelet Representations

Questions

1) Could the authors clarify how compositionality with discrete wavelet transforms leads to better explainability? Some examples or related work would help. I can see some intuitive argumentation but I think it has not been articulated in the paper.

Given the challenging nature of composing 'semantic' primitives for object recognition, the paper presents DWT as a potential framework for analyzing compositionality in ViTs.

• The analysis presented in the paper explains the relative importance given to the individual frequency bands by the ViT models. This may lead to understanding the object recognition task from the frequency domain perspective.
• The presented compositionality framework enables the incorporation of any specific domain knowledge with respect to the primitives (in this case, the DWT sub-bands) towards the downstream task. In other words, one can train models that can give specific weights to the sub-bands while learning to solve the task if needed.
• In the future, the community may bridge the frequency-to-semantics gap with the availability of sophisticated tools so that the presented framework can directly compose to the semantics

2) The observation that convex combination of sub-bands in image space alters pixel values significantly is interesting. Why does it not degrade accuracy? Is it because of normalization?

This is definitely interesting. Probing the combinations themselves revealed that the $l2$ norm between the convex combination and the original image is lower than the other two combinations. We hesitate to say it is because of normalization and this would require further analysis to answer.

We thank the reviewer for their assessment and we hope to answer their questions point by point.

Weaknesses

1.a) I believe the author's experimental setup is not using sufficient data. For reference, papers such as ViT-NeT which explores interpretability of ViT uses three datasets each of which is 10-20k images total. I'd suggest authors scale up their datasets - if we choose to use ImageNet-1k, a test set of at least 10 images per class would be more convincing. Currently, with a 15% test fraction, each class gets 1-2 images. I hesitate to make conclusions based on such a small test set per class.

We thank the reviewer for this comment. Please note that the parameters we are training for the compositional model are merely a handful. Also, we believe ImageNet is one of the most complex datasets that offers a lot of variety among the samples. In the revision, we scaled up the dataset to 50000 samples (1000 classes,50 images/class) with a 60:20:20 train:val:test split. So the test set contains 10000 samples (each class gets 10 images) and repeated some of the experiments. Due to time constraints, we could only carry out the experiments for level 1 decomposition. We plan to carry out the rest of the level 2 experiments with the scaled-up dataset. The results for Level 1 have been updated in the revised draft.

Comparing the accuracy of different composition models $g^{*}$ on the testset:

Relative accuracy of the learned composition models. Note that the target for the composed representation is the output predicted by the original image classifier (note the ground truth label).

1.b) No significant analysis is done on the test dataset, nor of the errors that composition leads to. It would be beneficial to provide some examples where composition lead to large error and some examples where the error is minimal. For example, Haar DWT is great when we have sudden transitions in signal - like sharp images. Perhaps blurrier examples or classes would cause larger error?

We appreciate the reviewer's point. We have conducted a preliminary experiment (Section 4.4 in the revised draft) to identify the samples on which the learned model fails. However, a deeper analysis to test the robustness of compositionality is currently out of scope of this paper. We would like to emphasize that the primary goal of this work is to present a framework to check if notions of compositionality are present in the representation space and the level 1 decomposition results provide strong support.

2) Authors do not inspect intermediate layers of ViT, which is a glaring hole to me - while it is true that most downstream tasks will use the penultimate layer of the ViT, I am still left wondering if all intermediate representations are composable, if its just some of them, is compositionality lost at some levels of the stack? I would recommend repeating the same analyses using intermediate layers’ representations.

The intermediate layer representation of each image is of size 197x768. The cls token is of size 1x768. Since the input to the classification layer is just the cls token, performing this analysis on the final encoder layer is easier. There is no trivial way to classify the intermediate layer's output. To perform the analysis on intermediate layers would require storing the entire 197x768 representation for each image and its subsequent wavelet components in order to train the model. Due to memory and computational limitations we could not perform the analysis for the intermediate representations.

We thank the reviewer for their constructive feedback. We hope to answer their queries point by point.

Weaknesses

1) My main concern is that the presented results do not convincingly demonstrate compositionality. Rather than defining true combinations of wavelet primitive representations, it appears that the learned weights mainly select the low-pass filtered image (Table 3). Indeed, it is not particularly surprising that the images in Figure 5 perform similarly to the original images.

We thank the reviewer for this comment. As mentioned in section 3.1 of the draft, our motivation for a learned composition of the wavelet primitives instead of the true combination (summing each primitive) of primitives in the representation space is based on our initial experiments of SSIM scores and CKA plots. These results show that the “true” combination does not demonstrate compositional behavior. We then speculate that if any compositional behavior does exist, could it be learned? Our results confirm that such a composition exists for the final encoder layer for level 1 DWT decomposition. To answer the reviewer’s concern as to the selection of mainly the low-pass filtered image, we conducted an experiment to see how the low-pass filtered image’s representation performs compared to the original image’s representation. We take 10 images per class from all the 1000 classes in ImageNet-1K and present the results.

Although the learned composition pays more attention to the low-filtered component, the other components are also important. This clearly signifies that a composition of all these components does indeed approximate the original representation much better than the low-pass component.

2) Compositionality is typically more valuable when components are semantic rather than appearance-based. It is doubtful that wavelets would exhibit compositional properties in the final layers of a model, where higher-level concepts are typically captured; instead, this is more likely to occur in lower-level layers. Furthermore, the idea that wavelets are a good basis for compositional representations is not really explored, and no other decomposition methods are considered or compared.

We understand that probing compositionality is done by identifying simpler concepts that can be composed into complex ideas. In the NLP setting, it is easy to find this analogy, which is that words can form simpler concepts composed into complex sentences. However, it is difficult to break down in the context of vision when the complex ideas are images. The primitive set required for the framework is very challenging to identify since the image space is continuous. We chose the DWT to extract the primitives because of the sound mathematical proof that wavelet decompositions have perfect reconstruction. We check for a homomorphism between the input and representation space on the foundation that the DWT has perfect reconstruction in the input space. To the best of our knowledge, we are the first to probe for compositionality with the framework presented in our paper using wavelets. Other decomposition methods, such as Fourier transform or Discrete Cosine Transform (DCT), also exhibit lossless decomposition and reconstruction, but they do not preserve spatial information, i.e., the frequency spectrum of an image is not visually meaningful. Hence, we chose DWT as the method for decomposition for our framework.

Questions

1) Could the authors clarify why they believe their results demonstrate compositionality, why the DWT is considered a suitable feature decomposition, and why compositional behavior would be expected in the last layer?

The framework presented in the paper divides the input image into its DWT primitives. The original image and its primitives are passed through the model separately. The individual primitives are linearly combined using the weights of the learned compositional model and is then passed through the final classification layer. The results demonstrate that composed representations perform similarly to the original image representation. We chose the DWT to extract the primitives because of the sound mathematical foundation that wavelet decompositions have perfect reconstruction. On the basis that a perfect composition exists in the input space, we investigate if a homomorphism exists between the input and representation space learned by the ViTs. While our analysis demonstrates the notions of compositionally (as per the framework introduced by Andreas 2019) with respect to the DWT primitives, more needs to be understood about why it is manifested.

2) l.299: "The target is the final classification layer output of the original image" does this mean using an L2 loss?

We use the Cross Entropy loss between the classification layer output of the original image’s representation (i.e., predicted soft labels) and the classification layer output of the composed representation.

3) l.445: Why is the analysis conducted across different settings?

We regret this confusion. The analysis was conducted using the same test set. At the time of submission, we could only run a part of the test set for some of the results. We have now revised the draft such that the updated results follow the same setting.

We thank the reviewer for their detailed response and valuable feedback. We hope to answer all of their points in similar fashion.

Weaknesses

Clarity

1) Figure 1: What does it mean here for the original image’s representation to be compared to the composed image representation? Does it mean that the maps we see in the figure are the result of the comparison, or are these just the composed representations and the comparison was done outside of the figure?

The SSIM maps shown in the figure are the results after comparison of the original representation and the composed representation at every layer of the encoder. We have included the clarification in the caption of the figure in the revised draft.

2) Section 3.1 ViT: What was the ViT model used for this experiment? What was it pre-trained on? The section does not explicitly specify these experimental conditions.

We thank the reviewer for pointing this out, and we have included the clarifications (L321-L322) in the revised draft. To answer this query, the models used are ViT- Base and ViT Large, both of which are pre-trained on the Imagenet-21k dataset.

3) L214-L215: What is the purpose of f_eta ? I understand that this is defined in Andreas 2019, but its definition is missing from this paper. I think explicitly defining it and its role would improve the clarity of this section.

We thank the reviewer for pointing this out, and we have included the clarifications in the revised draft. The clarifications are included at L200-L201 and L225-L226 in the revised draft. To answer the question $\hat{f_{\eta}}$ is a compositional approximation of the complex model $f$ . It can be used to measure the compositionality of the model.

4) What do the C’s and a’s represent?

The Ca's and Cd's represent the approximate and detail coefficients for each level of decomposition. The approximate is the output of a low pass filter and the detail is the output of a high pass filter. For the next level, the approximate is passed through the low pass and high pass filter to get the coefficients for the next level.

5) What are LL, LH, HL, and HH?

LL,LH,HL and HH are the corresponding Low-Low, Low-High, High-Low and High-High subbands from the wavelet decomposition. We have edited the figure and corrected the labeling. The labels for each image now follow the previously defined notations.

Soundness

1) Section 3.1: I believe this section aims to show that simply adding the DWT wavelets with equal weights does not yield a correct composition. However, this claim is supported by showing results over a single image from a single ViT model. I believe this experiment would need a much larger sample size over images/models in order to make such a broad claim.

This is a minor confusion. The SSIM map (Figure 1) was computed over a representative image but plot (in the initial submission) showing the CKA scores (Figure 2) takes 200 images and averages their performance over all encoder layers. Taking the reviewer’s point into account we have conducted another experiment by taking a total of 10000 images (10 images/per class from imagenet-1k dataset). The current figure in the revised draft shows these results.

2) L216-217: The original compositionality formulation from Andreas 2019 is modified to shift the application of the encoder function from the input space to an arbitrary intermediate representation space within the model. If I’m understanding correctly, doesn’t this violate the core premise of the problem statement? The purpose of compositionality tests is to find homomorphisms between the input space and the representation space, this is important because the input space comes from the data generating function. I’m not sure if it makes sense for this test to be defined for the transformation from one hidden layer to another. Minor, but I would suggest modifying the sentence “instead of drawing exact parallels, we tweak this statement to suit our analysis”, since the wording gives the impression that the formulation was modified to suit the narrative of the paper, rather than the needs of the original empirical question being asked.

We thank the reviewer for this observation. We regret the confusion caused by overlooking it. The proposed framework investigates the homomorphism from the input space to the embedding space learned by the ViTs. We have rectified the mistake in the revised draft (L228-L234). To clarify, the model is indeed taking inputs from the input space (which is the image space), whose derivatives (wavelet decompositions) are fed to the model $f$ . To check for the compositionality of layer $l$, the encoder representations of that layer (transformation from the input to that layer) are fed to the composition function $\hat{f_{\eta}}$