CaFT: Clustering and Filter on Tokens of Transformer for Weakly Supervised Object Localization

We use the ViT model as backbone and directly use the output of backbone as final classification result.

We take the tokens of last three blocks out as $map_{0}$,$map_{1}$,$map_{2}$ respectively, and the position embedding parameter of ViT as $map_{p}$. Then we merge these four maps to $map_{m}$ as

where $\alpha_{i}$ is the merge ratio.

In training process, we have designed to have three steps on $map_{m}$, as is shown in Fig.1.

1. (a)

   At the beginning of training, we input the merged tokens of $map_{m}$ to the clustering module, in which these tokens are clustered into 3 categories by k-means algorithm and the tokens with the same category as class token are chosen to be foreground. The next step is to binarize the category distribution map according to the foreground and utilize classical image filter to reduce some noises. Up to this point, the initial masks of images have been generated.

2. (b)

   The next step is to take the initial masks of Clustering as pseudo targets for Attention Filter 1 (AtF₁), which is a two-category classifier of foreground and background.

3. (c)

   After finishing training of AtF₁, CaFT divides the input image into four parts, inputs them into the model respectively and merges their masks together to a more refined mask. A new AtF₂ with the same structure as AtF₁ is trained on the pseudo label of refined masks.

In inference, the $map_{m}$ is directly input to AtF₂, and the output mask of AtF₂ is used to generate one bounding box for one image.

ViT firstly separates the image into patches, then these patches are flattened and put through a linear layer to become n tokens. Different from CNN network, ViT adds an extra class token and input it to network with other tokens from images. In addition, ViT adds a position embedding for every tokens before being input. Experiments of [3] show that position embedding is related to the corresponding position of tokens. The major structure of ViT is encoder, in which tokens will go through cascaded blocks consisted of a multi-head self-attention layer and a Multilayer Perceptron (MLP). After propagation in encoder, the network outputs the same number of tokens t_∗,t₁,…,t_n-1,t_n, and uses the class token t_∗ to inference the category.

A little different from original ViT model, we reserve output of the last three blocks of ViT for the subsequent process.

Classical CAM-based methods usually need to set a handcrafted fixed threshold to decide the foreground and background. For different inputs, the most suitable threshold to select foreground are different, as shown in Fig.3. However, under the previous frameworks, it is able to only choose a compromised value according to the performance on the whole dataset. In other words, a fixed design cannot extract the accurate boundary of foreground and background. Because of the compromised boundary, the predicted region is unable to cover the entire object, so it is inclined to focus on the most discriminative region.

Therefore, we decide to avoid this handcraft process. We find the eigenvectors that make up the feature map have apparent correlation with each other. But the similarity is still a continuous value and has different magnitude among images.

So, we apply clustering method to this task and choose the common k-means algorithm to automatically learn the boundary of foreground and background on the one-image level. At the beginning of experiment, as shown in Fig.4, we use the feature map of CNN-based network (we use ResNet50 [5] in this experiment) to cluster category of each point. However, we find the result of clustering is unstable with different random seeds which influence the initial state of clustering, as shown in the top line in Fig.4(b). By computing the similarity matrix of eigenvectors, we find the similarity between vectors of CNN are commonly low, in Fig.4(c). Moreover, in Fig.4(d), we can see the vector tends to have higher similarity with the vectors among it and lower similarity with the vectors far from it even if they are all regions of object. Most important of all, there is a fatal problem of CNN that although we can cluster the categories, we have no way to decide which cluster belongs to the object’s region.

It is easy to solve above problems by the Vision Transformer (ViT). Firstly, due to the special design of the extra class token, we find there are obvious correlations between class token and other tokens covering the object. And these correlations have a characteristic of long range, which is because of the different topological structure of CNN’s eigenvectors and ViT’s tokens. As is shown in Fig.5, eigenvectors are indirectly connected, while tokens acquire direct information from each other. From the similarity curve with the same image in Fig.4(c), it is obvious that the similarity of tokens has bigger variance than CNN’s vectors, which means a higher degree of difference between clustering classes. Therefore, the result of clustering is more stable for ViT’s tokens.

To verify the feasibility of using class token as flag of object cluster, we compute the mean Euclidean distance between class token and center of clustering of sampling images. It (10.653) is close to the mean distance within cluster to center of clustering ($D_{ic}$) and is much less than clustering radius ($D_{r}$) of $map_{m}$ shown in Table 1. This proves that the class token belongs to the cluster of object region.

As the Fig.7(c) shows, after choosing the cluster of class token as positive region, the positive mask has almost covered the entire object but there are some noisy points in the background. So, we use simple smoothing filter of picture processing to suppress these noises, which outputs Fig.7(d). Although the mask covers the entire object, the edge of mask does not perfectly fit to the object. Therefore, we proposed the next step of Attention Filter to fine tune the edge result and replace the clustering to accelerate the inference.

Some CAM-based methods usually separate the training of WSOL into several steps. PSOL [16] uses the DDT [13] algorithm to generate bounding boxes and considers these boxes as pseudo box-labels to train the box-regression network in the second stage. Inspired by this, we have designed a multi-stage training process. Different from regression of bounding boxes, we regard the subsequent training as segmentation training and take the output mask of clustering as pseudo label.

Through the previous clustering we obtain a set of masks covering the object but have some noises of background. We use Gaussian filter on the noisy mask to smooth part of these points. After that, it already has the ability to predict the bounding box, and the model with backbone ViT-B can achieve 70.16% and 57.59% GT-Known accuracy on CUB-200 and ImageNet-1k respectively.

However, it is not enough for a high-quality WSOL model. Besides the noise of background, the clustering process has a low inference speed because of the iteration of clustering algorithm. According to our early experiments on CNN, we find that there is less noise in the clustering result of CNN. Besides, convolution on the feature map with the same kernel is able to increase the resolving power of foreground and smooth the noise. Therefore, we use a shallow convolution head (Fig.2) following the backbone, which directly extracts the attention mask of the object from merged tokens $map_{m}$. This module plays an analogous role as the filter in image denoising, so we named it as Attention Filter (AtF).

Take the ViT-B-384 for example, after obtaining $map_{m}$, we pull the class token out and send the rest as feature map with shape (D, H, W) to cascaded convolution layers (AtF₁) which outputs an attention mask with shape (2, H, W). D means the embedding dimension of ViT. The 2 channels of output feature map correspond to two classes as background and foreground, and after choosing the points of foreground, AtF₁ outputs the (1, H, W) mask which could predict the bounding box of object by finding the external border. During training AtF₁, the ViT backbone is frozen, which can prevent the classification accuracy of backbone from being influenced by training AtF₁ and also avoid the ViT backbone with strong fitting ability from over-fitting the training set and the noise of the pseudo label.

After being trained on the initial mask from clustering, AtF₁ can achieve a big rise on location accuracy, from 70.16% to 85.64% on CUB-200. The output mask is shown in Fig.7(e). The $1\times 1$ convolution is similar to fully connected layer for each token, so it is essentially eigenspace transformation of tokens to enhance the feature. After transformation, it is able to separate the vague vectors around the cluster boundary. We do the same clustering process as original tokens of $map_{m}$ on feature maps of convolution layers in AtF (Conv1, Conv2). Table 1 shows the comparison of following metrics. The mean within-cluster distance to center of clustering ($D_{ic}$), as

where Euclidean(*, *) is computing Euclidean distance, $C_{o}$ is the set of points in cluster of object region with size $n_{o}$, and $c_{o}$ is the center of cluster of object region. Clustering radius ($D_{r}$), as

Distance between centers of clustering ($D_{cc}$) and $D_{cc}$ divided by $D_{ic}$ and $D_{r}$, as

where $k$ is the number of clusters in k-means algorithm, $c_{q}$ is the center of cluster $q$. Calinski-Harabasz score [1] ($Score$), which is defined as

where $n_{E}$ is the size of the whole tokens set $E$, $W_{k}$ is the trace of the within-cluster dispersion matrix and $B_{k}$ is trace of the between group dispersion matrix defined by

with $c_{E}$ is the center of $E$ and $n_{q}$ is the number of points in cluster $q$. Calinski-Harabasz score evaluates the quality of clustering by within-cluster and among-cluster variance. A higher Calinski-Harabasz score relates to a model with better defined clusters.

Through AtF₁, $D_{ic}$ and $D_{r}$ descend, $D_{cc}/D_{ic}$, $D_{cc}/D_{r}$ and $Score$ increase, which means the degree of within-cluster aggregation is higher, the spacing between clusters is larger and a more stable clustering result. The visual result is shown in Fig.6. Therefore, AtF is able to enhance the feature and reduce the noisy points in the mask.

To further enhance the result, we divide the training image into four parts and use the AtF₁ to predict four masks. Then these four masks are stitched into a more refined mask which is the pseudo label for AtF₂. Compared to directly using the result mask of AtF₁, AtF₂ trained on the refined mask performs much better. Although the box prediction on refined mask is not much better on CUB-200 (even worse on ImageNet-1K), in Table 2. Because the detail of edge preforms better in refined mask while there are some extra sporadic noisy points, which is shown in Fig.7(f). However, as Fig.7(g) shows, AtF₂ is able to smooth these noisy points and obtain a more accurate edge information from refined masks at the same time. While the AtF₂ trained on AtF₁ mask tends to drop into the local minimum as AtF₁ because they have similar structure. These are reasons why AtF₂ trained on refined mask has a big rise from 87.11% to 94.17% and surpasses the AtF₂ trained on AtF₁ mask by a large margin.

By the way, We can only use the divide-method after obtaining AtF₁ rather than when clustering, because the divided part of images may not contain objects and the result of clustering from blank part is unauthentic.

CaFT is based on the ViT backbone [3], which is pre-trained on the ImageNet-1K [7]. Each input image is directly resized to (384, 384). For experiment on CUB-200, we use the pre-trained weight of ViT to train the classification backbone on CUB-200, the following training never changes parameters of the backbone. The top-1 classification accuracy of the backbone is around 88.50% and 83.70% on CUB-200 and ImageNet-1K respectively.

We use k-means algorithm in package scikit-learn, and set 3 categories as the target with other parameters in default. On CUB-200 we use the whole training set of 5994 images. For experiment on ImageNet-1k, because of the high computing pressure of k-means algorithm, we just randomly sample 10 images from each class in training set to generate pseudo mask label, and train the AtF on this mini set (10000 images in total).

We use SGD optimizer for all training. For training AtFs, we set the learning rate as 0.1 with cosine scheduler and 20 epochs for CUB-200, and 5 epochs for ImageNet-1K.

We have tested the results of ViT-B and ViT-L model. For ViT-B model, we follow the integral process with clustering, AtF₁, refined mask and AtF₂; while for ViT-L, we directly train the AtF₂ using the pseudo mask labels generated by ViT-B model.

In this section, we will verify the components of CaFT by ablation study on CUB-200 dataset, and introduce the effect of some details.