SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Avani Gupta $^{1,2\dagger}$  Saurabh Saini $^{2}$  P J Narayanan $^{2}$
$^{1}$ M42, UAE  $^{2}$ IIIT Hyderabad, India
{avani.gupta, saurabh.saini}@research.iit.ac.in
pjin@iiit.ac.in
Humans use abstract concepts for understanding instead of hard features. Recent interpretability research has focused on human-centered concept explanations of neural networks. Concept Activation Vectors (CAVs) estimate a model's sensitivity and possible biases to a given concept. In this paper, we extend CAVs from post-hoc analysis to ante-hoc training in order to reduce model bias through fine-tuning using an additional Concept Loss. Concepts were defined on the final layer of the network in the past. We generalize it to intermediate layers using class prototypes. This facilitates class learning in the last convolution layer which is known to be most informative. We also introduce Concept Distillation to create richer concepts using a pre-trained knowledgeable model as the teacher. Our method can sensitize or desensitize a model towards concepts. We show applications of concept-sensitive training to debias several classification problems. We also use concepts to induce prior knowledge into IID, a reconstruction problem. Concept-sensitive training can improve model interpretability, reduce biases, and induce prior knowledge. Please visit https://avani17101.github.io/Concept-Distillation/ for code and more details.
EXplainable Artificial Intelligence (XAI) methods are useful to understand a trained model's behavior [74]. They open the black box of Deep Neural Networks (DNNs) to enable post-training identification of unintended correlations or biases using similarity scores or saliency maps. Humans, however, think in terms of abstract concepts, defined as groupings of similar entities [27]. Recent efforts in XAI have focused on concept-based model explanations to make them more aligned with human cognition. Kim et al. [32] introduce Concept Activation Vectors (CAVs) using a concept classifier hyperplane to quantify the importance given by the model to a particular concept. For instance, CAVs can determine the model's sensitivity on 'striped-ness' or 'dotted-ness' to classify Zebra or Cheetah using user-provided concept samples. They measure the concept sensitivity of the model's final layer prediction with respect to intermediate layer activations (outputs). Such post-hoc analysis can evaluate the transparency, accountability, and reliability of a learned model [14] and can identify biases or unintended correlations acquired by the models via shortcut learning [32, 5].
The question we ask in this paper is: If CAVs can identify and quantify sensitivity to concepts, can they also be used to improve the model? Can we learn less biased and more human-centered models? In this paper, we extend CAVs to ante-hoc model improvement through a novel concept loss to desensitize/sensitize against concepts. We also leverage the broader conceptual knowledge of a large pre-trained model as a teacher in a concept distillation framework for it.
†Work done while at IIIT Hyderabad
Figure 1: Overview of our approach: The generic conceptual knowledge of a capable teacher can be distilled to a student for performance improvement through bias removal and prior induction.
Figure 2: Categorization of related works with ours highlighted.
XAI has been used for ante-hoc model improvement during training [35, 30, 4]. They typically make fundamental changes to the model architecture or need significant concept supervision, making extensions to other applications difficult. For example, Koh et al. [35] condition the model by first predicting the underlying concept and then using it for class prediction. Our method can sensitize or desensitize the model to user-defined concepts without modifications to the architecture or direct supervision.
Our approach relies on the sensitivity of a trained model to human-specified concepts. We want the model to be sensitive to relevant concepts and indifferent to others. For instance, a cow classifier might be focusing excessively on the grass associated with cow images. If we can estimate the sensitivities of the classifier to different concepts, we can steer it away from irrelevant concepts. We do that using a concept loss term  $L_{C}$  and fine-tuning the trained base model with it. Since the base models could be small and biased in different ways, we use concept distillation using a large, pre-trained teacher model that understands common concepts better.
We also extend concepts to work effectively on intermediate layers of the model, where the sensitivity is more pronounced. Kim et al. [32] measure the final layer's sensitivity to any intermediate layer outputs. They ask the question: if any changes in activations are done in the intermediate layer, what is its effect on the final layer prediction? They used the final layer's loss/logit to estimate the model sensitivity as their interest was to study concept sensitivities for interpretable model prediction. We, on the other hand, aim to fine-tune a model by (de)sensitizing it towards a given concept which may be strongest in another layer [3]. Thus, it is crucial for us to measure the sensitivity in any layer by evaluating the effect of the changes in activations in one intermediate layer on another. We employ prototypes or average class representations in that layer for this purpose. Prototypes are estimated by clustering the class sample activations [12, 72, 41, 49, 48, 63, 43]. Our method, thus, allows intervention in any layer.
In this paper, we present a simple but powerful framework for model improvement using concept loss and concept distillation for a user-given concept defined in any layer of the network. We leverage ideas from post-hoc global explanation techniques and use them in an ante-hoc setting by encoding concepts as CAVs via a teacher model. Our method also admits sample-specific explanations via a local loss [52] along with the global concepts whenever possible. We improve state-of-the-art on classification problems like ColorMNIST and DecoyMNIST [51, 16, 52, 5], resulting in improved accuracies and generalization. We introduce and benchmark on a new and more challenging TextureMNIST dataset with texture bias associated with digits. We demonstrate concept distillation on two applications: (i) debiasing extreme biases on classification problems involving synthetic MNIST datasets [42, 16] and complex and sensitive age- vs.-gender bias in the real-world gender classification on BFFHQ dataset [34] and (ii) prior induction by infusing domain knowledge in the reconstruction problem of Intrinsic Image Decomposition (IID) [38] by measuring and improving disentanglement of albedo and shading concepts. To summarize, we:
- Extend CAVs from post-hoc explanations to ante-hoc model improvement method to sensitize/desensitize models on specific concepts without changing the base architecture.
- Extend the model CAV sensitivity calculation from only final layer to any layer and enhance it by making it more global using prototypes.
- Introduce concept distillation to exploit the inherent knowledge of large pretrained models as a teacher in concept definition.
- Benchmark results on standard biased MNIST datasets and on a challenging TextureMNIST dataset that we introduce.
- Show application on a severely biased classification problem involving age bias.
- Show application beyond classification to the challenging multi-branch Intrinsic Image Decomposition problem by inducing human-centered concepts as priors. To the best of our knowledge, this is the first foray of concept-based techniques into non-classification problems.
The related literature is categorized in Fig. 2. Post-hoc (after training) explainability methods include Activation Maps Visualization [59], Saliency Estimation [54], Model Simplification [69], Model Perturbation [17], Adversarial Exemplar Analysis [22], etc. See recent surveys for a comprehensive discussion [74, 15, 45, 65, 2]. Different concept-based interpretability methods are surveyed by Hitzler and Sarker [27], Schwalbe [58], Holmberg et al. [28].
The ante-hoc model improvement techniques can be divided into zero-shot [70], few-shot [67], and multi-shot (domain-adaptation [66]) categories based on the amount of test distribution needed. In zero-shot training, the model never sees the distribution of test samples [70], while in few-shot, the model has access to some examples from the distribution of the test set [67]. Our method works with abstract concept sets (different than any test sample), being essentially a zero-shot method but can take advantage of few-shot examples, if available, as shown in results on BFFHQ subsection 4.1. We restrict the discussion to the most relevant model improvement methods.
Ross et al. [52] penalize a model if it does not prefer the Right answers for Right Reasons (RRR) by using explanations as constraints. EG Erion et al. [16] augment gradients with a penalizing term to match with user-provided binary annotations. More details are available in relevant surveys [18, 68, 24, 8].
Based on the scope of explanation, the existing XAI methods can be divided into global and local methods [74]. Global methods [32, 20] provide explanations that are true for all samples of a class (e.g. 'stripiness' concept is focused for the zebra class or 'red color' is focused on the prediction of zero) while local methods [1, 59, 60, 53] explain each of samples individually often indicating regions or patches in the image which lead the model to prediction (example, this particular patch (red) in this image was focused the most for prediction of zero). Kim et al. [32] quantify concept sensitivity using CAVs followed by subsequent works [62, 73, 57]. Kim et al. [32] sample sensitivity is local (sample specific) while they aggregate class samples for class sample sensitivity estimation, which makes their method global (across classes). We enhance their local sample sensitivity to the global level via using prototypes that capture class-wide characteristics by definition [12].
Only a few concept-oriented deep learning methods train neural networks with the help of concepts [56, 55]. Concept Bottleneck Models (CBMs) [35], Concept-based Model Extraction (CME) [30], and Self-Explaining Neural Networks (SENNs) [4] predict concepts from inputs and use them to infer the output. CBMs [35] require concept supervision, but CME [30] can train in a partially supervised manner to extract concepts combining multiple layers. SENN [4] learns interpretable basis concepts by approximating a model with a linear classifier. These methods require architectural modifications of the base model to predict concepts. Their primary goal is interpretability, which differs from our model improvement goal using specific concepts without architectural changes. Closest to our work is CIArC [5], which leverages CAVs to manipulate model activations using a linear transformation to remove artifacts from the final results. Unlike them, our method trains the DNN to be sensitive to any specific concept, focusing on improving generalizability for multiple applications.
DFA [40] and EnD [64] use bias-conflicting or out-of-distribution (OOD) samples for debiasing using few-shot generic debiasing. DFA uses a small set of adversarial samples and separate encoders for learning disentangled features for intrinsic and biased attributes. EnD uses a regularization strategy to prevent the learning of unwanted biases by inserting an information bottleneck using pre-known bias types. It is to be noted that both DFA and EnD are neither zero-shot nor interpretability-based methods. In comparison, our method works in both known bias settings using abstract user-provided
Figure 3:  $v_{c}^{l}$  is calculated as normal to the separating hyperplane of concept set activations (textures  $C$  vs. random set  $C'$  here). A model biased towards  $C$  will have its class samples's loss gradient  $\nabla L_{p}$  along  $v_{c}^{l}$  (measured by sensitivity  $S_{C,k,l}(x)$ ). To desensitize the model for  $C$ , we perturb  $\nabla L_{p}$  to be parallel to the decision boundary by minimizing the cosine of the projection angle.
Figure 4: Datasets used: ColorMNIST (top row), TextureMNIST (next row), DecoyMNIST (third row), and BFFHQ (bottom rows). Concepts used include color, textured and gray patches, and bias-conflicting samples shown on the right.
concept sets and unknown bias settings using the bias-conflicting samples. Our Concept Distillation method is a concept sensitive training method for induction of concepts into the model, with debiasing being an application. Intrinsic Image Decomposition (IID) [11, 19] involves decomposing an image into its constituent Reflectance  $(R)$  and Shading  $(S)$  components [38, 46] which are supposed to be disentangled [46, 23].
We use CAVs for inducing R and S priors in a pre-trained SOTA IID framework [44] with improved performance. To the best of our knowledge, we are the first to introduce prototypes for CAV sensitivity enhancement. However, prototypes have been used in the interpretability literature before to capture class-level characteristics [33, 71, 31] and also have been used as pseudo-class labels before [12, 72, 41, 49, 48, 63, 43]. Some recent approaches [29, 61] use generative models to generate bias-conflicting samples (e.g. other colored zeros in ColorMNIST) and train the model on them to remove bias. Specifically, Jain et al. [29] use SVMs to find directions of bias in a shared image and language space of CLIP [50] and use Dall-E on discovered keywords of bias to generate bias-conflicting samples and train the model on them. Song et al. [61] use StyleGAN to generate bias-conflicting samples. Training on bias-conflicting samples might not always be feasible due to higher annotation and computation costs.
One significant difference between such methods [29, 61] is that they propose data augmentation as a debiasing strategy, whereas we directly manipulate the gradient vectors, which is more interpretable. Moayeri et al. [47] map the activation spaces of two models using the CLIP latent space similarity. Due to the generality of the CLIP latent space, this approach is helpful to encode certain concepts like 'cat, dog, man,' but it is not clear how it will work on abstract concepts with ambiguous definitions like 'shading' and 'reflectance' as seen in the IID problem described above.
3 Concept Guidance and Concept Distillation
Concepts have been used to explain model behavior in a post-hoc manner in the past. Response to abstract concepts can also demonstrate the model's intrinsic preferences, biases, etc. Can we use concepts to guide the behavior of a trained base model in desirable ways in an ante-hoc manner? We describe a method to add a concept loss to achieve this. We also present concept distillation as a way to take advantage of large foundational models with more exposure to a wide variety of images.
3.1 Concept Sensitivity to Concept Loss
Building on Kim et al. [32], we represent a concept  $C$  using a Concept Activation Vector (CAV) as the normal  $v_{C}^{l}$  to a linear decision boundary between concept samples  $C$  from others  $C'$  in a layer  $l$  of the model's activation space (Fig. 3). The model's sensitivity  $S_{C,l}(\boldsymbol{x}) = \nabla L_{o}\left(f_{l}(\boldsymbol{x})\right)\cdot v_{C}^{l}$  to  $C$  is
the directional derivative of final layer loss  $L_{o}$  for samples  $x$  along  $v_{C}^{l}$  [32]. The sensitivity score quantifies the concept's influence on the model's prediction. A high sensitivity for color concept may indicate a color bias in the model.
These scores were used for post-hoc analysis before ([32]). We use them ante-hoc to desensitize or sensitize the base model to concepts by perturbing it away from or towards the CAV direction (Fig. 3). The gradient of loss indicates the direction of maximum change. Nudging the gradients away from the CAV direction encourages the model to be less sensitive to the concept and vice versa. For this, we define a concept loss  $L_{C}$  as the absolute cosine of the angle between the loss gradient and the CAV direction
$$
L _ {C} (\boldsymbol {x}) = | \cos (\nabla L _ {o} (f _ {l} (\boldsymbol {x})), \boldsymbol {v} _ {C} ^ {l}) |, \tag {1}
$$
which is minimized when the CAV lies on the classifier hyperplane (Fig. 3). We use the absolute value to not introduce the opposite bias by pushing the loss gradient in the opposite direction. A loss of  $(1 - L_{C}(x))$  will sensitize the model to  $C$ . We fine-tune the trained base model for a few epochs using a total loss of  $L = L_{o} + \lambda L_{C}$  to desensitize it to concept  $C$ , where  $L_{o}$  the base model loss.
3.2 Concepts using Prototypes
Concepts can be present in any layer  $l$ , though the above discussion focuses on the sensitivity calculation of the final layer using model loss  $L_{o}$ . The final convolutional layer is proven to learn concepts better than other layers [3]. We can estimate the concept sensitivity of any layer using a loss for that layer. How do we get a loss for an intermediate layer, as no ground truth is available for it?
Class prototypes have been used as pseudo-labels in intermediate layers before [12, 72, 41]. We adapt prototypes to define a loss in intermediate layers. Let  $f_{l}(x)$  be the activation of layer  $l$  for sample  $x$ . We group the  $f_{l}(x)$  values of the samples from each class into  $\mathbf{K}$  clusters. The cluster centers  $P_{i}$  together form the prototype for that class. We then define prototype loss for each training sample  $x$  using the prototype corresponding to its class as
$$
L _ {p} (x) = \frac {1}{K} \sum_ {k = 1} ^ {K} \| f _ {l} (x) - P _ {k} \| ^ {2}. \tag {2}
$$
We use  $L_{p}$  in place of  $L_{o}$  in Eq. 1 to define the concept loss in layer  $l$ . The prototype loss facilitates the use of intermediate layers for concept (de)sensitization. Experiments reported in Tab. 1 confirm the effectiveness of  $L_{p}$ . Prototypes also capture sample sensitivity at a global level using all samples of a class beyond the sample-specific levels. We update the prototypes after a few iterations as the activation space evolves. If  $P^{n}$  is the prototype at Step  $n$  and  $P^{c}$  the cluster centres using the current  $f_{l}(x)$  values, the next prototype is  $P_{k}^{n + 1} = (1 - \alpha)P_{k}^{n} + \alpha P_{k}^{c}$  for each cluster  $k$ .
3.3 Concept Distillation using a teacher
Concepts are learned from the base model in the above formulation. Base models may have wrong concept associations due to their training bias or limited exposure to concepts. Can we alleviate this problem using a larger model that has seen vast amounts of data as a teacher in a distillation framework?
We use the DINO [13], a self-supervised model trained on a large number of images, as the teacher and the base model as the student for concept distillation. The teacher and student models typically have different activation spaces. We map the teacher space to the student space before concept learning. The mapping uses an autoencoder [26] consisting of an encoder  $E_{M}$  and a decoder  $D_{M}$  (Fig. 5). As a first step, the autoencoder (Fig. 5) is trained to minimize the loss  $L_{D_M} + L_{E_M}$ .  $L_{D_M}$  is the pixel-wise L2 loss between the original  $(f_t)$  and decoded  $(\hat{f}_t)$  teacher activations and  $L_{E_M}$  is the pixel-wise L2 loss between the mapped teacher  $(\hat{f}_s)$  and the student  $(f_s)$  activations. The mapping is learned over the concept set of images  $C$  and  $C'$ . See the dashed purple lines in Fig. 5.
Next, we learn the CAVs in the distilled teacher space  $\hat{f}_s$ , keeping the teacher, student, and mapping modules fixed. This is computationally light as only a few (50-150) concept set images are involved. The learned CAV is used in concept loss given in Eq. 1. Please note that  $E_M$  is used only to align the two spaces and can be a small capacity encoder, even a single layer trained in a few epochs.
Figure 5: Our framework comprises a concept teacher and a student classifier and has the following four steps: 1) Mapping teacher space to student space for concepts  $C$  and  $C'$  by training an autoencoder  $E_M$  and  $D_M$  (dotted purple lines); 2) CAV  $(v_c^l)$  learning in mapped teacher space via a linear classifier LC (dashed blue lines); 3) Training the student model with Concept Distillation (solid orange lines): We use  $v_c^l$  and class prototypes loss  $L_p$  to define our concept distillation loss  $L_c$  and use it with the original training loss  $L_o$  to (de)sensitize the model for concept  $C$ ; 4) Testing where the trained model is applied (dotted red lines)
Algorithm 1 Concept Distillation Pipeline
Given: A pretrained student model which is to be fine-tuned for concept (de)sensitization and a pretrained teacher model which will be used for concept distillation. Known Concepts  $C$  and negative counterparts (or random samples)  $C'$ , student training dataset  $\mathcal{D}$ , and student bottleneck layer  $l$ , #iterations to recalculate CAVs cav_update_frequency, #iterations to update prototypes proto_update_frequency.
1: Concept Distillation:
2: For all class samples in  $\mathcal{D}$ , estimate class prototypes  $P_{k\in \{0,K\}}^{0}$  with K-means.
3: Current iteration  $n = 0$ , initial prototypes  $P^0 = P^c$ .
4: While not converge do:
5: If  $n = 0$  or (update_cavs and  $n\%$  cav_update_frequency  $= 0$ ) then:
6: Learn Mapping module:
7: Forward pass  $x \in C \cup C'$  from Teacher and Student to get their concept activations  $f_t$  and  $f_s$ .
8: Learn the mapping module as autoencoders  $E_{M}$  and  $D_{M}$ .
9: CAV learning in mapped teacher's space:
10:  $\mathbf{v}_C^l$  learned by binary linear classifier as normal to decision boundary of  $E_M(f_t(x))$  for  $x\in C$  vs  $E_{M}(f_{t}(x^{\prime}))$  for  $x^{\prime}\in C^{\prime}$
11: If  $n\%$  proto_update_freqency = 0 and  $n \neq 0$  then:
12: Estimate new class prototypes  $P_{k\in \{0,K\}}^{c}$  with K-means.
13: Weighted Proto-type  $P_{k}^{n + 1} = (1 - \alpha)P_{k}^{n} + \alpha P_{k}^{c}$ .
14: Else:
15:  $P_{k}^{n + 1} = P_{k}^{n}$
16: Train student with loss  $L_{C} + L_{o}$
17:  $n + = 1$
We demonstrate the impact of the concept distillation method on the debiasing of classification problems as well as improving the results of a real-world reconstruction problem. Classification experiments cover two categories of Fig. 2: the zero-shot scenario, where no unbiased data is seen, and the few-shot scenario, which sees a few unbiased data samples. We use pre-trained DINO ViT-B8 transformer  $[13]^*$  as the teacher. We use the last convolution layer of the student model for concept distillation. The mapping module ( $< 120K$  parameters,  $< 10\mathrm{MB}$ ), CAV estimations (logistic regression,  $< 1\mathrm{MB}$ ), and prototype calculations all complete within a few iterations of training, taking 15-30 secs on a single 12GB Nvidia 1080 Ti GPU*. Our framework is computationally light. In our experimentation reported below, we fix the CAVs as initial CAVs and observe similar results as
*More details in supplementary
Table 1: Main components of our method shown on ColorMNIST dataset: Ours usage of teacher and proto-types yields the best performance.
Teacher?	Prototype?	Accuracy
X	X	9.96
X	✓	26.97
✓	X	30.94
✓	✓	50.93
Table 2: TCAV scores of bias concept: Our concept sensitive training significantly decreases the sensitivity of model towards bias.
Dataset	Concept	Base Model	Ours
ColorMNIST	Color	0.52	0.21
DecoyMNIST	Spatial patches	0.57	0.45
TextureMNIST	Textures	0.68	0.43
BFFHQ	Age	0.78	0.13
varying them. Also, our concept sensitive fine-tuning method converges quickly, and hence updating CAVs in every few iterations does not help*. Where relevant, we report average accuracy over five training runs with random seeds ( $\pm$  indicates variance).
4.1 Concept Sensitive Debiasings