# CoLiDR: Concept Learning using Aggregated Disentangled Representations

Sanchit Sinha<EMAIL_ADDRESS>University of VirginiaCharlottesvilleVAUSA
,  Guangzhi Xiong<EMAIL_ADDRESS>University of
VirginiaCharlottesvilleVAUSA and  Aidong Zhang<EMAIL_ADDRESS>University
of VirginiaCharlottesvilleVAUSA

(2024)

###### Abstract.

Interpretability of Deep Neural Networks using concept-based models offers a
promising way to explain model behavior through human-understandable concepts.
A parallel line of research focuses on disentangling the data distribution
into its underlying generative factors, in turn explaining the data generation
process. While both directions have received extensive attention, little work
has been done on explaining concepts in terms of generative factors to unify
mathematically disentangled representations and human-understandable concepts
as an explanation for downstream tasks. In this paper, we propose a novel
method CoLiDR - which utilizes a disentangled representation learning setup
for learning mutually independent generative factors and subsequently learns
to aggregate the said representations into human-understandable concepts using
a novel aggregation/decomposition module. Experiments are conducted on
datasets with both known and unknown latent generative factors. Our method
successfully aggregates disentangled generative factors into concepts while
maintaining parity with state-of-the-art concept-based approaches.
Quantitative and visual analysis of the learned aggregation procedure
demonstrates the advantages of our work compared to commonly used concept-
based models over four challenging datasets. Lastly, our work is generalizable
to an arbitrary number of concepts and generative factors - making it flexible
enough to be suitable for various types of data.

disentanglement, xai, concept learning, generalization

††journalyear: 2024††copyright: rightsretained††conference: Proceedings of the
30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining; August
25–29, 2024; Barcelona, Spain††booktitle: Proceedings of the 30th ACM SIGKDD
Conference on Knowledge Discovery and Data Mining (KDD ’24), August 25–29,
2024, Barcelona, Spain††doi: 10.1145/3637528.3671938††isbn:
979-8-4007-0490-1/24/08††ccs: Mathematics of computing Probability and
statistics

##  1\. Introduction

The increasing proliferation of Deep Neural Networks (DNNs) has revolutionized
multiple diverse fields of research such as vision, speech, and language (He
et al., 2016; Vaswani et al., 2017). Given the black-box nature of DNNs,
explaining DNN predictions has been an active field of research that attempts
to impart transparency and trustworthiness in their decision-making processes.
Recent research has categorized explainability into progressively increasing
levels of granularity. The most fine-grained approaches attempt to assign
importance scores to the raw features (e.g. pixels) extracted from the data,
while less granular approaches assign importance scores to data points (sets
of features). Explaining DNNs using concepts provides the highest level of
abstraction, as concepts are high-level entities shared among multiple similar
data points that are aligned with human understanding of the task at hand.
This makes concept explanation much more global in nature. Many recent
approaches for concept-based explanations have attempted to either 1) infer
concepts post-hoc from trained models (Kim et al., 2018) or 2) design
inherently explainable concept-based models (Alvarez Melis and Jaakkola,
2018), such as the concept bottleneck model (CBM) (Koh et al., 2020).

A parallel field of research on disentanglement representation learning
(Bengio et al., 2013; Higgins et al., 2017; Kim and Mnih, 2018; Chen et al.,
2018) attempts to learn a low-dimensional data representation where each
dimension independently represents a distinct property of the data
distribution. These approaches learn mutually independent generative factors
of data by estimating their probability distribution from observed data. Once
the probability distribution of the generative factors is estimated, a given
sample can be theoretically decomposed and re-generated from its generative
factors. Due to their ability to uncover the underlying generative factors,
disentanglement approaches are considered highly interpretable.

Figure 1. Schematic overview of the proposed CoLiDR approach. The input data
distribution $\mathbf{X}$ is first disentangled into $k$ mutually independent
generative factors (GFs). Subsequently, the GFs are aggregated into concepts.
Note that the concepts are modeled as a set of annotated concepts
corresponding to concepts with annotation by humans and a separate set of
concepts that are useful for prediction, but are unannotated. Finally, the
concepts are utilized for predicting the task label $\mathbf{Y}$.

Present state-of-the-art approaches do not effectively unify disentangled
representation learning with concept-based approaches. Approaches like
GlanceNets (Marconato et al., 2022) attempt to align concepts with
disentanglement with strong assumptions, which are not valid for real-world
datasets. To address the above issues, in this paper, we propose Concept
Learning using Aggregated Disentangled Representations (CoLiDR), a self-
interpretable approach that combines disentangled representation learning with
concept-based explainability. Specifically, CoLiDR learns disentangled
generative factors using a disentangled representation learning module,
followed by the aggregation of learned disentangled representations into
human-understandable concepts using a novel aggregation/decomposition module
and subsequently a task prediction module that maps concepts to task labels.
Our experiments show that interventions on learned concept representations can
fix wrongly classified samples, which makes CoLiDR useful for model debugging.
Figure 1 demonstrates the schematic overview of CoLiDR, where the learned
disentangled generative factors from data are aggregated into concepts, which
are further utilized for task prediction. Specifically, our main contributions
are as follows: 1) Propose CoLiDR framework, a novel method to aggregate
disentangled representations into human understandable concepts, while
achieving at-par performance to other commonly utilized concept-based models
such as CBMs and GlanceNets, 2) Improve model debugability by test-time
interventions, and 3) Provide a flexible framework which can be generalized to
arbitrary number of generative factors and concepts.

##  2\. Related Work

###  2.1. Related Work on Disentangled Representation Learning

Disentangled representation learning has long been a fascinating study aimed
at separating distinct informational factors of variations in real-world data
(Bengio et al., 2013; Locatello et al., 2019). Due to its probabilistic
framework and flexibility to customize training objectives, the variational
autoencoder (VAE) (Kingma and Welling, 2014) is a commonly used architecture
in the study of disentangled representation learning, which can capture
different factors of variation with its encoder and decoder. Based on the
traditional VAE, Higgins et al. (2017) proposed a variant, $\beta$-VAE, which
introduces an additional hyperparameter to scale the importance of the
regularization term. By changing the weight of the regularization term,
$\beta$-VAE can control the trade-off between the reconstruction of inputs and
the disentanglement of latent variables. Instead of optimizing the
Kullback–Leibler (KL) divergence between the latent distribution and standard
Gaussian prior for disentanglement as done by VAE and $\beta$-VAE, FactorVAE
(Kim and Mnih, 2018) and $\beta$-TCVAE (Chen et al., 2018) further decompose
the regularization term and propose to directly penalize the total correlation
between latent variables, which are shown to better disentangle the variables.

While early studies in disentanglement representation learning attempted to
learn independent latent variables by modifying the training objective of VAE
(Higgins et al., 2017; Kim and Mnih, 2018; Chen et al., 2018), Locatello et
al. (2019) showed that it is impossible to learn identifiable disentangled
latent variables without any supervision. A series of subsequent studies were
then proposed to learn disentangled representations with better
identifiability using different kinds of supervision (Locatello et al., 2020;
Shu et al., 2020; Shen et al., 2022).

###  2.2. Related Work on Concept Based Explanations

Due to the black-box nature of deep learning models, various approaches have
been explored to provide explanations for the outcome of deep neural networks
(Koh et al., 2020; Pedapati et al., 2020; Jeyakumar et al., 2020; Heskes et
al., 2020; O’Shaughnessy et al., 2020). One important direction is to
interpret the models with intuitive and human-understandable concepts, which
are usually high-level abstractions of the input features. Various attempts
have been made to automatically learn the concepts for different tasks (Kim et
al., 2018; Ghorbani et al., 2019; Yeh et al., 2020; Wu et al., 2020; Goyal et
al., 2019). Among them, the concept bottleneck model (CBM) (Koh et al., 2020)
is a commonly used approach to incorporate the learning of concepts into deep
neural networks. By constructing a low-dimensional intermediate layer in the
models, CBMs are able to capture the high-level concepts that are related to
the downstream tasks. Numerous studies have been carried out to adapt CBMs for
tasks in various domains (Sawada and Nakamura, 2022; Jeyakumar et al., 2022;
Pittino et al., 2021; Bahadori and Heckerman, 2020; Marconato et al., 2022).
The most relevant works to our proposed approach are Concept Bottleneck Models
(Koh et al., 2020), Concept Embedding Models (Zarlenga et al., 2022), CLAP
(Taeb et al., 2022) and GlanceNets (Marconato et al., 2022), details of which
are discussed in the next Section.

##  3\. Unifying Concept-based Modeling with Variational Inference

Problem setup. The concept learning problem is characterized as a two-step
process (Koh et al., 2020) for a given training sample
$\\{\mathbf{x},\mathbf{y},\mathbf{C}\\}$. Given an observation $x$, a concept-
based model learns a function to map the input to its associated human-
understandable concepts $\mathbf{C}=\\{c_{1},\cdots,c_{N}\\}$. Subsequently, a
predictor will project the concept embeddings to $\hat{y}$, which is the
prediction for the label $y$ of the sample in a downstream task.

Assumptions. Following previous work in concept-based models (Koh et al.,
2020; Marconato et al., 2022), concept models follow two fundamental
assumptions as follows: 1) the data distribution in the input space can be
accurately mapped to the distributions of concepts in the latent space, and 2)
the concept scores are necessary and sufficient to predict labels of samples
in downstream tasks.

###  3.1. Variational Inference-based Concept-learning for Greater
Interpretablity

Utilizing disentanglement approaches to explain the data generation process
has been well studied recently (Kingma and Welling, 2014; Higgins et al.,
2017). A typical disentanglement system attempts to learn a fundamental
representation of mutually independent generative factors (GFs) usually
modeled as random variables from independent probability distributions. The
Variational Inference (VI) process utilizes the learned generative factors to
emulate the data generation procedure. Even though variational inference has
been successfully utilized for controlled data generation with success
(Creswell et al., 2018; Higgins et al., 2017), utilizing them to improve the
interpretability of concept-based models is still relatively underexplored.
Our work attempts to fill this gap. Below, we list the advantages of using
variational inference-based approaches for concept learning:

  * •

Provides a controlled data generation process: In addition to emulating the
data generation process, VIs offer control via interventions over GFs - a
desirable property for stakeholders of a concept-based model.

  * •

Captures the continuous space mappings: Several works (Margeloiu et al., 2021;
Sinha et al., 2023) have demonstrated that learning discrete mapping of input
samples to concepts is susceptible to fragility. VI smoothens both the input
and concept spaces to learn a continuous, well-trained and robust mapping
function.

  * •

Easier to isolate confounds: Another well-documented problem in concept-based
models is the encoding of non-related confounds in the concept representations
(Yuksekgonul et al., 2022). VI methods are successful in isolating confounds -
hence reducing the effect of encoding noise, spurious correlations, etc. in
the concept representations.

Figure 2. A schematic view of the underlying assumptions considered across
SOTA concept-based models, CBM/CEM, CLAP, GlanceNet, and CoLiDR. The circles
in blue represent directly observable attributes, input sample $X$, task label
$Y$ and representative concepts $C$. The red circles represent learned
representations.

###  3.2. Comparison with Existing Approaches

We begin the discussion by presenting an anti-causal (Kilbertus et al., 2018)
model to visualize the data generation process of our proposed approach and
compare it to multiple comparable recent state-of-the-art approaches in Figure
2. Note that we represent the task labels as $\mathbf{Y}$, the set of concepts
as $\mathbf{C}$, the disentangled generative factors as $\mathbf{Z}$, and the
data distribution as $\mathbf{X}$. The single-edge arrows model the dependency
between distributions while the double-ended arrows represent variational
inference - modeling the generative process.

Comparisons to CBM/CEM. Concept Bottleneck Models (Koh et al., 2020) and
Concept Embedding Models (Zarlenga et al., 2022) do not incorporate
variational inference. Nevertheless, they visualize the data-generative
process as being conditioned on the annotated concepts and also assume that
the task labels are entirely conditioned on the concepts.

Comparisons to CLAP. CLAP (Taeb et al., 2022) is one of the first works
utilizing variational inference to model the data generation process. CLAP
considers a part of the disentangled space representing relevant GFs to be
conditioned on the task labels and a part of it representing confounds to be
conditioned on independent normal distributions. However, CLAP does not
utilize human-annotated concepts - making the learned disentangled space
$\mathbf{Z}$ unidentifiable (Refer (Locatello et al., 2020)). This makes CLAP
not comparable to our approach.

Comparisons to GlanceNets. GlanceNets (Marconato et al., 2022), to our
knowledge, is the only existing approach that attempts to bridge the gap
between VI and concept learning. The GlanceNets approach incorporates a
Variational Autoencoder based disentangling mechanism to learn the generative
factors (GFs) of the distribution and formulate the concept learning problem
as a one-to-one mapping between a subset of generative factors (or
disentangled representations) and human-understandable concepts. As shown in
Figure 2, the annotated concepts directly supervise a part of the learned
disentangled space $\mathbf{Z}$. However, in doing so it makes an implicit
assumption that the human-understandable concepts act as generative factors of
the data distribution. This is a strong assumption for two distinct reasons.
Firstly, multiple generative factors can contribute as a constituent of a
particular concept and multiple concepts can share the same generative
factors. Secondly, human-annotated concepts are by definition, abstract and
high-level and do not necessarily model the fine-grained data generation
process. Utilizing abstract concepts to supervise the disentangling process
may undermine the disentanglement procedure, as disentangled representations
are supposed to capture the low-level GFs of the data distribution instead of
high-level abstract human concepts. The model is shown to be effective on two
synthetic datasets, dSprites and MPI-3D, which are procedurally generated
using carefully curated GFs that correspond to human-understandable concepts.
In addition, CoLiDR can effectively generalize to datasets where GFs are
completely unknown. Unannotated concepts. Note that we model the concept set
as a union of known, annotated concepts, and unannotated concepts. This
modeling approach is common in unsupervised concept learning approaches like
(Alvarez Melis and Jaakkola, 2018), but has not yet been studied for
Variational Inference based concept learning. The assumption is that there
exists some concepts relevant to the task prediction but are not manually
annotated.

Table 1 lists the differences between the discussed approaches based on - (i)
the Presence of VI, (ii) the Presence of unsupervised GFs, (iii) the Presence
of supervised concepts, and (iv) the Presence of unannotated concepts
(Discussed in Methodology).

Method | Variational
| Unsupervised | Supervised | Unannotated
---|---|---|---|---
| Inference | GF | Concepts | Concepts
CBMs/CEMs | ✗ | ✗ | ✗ | ✗
GlanceNet | ✔ | ✗* | ✔ | ✗
CLAP | ✔ | ✔ | ✗ | ✔
CoLiDR | ✔ | ✔ | ✔ | ✔

Table 1. A summary of salient features of our method as compared to the
baselines. The asterisk represents partial supervision over the generative
factors by GlanceNets.

###  3.3. Drawbacks of Supervising GFs with Concepts

Although it might be tempting to utilize concepts as GFs like what GlanceNets
(Marconato et al., 2022) did, it is dangerous for a variety of reasons. In
this section, we provide the drawbacks associated with this line of thought
with a concrete and practical example. Let us consider as a sample from the
CelebA dataset (Liu et al., 2015) which consists of facial photographs of
celebrities. Each image in the dataset is annotated with binary concepts such
as the presence of “black hair”, “blonde hair”, “wavy hair”, “straight hair”
which are easily understandable by humans - implying they are abstract.
Consider the assumption (fundamental in (Marconato et al., 2022)) that each of
these concepts also align with the generative factors of the data distribution
as well.

Limited concept abstraction. As it can be seen, the assumption is flawed - two
identical images with straight hair and different colors (Black and Blonde)
would be sampled from completely different distributions despite sharing the
same type of hair. Images of two people having the same hair color but with
wavy and straight hair respectively would also be sampled from completely
different distributions despite sharing the same color. Hence, there is no
fine-grained control over the actual generative factors. Furthermore, as
concepts such as “attractiveness” are very subjective and contain a large
number of constituent underlying distributions, it is much more effective to
understand constituent distributions than directly aligning attractiveness
with a single disentangled representation.

Limited intra-concept hierarchy modeling. By definition of disentanglement,
generative factors and concepts would be independent of each other. Any
hierarchical relationship cannot be encoded in such a space. For example, the
concept “black hair” is ideally sampled from mutually independent
distributions representing “blackness” and multiple other distributions
constituting in construction of hair themselves. As such, “blackness” can also
be combined with distributions modeling facial hairs for “beard” concepts.

Hence, instead of directly aligning the concepts with disentangled
representations, it is more explainable and reasonable to understand the
actual distributions constituting a concept. Relaxing the assumption of
aligning disentangled latent space representations with concepts helps us to
better capture and explain constituent distributions in each concept - in turn
improving explainability.

##  4\. Methodology

Figure 3. Architecture of the proposed CoLiDR approach. CoLiDR consists of
three modules - the Disentangled Representations Learning (DRL) Module which
learns disentangled generative factors (top), the Aggregation/Decomposition
Module which learns to aggregate the generative factors into concepts and
subsequently decompose them back into generative factors (bottom-left) and the
Task Learning module (bottom-right) which utilizes the concepts to perform
task label prediction.

We first provide a broad overview of the proposed approach CoLiDR in Section
4.1. Section 4.2 details the architectures utilized for computing disentangled
representations. Subsequently, Section 4.3 introduces the
Aggregation/Decomposition module and the task prediction network. Finally,
Section 4.4 describes the training procedure and additional Disentangled
Representation (DR) Consistency loss which regularizes the
Aggregation/Decomposition module to correctly learn representations to concept
mapping.

###  4.1. Overview of CoLiDR

CoLiDR consists of three distinct modules as shown in Figure 3. The first
module learns the disentangled generative factors from the input data. The
second module aggregates the disentangled representations and maps them to
human-understandable concepts. The third module maps the concepts to the task
label.

###  4.2. Disentangled Representations Learning (DRL) Module

The first step of the proposed approach involves learning disentangled
representations corresponding to the various generative factors in the data.
As depicted in Figure 1, the generative factors form the basis of the data
generation process itself. Our DRL module learns the disentangled generative
factors that can be used for data generation. Suppose $z$ is the given
embedding of the generative factors and $x$ is the corresponding generated
data. The underlying data generation process $\mathit{p}_{\theta}(x|z)$ can be
estimated using a $\beta$-VAE, which estimates the maximum likelihood using
variational inference. Following (Kingma and Welling, 2014), we maximize the
Evidence Lower Bound (ELBO) to model the posterior distribution
$q_{\phi}(z|x)$ and the distribution of data generation $p_{\theta}(x|z)$ as
detailed below.

(1) $\small
ELBO=-\beta~{}\mathbb{D}_{KL}(q_{\phi}(\mathbf{z}|\mathbf{x})~{}||~{}p(\mathbf{z}))+\mathbb{E}_{q_{\phi}(z|x)}[log~{}p_{\theta}(\mathbf{x}|\mathbf{z})]$

The $\mathbb{E}$ term in Formula 1 comprises the reconstruction loss between
the input $x$ and predicted reconstruction $\hat{x}$, usually taken as Mean
Square Error (MSE). The prior distribution $\mathit{p}$ is usually taken as a
standard Gaussian distribution which encourages the covariance matrix of the
learned distribution to be diagonal, enforcing independence constrain (Higgins
et al., 2017; Kingma and Welling, 2014). The tunable hyperparameter $\beta$
acts as a quantitative measure of the extent of disentanglement.

$\beta$-VAE estimates $q_{\phi}$ using an encoder that maps $x$ from the input
observation space in $\mathbb{R}^{d}$ to the disentangled representation space
in $\mathbb{R}^{k}$, where $d$ and $k$ are the dimensions of the input and
latent space respectively. The distribution $p_{\theta}$ is estimated using a
decoder that maps the disentangled representation space back to the
observation space ($\mathbb{R}^{k}\rightarrow\mathbb{R}^{d}$). Specifically,
the encoder encodes the input as an estimated mean $z_{\mu}\in\mathbb{R}^{k}$
and a standard deviation vector $z_{\sigma}\in\mathbb{R}^{k}$, from which the
latent representation ${z}$ is sampled from the Gaussian multivariate
distribution $\mathcal{N}(z_{\mu},\text{diag}(z_{\sigma}^{2}))$.

###  4.3. Aggregation/Decomposition Module

While the DRL module learns disentangled generative factors in the latent
space, the factors may be too fine-grained to be aligned with human
understanding. However, these disentangled generative factors can be
considered as an independent basis of human-understandable concepts as shown
in Figure 1.

####  4.3.1. Aggregation of Generative Factors into Concepts

We propose the Aggregation module which aggregates the disentangled underlying
generative factors of data into human-understandable concepts. Specifically,
given a latent representation ${z}=[z_{1},\cdots,z_{k}]^{\top}$, each concept
$c_{i}\in\\{c_{i}\\}_{i=1}^{N}$ can be considered as a combination of all
disentangled factors $z_{1},\cdots,z_{k}$. To increase the expressiveness of
our model in learning concepts from the generative factors while keeping each
concept as a linear aggregation of different factors for interpretability, we
propose to encode each generative factor $z_{j}(j\in\\{1,\cdots,k\\})$ with an
independent neural network $a_{j}$, which learns the non-linear mapping from
$z_{j}$ to $z^{\prime}_{j}$ for concept learning. Subsequently, our model
learns the linear combinations of components in
$z^{\prime}=[z^{\prime}_{1},\cdots,z^{\prime}_{k}]^{\top}$ to represent high-
level concepts $c_{1},\cdots,c_{N}$ that are relevant to the downstream task.
Given the posterior distribution $q_{\phi}(z|x)$ in $\beta$-VAE, the learned
concepts can be formulated as

(2) $c=f(\mathcal{A}(z)),$

where $z\sim q(z|x)$, $f$ is a one-layer fully connected model, and
$\mathcal{A}$ is defined as

(3)
$\small\mathcal{A}(z)=[z^{\prime}_{1},\cdots,z^{\prime}_{k}]^{\top}=[a_{1}(z_{1}),a_{2}(z_{2}),\cdots,a_{k}(z_{k})]^{\top}.$

####  4.3.2. Concept Level Supervision

Given the annotation of $n$ manually defined concepts, we can supervise the
learning of the first $n$ concepts to be aligned with human knowledge.
Formally, given the estimated scores $c_{1},\cdots,c_{n}$ and human
annotations $l_{1},\cdots,l_{n}$, the supervision can be performed by
minimizing

(4) $\mathcal{L}_{con}=\sum_{i=1}^{n}\textsf{BCE}(c_{i},l_{i}),$

where BCE is the Binary Cross Entropy Loss. In addition to the supervised
learning of human-annotated concepts, our model is designed to automatically
capture the information of the remaining concepts
$c_{n+1},c_{n+2}\cdots,c_{N}$ from the input in an unsupervised manner as it
is not possible to annotate a truly closed set of concepts. Specifically, the
unsupervised learning of the unannotated concepts is performed through the
training of concept decomposition for the reconstruction of generative
factors, which are described in Sections 4.3.3 and 4.4.

####  4.3.3. Decomposing Concepts into Generative Factors

Corresponding to the aggregation module, we propose the decomposition module
$g$ which learns a mapping from the concept embedding
$c=[c_{1},\cdots,c_{N}]^{\top}$ back to transformed disentangled
representations ($\mathbf{z^{\prime}}$) which are further transformed back to
the original disentangled representations $\mathbf{z}$ with decoders
$\mathcal{D}=\\{d_{1},\cdots,d_{n}\\}$. Mathematically,

(5) $\displaystyle\hat{z}$
$\displaystyle=\mathcal{D}(g(c))=[d_{1}(g(c)),\cdots,d_{n}(g(c))]^{\top}$

where $\hat{z}$ is the estimated generative factors from the given concepts
$c$. Here the decomposition module acts as the inverse of the aggregation
module and maps concepts back to the disentangled representations. Instead of
using one decoder directly for the mapping from $c$ to $z$, we use independent
decoders to project each dimension in $z^{\prime}$ to the corresponding
dimension in $z$, which mitigates the problem of concept leakage from concepts
to unrelated generative factors (Margeloiu et al., 2021).

####  4.3.4. Task Prediction using Learned Concepts

Task prediction entails the mapping from learned concepts in $\mathbb{R}^{N}$
to the prediction of task labels in $\mathbb{R}^{m}$, where $m$ is the number
of categories in classification tasks. The model prediction can be formulated
as

(6) $p(y|{c})=\text{Softmax}(h(c))$

where the function $h$ is a shallow neural network. Although $h$ can be chosen
up to task requirements, we utilize a linear polynomial to estimate $h$ for
model interpretability, with only human-annotated concepts as the input to see
how well they can be used to explain the model prediction:

(7) $\mathit{h}({c})=w_{0}+\sum^{n}_{i=1}w_{i}\cdot c_{i},$

where $w_{0},w_{1},\cdots,w_{n}s$ are the bias and weights in the predictor.
Given the true labels $y$, the loss for the task learning is defined as

(8) $\mathcal{L}_{pred}=\Phi\left(p(y|c),y\right),$

where $\Phi$ can be (binary) cross-entropy loss or mean square loss depending
on the type of the downstream task.

###  4.4. Disentangled Representation Consistency and End-to-end Training

The three modules of CoLiDR (DRL, Aggregation/Decomposition, Task Learning)
are trained end-to-end. To maintain the consistency of learned representations
of generative factors, we enforce $z$ and $\hat{z}$ to be as similar as
possible. We propose to use a DR consistency loss ($L_{drc}$), which can be
formulated as:

(9) $\mathcal{L}_{drc}=\|z-\hat{z}\|_{2}^{2}$

For a given set of disentangled generative factors $\mathbf{z}$, there exists
a family of surjective functions $f$ \- which map from $\mathbf{z^{\prime}}$
to concepts $\mathbf{c}$. Consequently, there exists a family of functions $g$
inverse coupled with $f$, which maps from concepts $c$ to $z$. As the function
$f$ is subjective, computing a direct inverse is intractable and hence
requires enforced consistency through $\mathcal{L}_{drc}$. The function $g$ is
NOT a standalone family of functions as they are inverse of a surjective
function $f$.

In addition, we also encourage sparsity on the transformed representations
${z^{\prime}}$ to identify the most important generative factors that are
aggregated to compose the concepts and reduce the impact of non-relevant
factors. The overall training objective can be given by:

(10)
$\mathcal{L}=ELBO+\lambda_{1}\mathcal{L}_{con}+\lambda_{2}\mathcal{L}_{pred}+\lambda_{3}\mathcal{L}_{drc}+\lambda_{4}\|z^{\prime}\|_{1}.$

##  5\. Experimental Setup

###  5.1. Dataset Descriptions

  * •

D-Sprites (Matthey et al., 2017): D-Sprites consists of procedurally generated
samples from six independent generative factors. Each object in the dataset is
generated based on two categorical factors (shape, color) and four numerical
factors (X position, Y position, orientation, scale). The six factors are
independent of each other. The dataset consists of 737,280 images. We randomly
split data into train-test sets in a 70/30 split.

  * •

Shapes3D (Burgess and Kim, 2018): Shapes3D consists of synthetically generated
samples from six independent generative factors consisting of color (hue) of
floor, wall, object (float values) and scale, shape, and orientation in space
(integer values). The dataset consists of 480,000 images. We randomly split
the data into train and test sets in a 70/30 split.

  * •

CelebA (Liu et al., 2015): CelebA consists of about 200,000 178 × 218 sized
RGB images of center-aligned facial photographs of celebrities. The faces are
annotated with 40 binary concepts like hair color, smile, attractiveness, etc.
Some of the features in the set are simple and observable like color (black,
blonde) and style of hair (wavy, bangs). However, many concepts are abstract
and subjective like attractiveness, heavy makeup, etc. We only consider the
objective concepts for experiments. We center-crop the images to 148x148 and
subsequently resize them to 64x64.

  * •

AWA2 (Lampert et al., 2013): Animals with Attributes-2 consists of 37,322
images of a combined 50 animal classes with 85 binary concepts like number of
legs, presence of tail, etc. We remove certain subjective concepts such as
“eats fish”. AWA2 is neither centered nor cropped and consists of significant
background noise, making it significantly harder to disentangle. We resize all
images to 64x64 and combine the train and test splits. We use 70% of the data
for training and 30% for testing.

###  5.2. Dataset Task Descriptions

Synthetic datasets: As dSprites and Shapes3D datasets are procedurally
generated, they do not contain an inherent downstream task, hence we construct
downstream tasks using combinations of GFs Similar to (Marconato et al., 2022)
and (Raman et al., 2023). For each task, we consider two GFs at random and a
sample has the label “$1$” when all factors satisfy a pre-defined criterion
and. For categorical factors, we consider the presence of exact values as
Truth, while for continuous factors we use a threshold. More details on task
construction can be found in the Appendix.

Real-world datasets: For the CelebA dataset, the downstream task is cluster
assignment (Marconato et al., 2022). For AWA2, the downstream task is
classification.

|  | Task Accuracy ($\uparrow$) | Concept Error ($\downarrow$)
---|---|---|---
|  | dSprites | Shapes3D | CelebA | AWA2 | dSprites | Shapes3D | CelebA | AWA2
wo/ dis. | CBM/CEM | 0.924 | 0.913 | 0.828 | 0.531 | 0.068 | 0.0910 | 0.0378 | 0.33
| GlanceNet | 0.931 | 0.915 | 0.814 | 0.517 | 0.070 | 0.0910 | 0.0379 | 0.44
w/ dis. | CoLiDR - VAE | 0.931 | 0.911 | 0.821 | 0.515 | 0.070 | 0.0912 | 0.0377 | 0.39
| CoLiDR - $\beta$-VAE (no $\mathcal{L}_{drc}$) | 0.921 | 0.911 | 0.819 | 0.515 | 0.071 | 0.0919 | 0.0381 | 0.39
| CoLiDR - $\beta$-VAE | 0.929 | 0.916 | 0.828 | 0.521 | 0.069 | 0.0910 | 0.0375 | 0.36
Table 2. Average task accuracy and concept errors across four datasets for
concept-based models with disentanglement learning (GlanceNet, CoLiDR - VAE,
CoLiDR - $\beta$-VAE) and models without disentanglement learning (CBM).
“CoLiDR - VAE” is a version of our model with a vanilla VAE (parameterized by
$\beta=1$) while “CoLiDR - $\beta$-VAE” is another version with a $\beta$-VAE
(parameterized by various $\beta$s) discussed in Appendix. Concept Errors are
reported for dSprites and Shapes3D as RMSE while for CelebA and AWA2 as 0-1
error. Best results for models with disentanglement learning are marked in
bold.

###  5.3. Model Implementation Details

Disentangled representations learning (DRL) module. We utilize two different
Variational Autoencoder architectures for estimating disentangle
representations in the DRL module. We utilize a standard VAE (Kingma and
Welling, 2014) and a $\beta$-VAE (Higgins et al., 2017). Even though similar
in formulation, a $\beta$-VAE is parameterized by a tunable hyperparameter
$\beta$ which controls the strength of disentanglement. For both VAEs, the
encoder is a 5-layer CNN with BatchNorm and LeakyReLU as the activation
function. The decoder is modeled symmetrically to the encoder with five
Transpose Convolutional Layers. The size of the latent space $k$ is set as 64
for d-Sprites and Shapes3D and 512 for CelebA and AWA2 datasets.

Aggregation/Decomposition module. The Aggregation module is composed of the
set of transformation operations $\mathcal{A}$ and mapping function between
transformed representations and concepts $f$. We model each of the $k$ neural
networks $a_{i}$ for $i\in\\{0,1,..,k\\}$ as a 3-layer network. For dSprites
and Shapes3D the network consists of layers sized [64,64,1] and for CelebA and
AWA2 the network consists of layers sized [512,512,1] with ReLU as the
activation function. The function $f$ is modeled as a single fully connected
layer. Similarly, the Decomposition module is composed of the set of inverse
transformation operations $\mathcal{D}$ and mapping function between concepts
and transformed representations $g$. We model each of the $k$ neural networks
$d_{i}$ for $i\in\\{0,1,..,k\\}$ as a 3-layer network. For dSprites and
Shapes3D the network consists of layers sized [1,64,64] and for CelebA and
AWA2 the network consists of layers sized [1,512,512] with ReLU as the
activation function. The function $g$ is modeled as a single fully connected
layer.

Task prediction module. As proposed in (Koh et al., 2020) and (Marconato et
al., 2022), we utilize a single fully connected layer mapping from the
concepts to predictions ($h$). The final output is passed through a softmax
layer to compute probability scores for the label.

###  5.4. Evaluation Setup

Task and concept accuracy. For a Concept-Based Model to be deemed effective,
it is required to be at par on performance with non-inherently explainable
models (or black box models). We measure the task accuracy and compare it
against a standard Deep Neural Network formed using the same encoder as the
VAE and a fully connected layer as the classification head.

Next, we compare the concept-accuracy of CoLiDR against CBMs as proposed in
(Koh et al., 2020). For datasets with binary concepts, accuracy is reported as
0-1 error (misclassification). For datasets with non-binary concepts, Root
Mean Square Error (RMSE) is reported.

Disentangled representation aggregation performance. Effective aggregation of
disentangled representations serves as the most important desiderata for
CoLiDR. However, it is not straightforward to understand the aggregation
effect. Unlike supervised disentanglement where each dimension of the
disentangled representations is forced to correspond to a concept, CoLiDR
aggregates multiple dimensions into a concept. Hence, instead of identifying
individually significant dimensions as concepts, it is important to consider a
set of representative dimensions from the aggregation module.

Due to VAE’s inherent disentanglement procedure, we can safely assume the
mutual independence of dimensions among the learned disentangled
representations. Hence, the latent dimensions themselves can be thought of as
features of the aggregation module. In effect, the problem of selecting a
representative set of dimensions is identical to assigning importance scores
to the most important features in a deep neural network. Multiple post-hoc
interpretability methods are proposed (Sundararajan et al., 2017; Ribeiro et
al., 2016). We utilize Integrated Gradients (IG)(Sundararajan et al., 2017) to
assign importance scores to each dimension. For an input $x$ and a baseline
$x^{\prime}$ (zero-vector), IG computes attribution scores for each feature
$i$ using the following path integral:

(11) $IG(x)=(x_{i}-x^{\prime}_{i})\cdot\int^{1}_{\alpha=0}\frac{\partial
F(x^{\prime}+\alpha\cdot(x-x^{\prime}))}{\partial x_{i}}d\alpha$

We utilize Captum (https://captum.ai/) to compute attributions using IG. The
values of the attributions are normalized and their absolute values are
assigned as the final attributions. Once the attribution scores are computed,
we utilize three distinct methods to evaluate the effects of the
representative set of disentangled representations.

  * •

GradCAM (Selvaraju et al., 2017) visualizations: For each dimension in the
top-k most important dimensions, the GradCAM attribution visualization plots
are plotted.

  * •

Latent space traversal: For concepts with one dominant dimension in the
representative set, we linearly interpolate the normalized latent space
between two terminal values to identify visual cues in each generated image

  * •

Oracle Classification: Even though the aforementioned visualizations provide
qualitative evaluation, a robust quantitative evaluation is required to
assuage the effects of confirmation bias. To achieve this, we train a Oracle
Network to automatically classify images to the associated concepts with high
accuracy. Subsequently, we compare heatmaps generated using CBMs and CoLiDR
with the actual ground-truth annotated bounding boxes and report the average
Intersection over Union (IoU) scores which measure the ratio of overlap to the
combined are between two bounding boxes (Refer Appendix for mathematical
formulation). We utilize a fine-tuned VGG-16 model as the backbone of the
oracle for each dataset. For more details, refer to Appendix.

Concept decomposition performance (intervention). Another desirable
desideratum of a Concept-based model is its ease of debugging. For a
misclassified sample, fixing the concept annotation by a domain expert should
be able to correct the model predictions. Recall that the decomposition module
decomposes concepts back into disentangled representations. We define an
intervention to be successful for a wrongly classified sample, replacing the
predicted concept score with the ground truth concept annotations changes the
wrong prediction label to the correct ground truth label $y$.

###  5.5. Evaluation Metric Descriptions

####  5.5.1. Oracle Training

We utilize a fine-tuned VGG-16 which contains 5 blocks of convolutional layers
followed by a max pooling layer at the end of each block. The final output is
passed through a 3-layer fully connected network to perform concept
prediction. The model is trained similarly to the concept network where each
concept is weighted by its relative occurrence in the dataset.

####  5.5.2. Intersection over Union Calculation (IoU)

Intersection over Union is defined as the ratio of overlap between two spans
of areas on an image. Assuming two bounding boxes $A$ and $B$, IoU is defined
as:

(12) $\text{IoU}=\frac{A\cap B}{A\cup B}$

We calculate IoU’s between heatmaps generated using GradCAM. The selected
pixels above a threshold are selected (255 value grayscale CAM maps above
150). As the heatmaps can be irregular in shape, we consider each selected
pixel as part of a set on a fixed image size.

##  6\. Results and Discussion

###  6.1. Task and Concept Accuracy

In Table 2, we compare our model against Concept Bottleneck Models (CBMs) (Koh
et al., 2020), which are trained without disentanglement learning, and
GlanceNets (Marconato et al., 2022), which involves the learning of
disentangled factors, on the average task performance and concept errors. We
implement CBMs by replacing the VAE with a standard Autoencoder and the task
learning module and the Aggregation/Decomposition module are replaced by
identity functions. For GlanceNets, we provide supervision on not only the
learned concepts but also a part of disentangled latent space (represented as
$\mathbf{G_{V}}$ in (Marconato et al., 2022)). Differently, we do not utilize
the Open-Set Recognition Mechanism as we do not specifically study concept
leakage.

As can be seen in Table 2, CoLIDR variants perform the best among all concept-
based models with disentanglement learning and even outperforms CBM on 3 out
of 4 datasets. The results show that CBM performs better than GlanceNet and
CoLiDR on the AWA2 dataset (0.531). A possible reason for this observation
stems from the fact that disentanglement performance on the AWA2 dataset is
not effectively captured by VAEs (Refer to Appendix) due to less training data
and significant noise in the dataset. Furthermore, the performance of
GlanceNet and CoLIDR is at par on dSprites and Shapes3D implying that datasets
that are ‘easier’ to disentangle yield similar performances across methods.

For the concept identification task, CoLiDR still achieves the best
performance with the lowest concept errors among all models with
disentanglement learning. Compared with the model without disentanglement
learning, Table 2 shows that our model performs comparable and sometimes
better than CBM on dSprtes, Shapes3D, and CelebA datasets. For the AWA2
datasets, CBM outperforms both CoLiDR and GlanceNet. We attribute this to the
fact that CBM can learn concepts very flexibly without the regularization of
disentanglement. However, this also leads to the problem of learning spurious
correlation instead of real semantics of concepts, as shown in Appendix. Note
that CoLiDR with no consistency regularization ($\mathcal{L}_{drc}$) performs
the worst on concept learning - implying strong disentanglement performance is
vital to concept learning.

###  6.2. Disentangled Representation Visualizations

Figures 4 and 5 respectively demonstrate the most important dimensions
constituting a concept. In the Figure 4, we demonstrate the concept
“brown_hair” and its associated dimensions with scores calculated using IG.
For example, in Figure 4 GradCAM heatmap visualizations on the two highest
computed constituent dimensions for a correctly identified concept
“straight_hair” (top) and “wavy_hair” are shown. We see that both dimensions
correctly correspond to the distributions constituting the hair of the people
in the image. Similarly, in Figure 5, the GradCAM visualizations corresponding
to the concepts “heart” (left) and “ellipse” (right) are shown. In addition,
we also provide a linear interpolation of the identified dimensions. As can be
seen, the dimensions are indeed responsible for the shapes of hearts and
ellipses.

Figure 4. GradCAM visualizations of the top-2 highest attributed dimensions
with respect to correctly predicted concepts “wavy_hair” (top) and
“straight_hair” (bottom) on 2 distinct samples from CelebA.

###  6.3. Comparision with Oracle Network

We partition a subset of concepts that can be easily understood in the CelebA
dataset and dub them ‘simple’ (Refer to Appendix for the exact splits). We
report average IoU values for correctly classified concepts on both ‘simple’
and all concepts in Table 3. Column-2, 3 and 4 list the average IoUs for
simple concepts across CBMs, CoLiDR with only the top-2 and top-5 dimensions
respectively. Heatmaps produced by CBMs (Margeloiu et al., 2021) show very low
IoU values implying they are not able to capture the spatiality of concepts
well. On the other hand, CoLIDR with $\beta$-VAEs give much higher IoUs than
CBMs - demonstrating that CoLiDR captures the spatial nature of concepts
extremely well. The more dimensions visualized (two vs five), the lower the
IoU goes - implying that most of the spatial concept information is encoded in
only a few dimensions.

| Oracle | CBM/CEM |  | CoLiDR
---
(dim=2)
| CoLiDR
---
(dim=5)
‘Simple’ IoU | 0.99 | 0.31 | 0.74 | 0.71
Overall IoU | 0.99 | 0.24 | 0.63 | 0.62
Table 3. IoU values for GradCAM heatmaps generated using CBMs (Column-2) as
compared to CoLiDR utilizing the top-2 dimensions (Column-3) and top-5
dimensions (Column-4).  Figure 5. GradCAM visualizations of the top-2 highest
attributed dimensions for the correctly predicted concept of shape for (top)
d-sprites and (bottom) Shapes3D. The interpolation along the highest
contributing dimension shows that the dimension effectively captures the shape
of the object in the image.

###  6.4. Test-time Intervention

Figure 6 shows the effect of intervening on a wrongly predicted concept by its
ground truth value. As the number of concepts vary for each dataset, we plot
the percentage of samples getting corrected after intervention (y-axis) v/s
percentage of concepts replaced by their ground truth values (x-axis). We see
that intervention is more successful on datasets with fewer concepts.

Figure 6. Test-time intervention for all datasets. CoLiDR Percentage of
samples correctly classified increases after intervention as the percentage of
concepts increases.

Additional results can be found in the Appendix.

Figure 7. GradCAMs Heatmaps corresponding to the most important dimensions of
correctly predicted concepts of 3 distinct images belonging to test set of
CelebA.

##  7\. Future Work

CoLiDR is a highly generalizable framework that can be plugged into any
disentanglement architecture. In this paper we demonstrate our results on
Variational Autoencoders, however many more complex architectures can be
utilized for learning disentangled representations. Generative Adversarial
Networks are another class of stochastic non-linear latent representation
learning framework which can be directly utilized as a DRL module. Another
avenue of exploration can go along the lines where concepts instead of being
independent of each other, are causally related - similar to (Heskes et al.,
2020).

##  8\. Conclusion

In this paper, we propose CoLiDR, a novel interpretable concept-based model
that learns concepts using aggregated disentangled representations. Empirical
results demonstrate that CoLiDR bridges the gap between concept learning and
disentangled representation learning by formatting human-understandable
concepts as aggregations of fundamental generative factors. The performance of
CoLiDR for task and concept prediction is on par with the concept-focused
models (CBM/CEM) that lack the explainability of generative factors. CoLiDR
can flexibly learn complex concepts as an aggregation of disentangled
generative factors which improves its performance on both task and concept
learning. Using an aggregation step, concepts learned using CoLiDR can be
constituents of multiple generative factors, giving a much finer overview of
interpretability.

## Acknowledgements

This work is supported in part by the US National Science Foundation under
grants 2217071, 2213700, 2106913, 2008208, 1955151. Any opinions, findings,
and conclusions or recommendations expressed in this material are those of the
author(s) and do not necessarily reflect the views of the National Science
Foundation.

## References

  * (1)
  * Alvarez Melis and Jaakkola (2018) David Alvarez Melis and Tommi Jaakkola. 2018.  Towards robust interpretability with self-explaining neural networks.  _Advances in neural information processing systems_ 31 (2018).
  * Bahadori and Heckerman (2020) Mohammad Taha Bahadori and David E Heckerman. 2020.  Debiasing concept bottleneck models with instrumental variables.  _arXiv preprint arXiv:2007.11500_ (2020).
  * Bengio et al. (2013) Yoshua Bengio, Aaron Courville, and Pascal Vincent. 2013.  Representation learning: A review and new perspectives.  _IEEE transactions on pattern analysis and machine intelligence_ 35, 8 (2013), 1798–1828.
  * Burgess and Kim (2018) Chris Burgess and Hyunjik Kim. 2018.  3D Shapes Dataset.  https://github.com/deepmind/3dshapes-dataset/.
  * Chen et al. (2018) Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. 2018.  Isolating sources of disentanglement in variational autoencoders.  _Advances in neural information processing systems_ 31 (2018).
  * Creswell et al. (2018) Antonia Creswell, Tom White, Vincent Dumoulin, Kai Arulkumaran, Biswa Sengupta, and Anil A Bharath. 2018.  Generative adversarial networks: An overview.  _IEEE signal processing magazine_ 35, 1 (2018), 53–65.
  * Ghorbani et al. (2019) Amirata Ghorbani, James Wexler, James Zou, and Been Kim. 2019.  Towards automatic concept-based explanations.  _Advances in Neural Information Processing Systems_ (2019).
  * Goyal et al. (2019) Yash Goyal, Amir Feder, Uri Shalit, and Been Kim. 2019.  Explaining classifiers with causal concept effect (cace).  _arXiv preprint arXiv:1907.07165_ (2019).
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016.  Deep residual learning for image recognition. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 770–778.
  * Heskes et al. (2020) Tom Heskes, Evi Sijben, Ioan Gabriel Bucur, and Tom Claassen. 2020.  Causal shapley values: Exploiting causal knowledge to explain individual predictions of complex models.  _Advances in neural information processing systems_ 33 (2020), 4778–4789.
  * Higgins et al. (2017) Irina Higgins, Loïc Matthey, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. 2017.  beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net.  https://openreview.net/forum?id=Sy2fzU9gl
  * Jeyakumar et al. (2022) Jeya Vikranth Jeyakumar, Luke Dickens, Luis Garcia, Yu-Hsi Cheng, Diego Ramirez Echavarria, Joseph Noor, Alessandra Russo, Lance Kaplan, Erik Blasch, and Mani Srivastava. 2022.  Automatic Concept Extraction for Concept Bottleneck-based Video Classification.  _arXiv preprint arXiv:2206.10129_ (2022).
  * Jeyakumar et al. (2020) Jeya Vikranth Jeyakumar, Joseph Noor, Yu-Hsi Cheng, Luis Garcia, and Mani Srivastava. 2020.  How can i explain this to you? an empirical study of deep neural network explanation methods.  _Advances in Neural Information Processing Systems_ (2020).
  * Kilbertus et al. (2018) Niki Kilbertus, Giambattista Parascandolo, and Bernhard Schölkopf. 2018.  Generalization in anti-causal learning.  _arXiv preprint arXiv:1812.00524_ (2018).
  * Kim et al. (2018) Been Kim, Martin Wattenberg, Justin Gilmer, Carrie Cai, James Wexler, Fernanda Viegas, et al. 2018\.  Interpretability beyond feature attribution: Quantitative testing with concept activation vectors (tcav). In _International conference on machine learning_. PMLR, 2668–2677.
  * Kim and Mnih (2018) Hyunjik Kim and Andriy Mnih. 2018.  Disentangling by factorising. In _International Conference on Machine Learning_. PMLR, 2649–2658.
  * Kingma and Welling (2014) Diederik P. Kingma and Max Welling. 2014.  Auto-Encoding Variational Bayes. In _2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , Yoshua Bengio and Yann LeCun (Eds.).  http://arxiv.org/abs/1312.6114
  * Koh et al. (2020) Pang Wei Koh, Thao Nguyen, Yew Siang Tang, Stephen Mussmann, Emma Pierson, Been Kim, and Percy Liang. 2020.  Concept bottleneck models. In _International Conference on Machine Learning_. PMLR, 5338–5348.
  * Lampert et al. (2013) Christoph H Lampert, Hannes Nickisch, and Stefan Harmeling. 2013.  Attribute-based classification for zero-shot visual object categorization.  _IEEE transactions on pattern analysis and machine intelligence_ 36, 3 (2013), 453–465.
  * Liu et al. (2015) Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. 2015.  Deep learning face attributes in the wild. In _Proceedings of the IEEE international conference on computer vision_. 3730–3738.
  * Locatello et al. (2019) Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. 2019.  Challenging common assumptions in the unsupervised learning of disentangled representations. In _international conference on machine learning_. PMLR, 4114–4124.
  * Locatello et al. (2020) Francesco Locatello, Michael Tschannen, Stefan Bauer, Gunnar Rätsch, Bernhard Schölkopf, and Olivier Bachem. 2020.  Disentangling Factors of Variations Using Few Labels. In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net.  https://openreview.net/forum?id=SygagpEKwB
  * Marconato et al. (2022) Emanuele Marconato, Andrea Passerini, and Stefano Teso. 2022.  Glancenets: Interpretable, leak-proof concept-based models.  _Advances in NIPS_ 35 (2022), 21212–21227.
  * Margeloiu et al. (2021) Andrei Margeloiu, Matthew Ashman, Umang Bhatt, Yanzhi Chen, Mateja Jamnik, and Adrian Weller. 2021.  Do Concept Bottleneck Models Learn as Intended?  _arXiv preprint arXiv:2105.04289_ (2021).
  * Matthey et al. (2017) Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. 2017.  dSprites: Disentanglement testing Sprites dataset.  https://github.com/deepmind/dsprites-dataset/.
  * O’Shaughnessy et al. (2020) Matthew O’Shaughnessy, Gregory Canal, Marissa Connor, Mark Davenport, and Christopher Rozell. 2020.  Generative causal explanations of black-box classifiers.  _Advances in Neural Information Processing Systems_ (2020).
  * Pedapati et al. (2020) Tejaswini Pedapati, Avinash Balakrishnan, Karthikeyan Shanmugam, and Amit Dhurandhar. 2020.  Learning Global Transparent Models Consistent with Local Contrastive Explanations.  _Advances in Neural Information Processing Systems_ 33 (2020).
  * Pittino et al. (2021) Federico Pittino, Vesna Dimitrievska, and Rudolf Heer. 2021.  Hierarchical Concept Bottleneck Models for Explainable Images Segmentation, Objects Fine Classification and Tracking.  _Objects Fine Classification and Tracking_ (2021).
  * Raman et al. (2023) Naveen Raman, Mateo Espinosa Zarlenga, Juyeon Heo, and Mateja Jamnik. 2023.  Do Concept Bottleneck Models Obey Locality?. In _XAI in Action: Past, Present, and Future Applications_.
  * Ribeiro et al. (2016) Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. 2016.  ” Why should i trust you?” Explaining the predictions of any classifier. In _Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining_. 1135–1144.
  * Sawada and Nakamura (2022) Yoshihide Sawada and Keigo Nakamura. 2022.  Concept bottleneck model with additional unsupervised concepts.  _IEEE Access_ 10 (2022), 41758–41765.
  * Selvaraju et al. (2017) Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. 2017.  Grad-cam: Visual explanations from deep networks via gradient-based localization. In _Proceedings of the IEEE international conference on computer vision_. 618–626.
  * Shen et al. (2022) Xinwei Shen, Furui Liu, Hanze Dong, Qing Lian, Zhitang Chen, and Tong Zhang. 2022.  Weakly supervised disentangled generative causal representation learning.  _The Journal of Machine Learning Research_ 23, 1 (2022), 10994–11048.
  * Shu et al. (2020) Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2020.  Weakly Supervised Disentanglement with Guarantees. In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net.  https://openreview.net/forum?id=HJgSwyBKvr
  * Sinha et al. (2023) Sanchit Sinha, Mengdi Huai, Jianhui Sun, and Aidong Zhang. 2023.  Understanding and enhancing robustness of concept-based models. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 37. 15127–15135.
  * Sundararajan et al. (2017) Mukund Sundararajan, Ankur Taly, and Qiqi Yan. 2017.  Axiomatic attribution for deep networks. In _International conference on machine learning_. PMLR, 3319–3328.
  * Taeb et al. (2022) Armeen Taeb, Nicolò Ruggeri, Carina Schnuck, and Fanny Yang. 2022.  Provable concept learning for interpretable predictions using variational inference.  _arXiv preprint arXiv:2204.00492_ 160 (2022).
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ 30 (2017).
  * Wu et al. (2020) Weibin Wu, Yuxin Su, Xixian Chen, Shenglin Zhao, Irwin King, Michael R Lyu, and Yu-Wing Tai. 2020.  Towards Global Explanations of Convolutional Neural Networks With Concept Attribution. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_. 8652–8661.
  * Yeh et al. (2020) Chih-Kuan Yeh, Been Kim, Sercan Arik, Chun-Liang Li, Tomas Pfister, and Pradeep Ravikumar. 2020.  On completeness-aware concept-based explanations in deep neural networks.  _Advances in Neural Information Processing Systems_ 33 (2020), 20554–20565.
  * Yuksekgonul et al. (2022) Mert Yuksekgonul, Maggie Wang, and James Zou. 2022.  Post-hoc concept bottleneck models.  _arXiv preprint arXiv:2205.15480_ (2022).
  * Zarlenga et al. (2022) Mateo Espinosa Zarlenga, Pietro Barbiero, Gabriele Ciravegna, Giuseppe Marra, Francesco Giannini, Michelangelo Diligenti, Zohreh Shams, Frederic Precioso, Stefano Melacci, Adrian Weller, et al. 2022\.  Concept embedding models.  _arXiv preprint arXiv:2209.09056_ (2022).

##  Appendix A Appendix

###  A.1. Downstream Task Descriptions

####  A.1.1. D-Sprites/Shapes3D

We create 6 distinct binary tasks for the D-Sprites dataset. Each task is
created by splitting the training set based on two factors. As the training
set is randomly sampled, the number of training points in each task is
different. However, the training set for each task is ensured to have balanced
labels.

####  A.1.2. CelebA

For CelebA dataset, we utilize the labels formed by clustering all the images
into 10 clusters using the clustering algorithm. The choice of the number of
clusters are made using TSNE visualization of the actual clusters. It is worth
noting that we utilize ten clusters instead of four in (Marconato et al.,
2022). This important change in task is made due to the fact that we utilize a
different subset of concepts (‘simple’ concepts) to model CoLiDR. Table 4
lists the set of concepts for which annotations are provided. We compute all
the experimental results on the simple set of concepts (left). The simple
concepts are chosen as they are objective in nature and can be easily
discerned by visual inspection of heatmaps.

Simple Concepts (20) | All Concepts (40)
---|---
Arched_Eyebrows, Bangs, Big_Lips, | 5_o_Clock_Shadow, Arched_Eyebrows, Attractive, Bags_Under_Eyes,
Black_Hair, Blond_Hair, Brown_Hair, | Bald, Bangs, Big_Lips, Big_Nose, Black_Hair, Blond_Hair, Blurry,
Double_Chin, Goatee, Gray_Hair, | Brown_Hair, Bushy_Eyebrows, Chubby, Double_Chin, Eyeglasses
Male, Mustache, Narrow_Eyes, Rosy_Cheeks, | Goatee, Gray_Hair, Heavy_Makeup, Male,
Sideburns, Smiling, Straight_Hair, | Mouth_Slightly_Open, Mustache, Narrow_Eyes, No_Beard, Oval_Face,
Wavy_Hair, Wearing_Earrings, Wearing_Hat | Pale_Skin, Pointy_Nose, Receding_Hairline, Rosy_Cheeks, Sideburns,
High_Cheekbones | Smiling, Straight_Hair, Wavy_Hair, Wearing_Earrings, Wearing_Hat,
| Wearing_Lipstick, Wearing_Necklace, Wearing_Necktie, Young
Table 4. List of Simple Concepts (left) and the entire concept set (Right) for
the CelebA dataset

####  A.1.3. AWA2

For AWA2 dataset, we keep the task the same as the identity classification of
50 different identities. The original dataset is annotated with 85 different
binary concepts. The list of concepts can be found
111https://cvml.ist.ac.at/AwA2/. For our experiments, we utilize a smaller
subset of 30 concepts. The subset is specifically chosen to only capture
concepts which are objective and human-understandable.

###  A.2. Implementation Details

####  A.2.1. Model Architecture

Disentangled Representations Learning (DRL) Module: Table 5 details the major
components of both architectures respectively. The standard VAE suffers from
the well-documented problem of (Higgins et al., 2017) over-regularization of
the disentangled representations $\mathbf{z}$, which hinders the quality of
visualized reconstructions $\mathbf{\hat{x}}$ due to the information
constraint imposed by the independent Gaussian priors. Hence, we evaluate
CoLiDR on both a standard VAE with $\beta=1$ and a $\beta$-VAE, where the
values of $\beta$ are tuned separately on each dataset. Specifically, we
employ the following values: dSprites: $\beta=0.025$, Shapes3D: $\beta=0.025$,
CelebA: $\beta=2.5e-5$ and AWA2: $\beta=1e-5$.

Aggregation/Decomposition Module: The Aggregation module $\mathcal{A}$ maps
between the latent representations $\mathbf{z}$ and concept set $\mathbf{C}$
while the Decomposition module $\mathcal{D}$ maps the concept set $\mathbf{C}$
to the latent representations $\mathbf{z}$ using a set of non linear
transforms $\\{a_{i}\\}$ and $\\{d_{i}\\}$ where i={0,..k}. We model each
transformation as a fully connected network mapping from a single dimension in
the latent space ($\mathbf{z}$) to the transformed latent space
($\mathbf{z^{\prime}}$) and vice-versa. The number of the intermediate layers
is 2 for dSprites/Shapes3D each of size 512, and the number of the
intermediate layers is 4 for CelebA/AWA2 each of size 512. The map between the
transformed latent space $z^{\prime}$ and $C$ and vice versa - is modeled as a
single feed-forward layer.

Task Prediction Module: The task prediction module utilizes a linear fully
connected layer mapping from the annotated concepts (Row-5 in Table 5) to the
number of task labels (Row-7 in Table 5). The final predictions are performed
after converting the model outputs into probabilities using the softmax
operator.

####  A.2.2. Training Hyperparameter Settings

Sequential Training: To avoid posterior collapse, we first train the VAE
without the influence of any other loss term, i.e.,
$\lambda_{1},\lambda_{2},\lambda_{3},\lambda_{4}=0$, only utilizing the ELBO
term (Equation 10). Subsequently, we only learn the concept mappings, i.e.,
$\lambda_{2}=0$ as proposed in Sequential training of (Koh et al., 2020).
Finally, we train the entire model architecture end-to-end using carefully
tuned hyperparameters $\lambda_{1},\lambda_{2},\lambda_{3},\lambda_{4}$ for
each dataset. The detailed values are given below:

  * •

dSprites: $\lambda_{1}=0.5,\lambda_{2}=0.2,\lambda_{3}=1,\lambda_{4}=0.1$

  * •

Shapes3D: $\lambda_{1}=0.5,\lambda_{2}=0.2,\lambda_{3}=1,\lambda_{4}=0.1$

  * •

CelebA: $\lambda_{1}=10,\lambda_{2}=0.08,\lambda_{3}=0.1,\lambda_{4}=0.01$

  * •

AWA2: $\lambda_{1}=5,\lambda_{2}=0.05,\lambda_{3}=0.1,\lambda_{4}=0.01$

Learning Rate and Scheduling: Training VAEs are notoriously highly sensitive
to the LR and optimizer. We utilize the Adam Optimizer with an initial
learning rate of 5e-3 for dSprites/Shapes3D, 1e-3 for CelebA and 5e-4 for
AWA2. The batch sizes are 64 each for dSprites and Shapes3D, while it is set
at 32 for CelebA and 16 for AWA2. We train the entire model architecture for
100 epochs with the first 50 epochs training the VAE, the subsequent 25 epochs
training the Aggregation/Decomposition module while the last 25 epochs train
the entire model, end-to-end.

###  A.3. Additional Visual Results

Figure 8 demonstrates the highest attributed latent dimensions associated with
correctly classified concepts. In particular. we compare the heatmaps
generated using CoLiDR with the heatmaps generated using CBMs. As can be seen
in Figure 8, the heatmaps corresponding to the most important dimensions of
blond hair, narrow eyes, and rosy cheeks correspond to their respective
correct positions in the image using CoLiDR as compared to CBMs. Similarly,
Figure 9(a) shows the heatmaps corresponding to correctly classified concepts
from the Shapes3D dataset. In addition, in Figure 9(b), we provide linear
interpolation of the most contributing latent dimensions corresponding to a
concept.

Figure 8. GradCAMs Heatmaps corresponding to a test image from CelebA. The Top
Row shows heatmaps computed on CBMs while the bottom row shows the heatmaps
corresponding to the most important dimensions of correctly predicted concepts
using CoLiDR. As can be seen: heatmaps from CoLiDR capture the spatiality of
concepts much better than CBMs.

(a) GradCAMs Heatmaps corresponding to the most important dimensions of
correctly predicted concepts of 3 distinct images belonging to test set of
Shapes3D.

(b) Linear interpolation of randomly sampled latent space corresponding to
images. The interpolated space corresponds well with the concepts it
contributes to the most.

Figure 9. Additional results on the Shapes3D dataset Component | dSprites/Shapes3D | CelebA/AWA2
---|---|---
Encoder | 3x | 5x
| [ Conv2D(filters, kernel=3, stride=2); | [ Conv2D(filters, kernel=3, stride=2)
| BatchNorm2D(); | BatchNorm2D();
| LeakyReLU()] | LeakyReLU()]
| filters in [32, 64, 64] | filters in [32, 64, 128, 256, 512]
Latent Dim | 64 | 512
Decoder | 3x | 5x
| [ ConvTranspose2D, filters, kernel=3, stride=2; | [ ConvTranspose2D, filters, kernel=3, stride=2;
| BatchNorm2D(); | BatchNorm2D();
| LeakyReLU()] | LeakyReLU()]
| filters in [64, 64, 32] | filters in [512, 256, 128, 64, 32]
Aggregation | 64x [Linear(1,64); ReLU();Linear(64,64); ] | 512x [Linear(1,512); ReLU(); Linear(512,512);
| ReLU(); Linear(64,1) | ReLU(); Linear(512,1)]
Concepts | Annot = 6; Total = 64 | Annot = 20 (CelebA), 30 (AWA2); Total = 512
Decomposition | 64x [Linear(1,512); ReLU(); Linear(512,1)] | 512x [Linear(1,512); ReLU(); Linear(512,512);
|  | ReLU(); Linear(512,512)]
Task Labels | 2 | 10 (CelebA), 50 (AWA2)
Table 5. CoLiDR module architecture details for the dSprites, Shapes3D
datasets (Left) and CelebA, AWA2 datasets (Right). We utilize identical
Encoder and Decoder stacks and identical input/output image sizes across all
datasets.