SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gaurav Bhatt $^{13}$ , Deepayan Das $^{2}$ , Leonid Sigal $^{13}$ , Vineeth N Balasubramanian $^{2}$
<sup>1</sup>The University of British Columbia, <sup>2</sup>Indian Institute of Technology Hyderabad  
<sup>3</sup>The Vector Institute, Canada
Intelligent systems possess a crucial characteristic of breaking complicated problems into smaller reusable components or parts and adjusting to new tasks using these part representations. However, current part-learners encounter difficulties in dealing with incidental correlations resulting from the limited observations of objects that may appear only in specific arrangements or with specific backgrounds. These incidental correlations may have a detrimental impact on the generalization and interpretability of learned part representations. This study asserts that part-based representations could be more interpretable and generalize better with limited data, employing two innovative regularization methods. The first regularization separates foreground and background information's generative process via a unique mixture-of-parts formulation. Structural constraints are imposed on the parts using a weakly-supervised loss, guaranteeing that the mixture-of-parts for foreground and background entails soft, object-agnostic masks. The second regularization assumes the form of a distillation loss, ensuring the invariance of the learned parts to the incidental background correlations. Furthermore, we incorporate sparse and orthogonal constraints to facilitate learning high-quality part representations. By reducing the impact of incidental background correlations on the learned parts, we exhibit state-of-the-art (SoTA) performance on few-shot learning tasks on benchmark datasets, including MiniImagenet, TieredImageNet, and FC100. We also demonstrate that the part-based representations acquired through our approach generalize better than existing techniques, even under domain shifts of the background and common data corruption on the ImageNet-9 dataset. The implementation is available on GitHub: https://github.com/GauravBh1010tt/DPViT.git
Many datasets demonstrate a structural similarity by exhibiting "parts" or factors that reflect the underlying properties of the data [15, 18, 21, 28, 31, 43, 54]. Humans are efficient learners who represent objects based on their various traits or parts, such as a bird's morphology, color, and habitat characteristics. Part-based methods learn these explicit features from the data in addition to convolution and attention-based approaches (which only learn the internal representations), making them more expressive [7, 21, 41, 52, 54]. Most existing part-based methods focus on the unsupervised discovery of parts by modeling spatial configurations [14, 21, 52, 54, 61], while others use part localization supervision in terms of attribute vectors [25, 41, 50, 56] or bounding boxes [7]. Part-based methods come with a learnable part dictionary that provides a direct means of data abstraction and is effective in limited data scenarios, such as few-shot learning. Furthermore, the parts can be combined into hierarchical representations to form more significant components of the object
*First author; Email: gauravbhatt.cs.iitr@gmail.com
(b) ViT with parts
Figure 1: Impact of incidental correlations on the interpretability of part learners. We visualize the attention maps projected by the learned part dictionaries. Figure 1(b) illustrates the ViT-S backbone featuring a learnable part dictionary. However, it encounters difficulties in correctly identifying significant elements like the laptop, giving more attention to the background instead. In contrast, the proposed DPViT method successfully detects the most crucial parts of the image even in the presence of incidental correlations.
(c) Proposed - DPViT
description [50, 55]. Part-based learning methods offer advantages in terms of interpretability and generalization, particularly in safety-critical domains like healthcare, aviation and aerospace industry, transportation systems (e.g., railways, highways), emergency response, and disaster management. These fields often face challenges in collecting large training samples, making part-based learning methods valuable. Furthermore, the ability to explain decisions becomes crucial in these contexts.
Various studies have indicated that correlations between image background and labels can introduce biases during machine learning model training [29, 32-34, 45, 46, 53, 59]. These correlations exist because specific background configurations create shortcuts for the model during training [32, 53]. While background information is crucial for decision-making, imbalanced background configurations can create unintended correlations, leading to undesirable outcomes. These correlations negatively impact the interpretability and generalization of part-based learners. For instance, let's consider a scenario where a laptop, a charger, and an iPod are on a table. In one case, let's examine the situation without any context or background information. Without background information, the model may struggle to understand the purpose and significance of components or parts such as the laptop, charger, and iPod on a table. It could fail to differentiate between these objects or grasp their functionalities, resulting in a lack of recognition and understanding. Conversely, suppose the model is predominantly trained with examples of these items on tables. In that case, it may overly focus on the background elements, such as the table itself, disregarding the individual entities. Thus it becomes essential to handle incidental correlations of image background to achieve a balanced understanding. Existing part-based approaches fail to handle incidental correlations that arise due to specific background signals dominating the training data (analogous to Figure 1(b)), thereby hampering their interpretability and generalization to limited data scenarios.
Having high-quality part representations is essential for achieving proficiency in part-based learning. In this context, quality pertains to the sparsity and diversity of the learned parts. Sparsity ensures that only a few parts are responsible for a given image, as images comprise a small subset of parts. Conversely, diversity in part representations prevents the parts from converging into a single representation and facilitates each part's learning of a unique data representation. Although incidental correlations can negatively affect learned parts' quality, the quality of part-based methods is a significant challenge that all part learners face. While previous studies have addressed the issue of part quality [50, 52], their solutions do not assure the sparsity and diversity of the learned parts, thereby failing to guarantee high-quality parts.
To solve the aforementioned challenges, we introduce the Disentangled Part-based Vision Transformer (DPViT), which is trained to be resilient to the incidental correlations of image backgrounds and ensures high-quality part representations. We propose a disentangled pre-training phase, which separates the generative process of the foreground and background information through a unique mixture-of-parts formulation. We impose structural constraints on the parts to ensure that the mixture-of-parts for foreground and background includes soft, object-agnostic masks without requiring direct supervision of part localization. The parts are learned to be invariant to incidental correlations using a self-supervised distillation fine-tune phase. To address the issue of the quality of learned parts, we impose sparse and spectral penalties on the part matrix to guarantee the high quality of learned part representations. We include an assessment of the sparse and spectral norms of the part matrix as a quantitative indicator of the learned parts' quality. Finally, we evaluate the effectiveness of our method on benchmark few-shot datasets, including MiniImagenet [35], TieredImageNet [40],
and FC100 [6]. To demonstrate the robustness of our proposed method to incidental correlations of backgrounds and common data corruptions, we use the benchmark ImageNet-9 dataset [53].
Our key contributions can be summarized as follows:
- We propose regularization techniques to disentangle the generative process of foreground and background information through a mixture-of-parts formulation. Additionally, we employ a self-supervised distillation regularization to ensure that the learned parts remain invariant to incidental correlations of the image background.
- We ensure the high quality of learned parts by employing both sparsity and spectral orthogonal constraints over the part matrix. These constraints prevent the parts from degenerating and encourage a diverse range of part representations.
- Apart from our evaluation of standard few-shot benchmark datasets, we also analyze the impact of incidental correlations of background and typical data distortions by utilizing the benchmark ImageNet-9 dataset [53].
Part-based learning. The advantages of learning part-based representations have been extensively researched in image recognition tasks [15, 18, 31, 38, 43, 49, 52, 55, 56]. Earlier methods attempted to learn parts by defining a stochastic generative process [19, 43]. Part-based methods have been broadly classified into unsupervised and supervised categories. Unsupervised methods concentrate on learning the spatial relationship between parts by using part dictionaries without the supervision of part localization [14, 15, 20, 21, 28, 49, 52, 54, 61]. In contrast, supervised part-based methods rely on the supervision of part localization through attribute vectors [16, 27, 41, 50] or part bounding box information [7]. In the literature, parts are also referred to as concepts when supervision about part localization is involved [7, 41].
Discovering parts in an unsupervised way is a more challenging scenario that is applicable to most practical problems. Part dictionaries help data abstraction and are responsible for learning implicit and explicit data representations. For example, [28] clustered DCNN features using part-based dictionaries, while [26] introduced a generative dictionary-based model to learn parts from data. Similarly, [20] uses part-based dictionaries and an attention network to understand part representations. The ConstellationNet [54] and CORL [21] are some of the current methods from the constellation family [14, 61], and use dictionary-based part-prototypes for unsupervised discovery of parts from the data. Our approach also belongs to this category, as we only assume the part structure without requiring any supervision of part localization.
**Incidental correlations of image background.** Image backgrounds have been shown to affect a machine learning model's predictions, and at times the models learn by utilizing the incidental correlations between an image background and the class labels [4, 32, 42, 45, 46, 53, 58, 59]. To mitigate this issue, background researchers have used augmentation methods by altering background signals and using these samples during the training [32, 53, 58]. [53] performed an empirical study on the effect of image background on in-domain classification. They introduce several variants of background augmentations to reduce a model's reliance on the background. Similarly, [58] uses saliency maps of the image foreground to generate augmented samples to reduce the effect of the image background. Recently, [32] showed the effectiveness of background augmentation techniques for minimizing the effect of incidental correlations on few-shot learning.
Unlike these methods, our approach does not depend on background augmentations but instead learns the process of generating foreground and background parts that are disentangled. Furthermore, our proposed approach is not sensitive to the quality of foreground extraction and can operate with limited supervision of weak foreground masks.
Few-shot learning and Vision Transformers. In recent years, few-shot learning (FSL) has become the standard approach to evaluate machine learning models' generalization ability on limited data [1, 2, 5, 13, 24, 36, 44, 47, 54]. Vision transformers (ViT) [11] have demonstrated superior performance on FSL tasks [10, 22, 23, 30, 48, 51], and self-supervised distillation has emerged as a popular training strategy for these models [8, 17, 22, 30, 51]. A recent trend involves a two-stage procedure where models are pretrained via self-supervision before fine-tuning via supervised learning [8, 17, 22, 30, 60]. For example, [17] leverages self-supervised training with iBOT [60] as a pretext task, followed by inner loop token importance reweighting for supervised fine-tuning. HCTransformer [22] uses attribute surrogates learning and spectral tokens pooling for pre-training vision transformers
and performs fine-tuning using cascaded student-teacher distillation to improve data efficiency hierarchically. SMKD [30] uses iBOT pre-training and masked image modeling during fine-tuning to distill knowledge from the masked image regions.
Our approach employs a two-stage self-supervised training strategy of vision transformers akin to [22, 30, 60]. However, unlike existing methods that focus on generalization in few-shot learning, our primary objective is to learn part representations that are invariant to incidental correlations of the image background. Our training procedure is designed to facilitate learning disentangled and invariant part representations, which is impossible through existing two-stage self-supervised pipelines alone.
3 Problem Formulation and Preliminaries
In few-shot classification, the aim is to take a model trained on a dataset of samples from seen classes  $\mathcal{D}^{seen}$  with abundant annotated data, and transfer/adopt this model to classify a set of samples from a disjoint set of unseen/novel classes  $\mathcal{D}^{novel}$  with limited labeled data. Formally, let  $\mathcal{D}^{seen} = \{(\mathbf{x},y)\}$ , where  $\mathbf{x} \in \mathbb{X}$  corresponds to an image and  $y \in \mathbb{Y}^{seen}$  corresponds to the label among the set of seen classes. We also assume that during training, we have limited supervision of class-agnostic foreground-background mask  $(\mathcal{M}_f, \mathcal{M}_b)$  for regularization during training, which can be easily obtained by any weak foreground extractor as a preprocessing step (following [53]). Please note that no mask information is required for  $\mathcal{D}^{novel}$  at inference.
We follow the work of [9, 60] on self-supervised training of ViTs to design our pretrain phase. During training, we apply random data augmentations to generate multiple views of the a given image  $x^{v} \in \mathcal{D}^{seen}$ . These views are then fed into both the teacher and student networks. Our student network, with parameters  $\theta_{s}$ , includes a ViT backbone encoder and a projection head  $\phi_{s}$  that outputs a probability distribution over K classes. The ViT backbone generates a [cls] token, which is then passed through the projection head. The teacher network, with parameters  $\theta_{t}$ , is updated using Exponentially Moving Average (EMA) and serves to distill its knowledge to the student by minimizing the cross-entropy loss over the categorical distributions produced by their respective projection heads.
$$
\mathcal {L} _ {c l s} = \mathbb {E} _ {(\mathbf {x}, y) \sim \mathcal {D} ^ {s e e n}} \mathcal {L} _ {c e} \left(\mathcal {F} _ {\phi} ^ {t} \left(\mathcal {F} _ {\theta} ^ {t} \left(\mathbf {x} ^ {\mathbf {1}}\right)\right), \mathcal {F} _ {\phi} ^ {s} \left(\mathcal {F} _ {\theta} ^ {s} \left(\mathbf {x} ^ {\mathbf {2}}\right)\right)\right). \tag {1}
$$
For inference, we use the standard  $M$ -way,  $N$ -shot classification by forming tasks  $(\mathcal{T})$ , each comprising of support set  $(\mathcal{S})$  and query set  $(\mathcal{Q})$ , constructed from  $\mathcal{D}^{novel}$ . Specifically, a support set consists of  $M \times N$  images;  $N$  random images from each of  $M$  classes randomly chosen from  $\mathbb{Y}^{novel}$ . The query set consists of a disjoint set of images, to be classified, from the same  $M$  classes. Following the setup of [47], we form the class prototypes  $(\mathbf{c}_m)$  using samples from  $\mathcal{S}$ . The class prototypes and learned feature extractor  $(\mathcal{F}_{\theta})$  are used to infer the class label  $\hat{y}$  for an unseen sample  $\mathbf{x}^q \in \mathcal{Q}$  using a distance metric  $d$ .
$$
\hat {y} = \underset {m} {\arg \max } d \left(\mathcal {F} _ {\theta} \left(\mathbf {x} ^ {q}\right), \mathbf {c} _ {m}\right); \mathbf {c} _ {m} = \frac {1}{N} \sum_ {\left(\mathbf {x}, y _ {m}\right) \in \mathcal {S}} \mathcal {F} _ {\theta} (\mathbf {x}). \tag {2}
$$
4 Proposed Methodology
Given an input sample  $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ , and a patch size  $f$ , we extract flattened 2D patches  $\mathbf{x}_{\mathsf{f}} \in \mathbb{R}^{N \times (F^2 \cdot C)}$ , where  $N$  is the number of patches generated and  $(F, F)$  is the resolution of each image patch. Similar to a standard ViT architecture [11], we prepend a learnable [class] token and positional embeddings to retain the positional information. The flattened patches are passed to multi-head self attention layers and MLP blocks to generate a feature vector  $z_{p} = MSA(x_{f})$ .
Next, we define the parts as part-based dictionaries  $\mathbf{P} = \{\mathbf{p}_k\in \mathbb{R}^{F^2\cdot C}\}_{k = 1}^K$ , where  $\mathbf{p}_k$  denotes the part-vector for the part indexed as  $k$ . The part-matrix  $(\mathbf{P})$  is initialized randomly and is considered a trainable parameter of the architecture. Note that the dimension of each part-vector is equal to the dimension of flattened 2D patches, which is  $F^2\cdot C$ . For each part  $\mathbf{p}_k$ , we compute a distance map  $\mathbf{D}^k\in \mathbb{R}^N$  where each element in the distance map is computed by taking dot-product between the part  $\mathbf{p}_k$  and all the  $N$  patches:  $\mathbf{D}^k = \mathbf{x_f}\cdot \mathbf{p}_k$ .
Using the distance maps  $\mathbf{D} \in \mathbb{R}^{N \times K}$ , we introduce a multi-head cross-attention mechanism and compute the feature vector:  $z_{d} = MCA(F_{\psi}(\mathbf{D}))$ , where  $F_{\psi}$  is an MLP layer which upsamples  $\mathbf{D}: K \to F^2 \cdot C$ . The cross-attention layer shares a similar design to self-attention layers; the only
Figure 2: Overview of proposed architecture - DPViT. We employ a learnable part dictionary to generate a formulation incorporating foreground and background information. The spatial distance maps, computed by the part dictionary, are utilized to determine the mixture of latent codes for foreground and background. Our transformer encoder comprises multi-head self-attention (MSA) and multi-head cross-attention (MCA) layers. The MSA layer takes embedded patches as input, while the MCA layers utilize the distance maps as input.
difference is the dimensions of input distance maps. The cross-attention helps contextualize information across part-dictionary and the spatial image regions and provides complementary properties to MSA layers. (Please refer to Appendix for experiments on complementary properties of MSA and MCA).
Finally, we add the output feature vectors of  $MSA$  and  $MCA$  to form the feature extractor  $\mathcal{F}_{\theta}$  defined in Eqn 1:
$$
\mathcal {F} _ {\theta} = \left[ z _ {p} \oplus z _ {d} \right] \tag {3}
$$
Disentanglement of foreground-background space using mixture-of-parts. We start by dividing the parts-matrix  $\mathbf{P} \in \mathbb{R}^{K \times F^2 \cdot C}$  into two disjoint sets: foreground set  $\mathbf{P}_{\mathbf{F}} \in \mathbb{R}^{n_f \times C}$  and background set  $\mathbf{P}_{\mathbf{B}} \in \mathbb{R}^{n_b \times C}$ , such that  $K = n_f + n_b$ .
Next, we construct latent variables to aggregate the foreground and background information using a mixture-of-parts formulation over the computed distance maps  $\mathbf{D}$ :
$$
L _ {F} = \sum_ {k \in n _ {f}} \alpha_ {k} \mathbf {D} ^ {k} + \delta_ {f}; L _ {B} = \sum_ {k \in n _ {b}} \beta_ {k} \mathbf {D} ^ {k} + \delta_ {b} \tag {4}
$$
where,  $\alpha_{k}$  and  $\beta_{k}$  are the learnable weights given to the  $k^{th}$  part-vector in the corresponding mixture, whereas  $\delta_f$  and  $\delta_b$  are Gaussian noises sampled from  $\mathcal{N}(0,1)$ . The Gaussian noise is added to the latent codes to ensure that mixture-of-parts are robust to common data distortions. Please note that the purpose of Gaussian noise is not to induce variability in the latent codes, as foreground information for a given image is deterministic.
Finally, our disentanglement regularization takes the form of an alignment loss between the latent codes and the class-agnostic foreground-background masks:
$$
\mathcal {L} _ {m i x} = \left| \left| \mathcal {I} \left(L _ {F}\right) - \mathcal {M} _ {f} \right| \right| _ {2} + \left| \left| \mathcal {I} \left(L _ {B}\right) - \mathcal {M} _ {b} \right| \right| _ {2} \tag {5}
$$
where,  $\mathcal{I}(L)$  is the bilinear interpolation of a given latent code  $L$  to the same size as  $\mathcal{M}$ .
During the architectural design phase, we employ distinct components (P) for each encoder block. Consequently,  $z_{p}$  and  $z_{d}$  are calculated in an iterative manner and subsequently transmitted to the subsequent encoder block. Regarding the computation of  $L_{mix}$ , we utilize the distance maps  $\mathbf{D}$  from the concluding encoder block. While it's feasible to calculate  $L_{mix}$  iteratively for each block and then aggregate them for a final  $L_{mix}$  computation, our observations indicate that this approach amplifies computational expenses and results in performance deterioration. As a result, we opt to compute  $L_{mix}$  using the ultimate encoder block.
Learning high-quality part representations. A problem with minimizing the mixture objective defined in Eqn 5 is that it may cause the degeneration of parts, thereby making the part representations less diverse. One solution is to enforce orthogonality on the matrix  $\mathbf{P}^{m\times n}$  by minimizing  $||\mathbf{P}^T\mathbf{P} - \mathbf{I}||$  similar to [50]. However, the solution will result in a biased estimate as  $m < n$ ; that is, the number of
parts  $(K)$  is always less than the dimensionality of parts  $(F^2\cdot C)$ . In our experiments, we observed that increasing  $K$  beyond a certain threshold degrades the performance as computational complexity increases, and is consistent with the findings in [54]. (Please refer to our Appendix section for experiments on the different values of  $K$ ). To minimize the degeneration of parts, we design our quality assurance regularization by minimizing the spectral norm of  $\mathbf{P}^T\mathbf{P} - \mathbf{I}$ , and by adding  $L_{1}$  sparse penalty on the part-matrix  $\mathbf{P}$ . The spectral norm of  $\mathbf{P}^T\mathbf{P} - \mathbf{I}$  has been shown to work with over-complete  $(m < n)$  and under-complete matrices  $(m\geq n)$  [3].
$$
\mathcal {L} _ {Q} \left(\lambda_ {s}, \lambda_ {o}\right) = \lambda_ {s} | | \mathbf {P} | | _ {1} + \lambda_ {o} \left[ \sigma \left(\mathbf {P} _ {\mathbf {F}} \cdot \mathbf {P} _ {\mathbf {F}} ^ {T} - \mathbf {I}\right) + \sigma \left(\mathbf {P} _ {\mathbf {B}} \cdot \mathbf {P} _ {\mathbf {B}} ^ {T} - \mathbf {I}\right) \right] \tag {6}
$$
where  $\mathbf{I}$  is the identity matrix,  $\lambda_{s}$  and  $\lambda_{o}$  are the regularization coefficients for sparsity and orthogonality constraints.  $\sigma (\mathbf{P})$  is the spectral norm of the matrix  $\mathbf{P}$  which is computed using the scalable power iterative method described in [3].
Disentangled Pretraining Objective. We pretrain DPViT using the following loss function:
$$
\mathcal {L} _ {P T} = \lambda_ {c l s} \mathcal {L} _ {c l s} + \lambda_ {m i x} \mathcal {L} _ {m i x} + \mathcal {L} _ {Q} (\lambda_ {s}, \lambda_ {o}) \tag {7}
$$
where  $\lambda_{cls},\lambda_{mix},\lambda_s$  , and  $\lambda_{o}$  are the weights given to each loss term and are tuned on the validation set.
Figure 3: Invariant fine-tuning of DPViT via distillation framework.
4.1 Invariant fine-tuning