Exploring Compositionality in Vision Transformers using Wavelet Representations

TL;DR

The paper investigates whether Vision Transformer embeddings exhibit compositional structure when viewed through input-dependent wavelet primitives. It proposes a post-hoc compositionality framework that uses the Discrete Wavelet Transform (DWT) to generate primitives and learns a composition function $g_{ta}$ to approximate ViT encoder outputs, focusing on the final layer. The results show that one-level DWT primitives enable approximate compositionality, with the learned model outperforming simple addition and remaining robust under distortions such as noise and JPEG compression. This work provides a new lens for interpreting ViT representations and suggests a path toward more explainable and robust vision transformers.

Abstract

While insights into the workings of the transformer model have largely emerged by analysing their behaviour on language tasks, this work investigates the representations learnt by the Vision Transformer (ViT) encoder through the lens of compositionality. We introduce a framework, analogous to prior work on measuring compositionality in representation learning, to test for compositionality in the ViT encoder. Crucial to drawing this analogy is the Discrete Wavelet Transform (DWT), which is a simple yet effective tool for obtaining input-dependent primitives in the vision setting. By examining the ability of composed representations to reproduce original image representations, we empirically test the extent to which compositionality is respected in the representation space. Our findings show that primitives from a one-level DWT decomposition produce encoder representations that approximately compose in latent space, offering a new perspective on how ViTs structure information.

Figures (6)

• Figure 1: Comparison of DFT, DCT, and DWT decompositions. Unlike DFT and DCT, whose subbands represent global frequency components, DWT produces spatially localized, visually interpretable subbands suitable for compositionality analysis.
• Figure 2: Tree structure of the Discrete Wavelet Transform (DWT). S represents the input signal. $Ca_{i},Cd_{i}$ represent the approximate and detail coefficients of $i^{th}$ level.
• Figure 3: SSIM maps for each channel (R,G,B). For each encoder layer output, the original image's representation is compared with the composed image representation. The SSIM maps shown here are after comparison. There is no immediate notion of compositionality present visually.
• Figure 4: CKA scores of original vs composed representations at various encoder layers of ViT-B averaged over 10K images. The learned combination has a much higher CKA score across layers compared to a simple sum.
• Figure 5: Overview of the proposed compositionality framework for ViTs. The figure presents learning the composition function for Level 1 DWT decomposition. $D_{a}$, $D_{b}$, $D_{c}$, $D_{d}$ are the coefficients of the wavelet decomposition discussed in \ref{['eq:sum']}.