Concept Whitening

The core mechanism of Concept Whitening is the orthogonal rotation of the latent space. After standardizing the activations to zero mean and unit variance, the module multiplies them by an orthogonal matrix. This matrix is optimized during training so that each axis of the latent space corresponds to a single, predefined concept. The orthogonality constraint ensures that the concepts are decorrelated and independent, meaning that activating one concept does not inadvertently influence another. This is a direct departure from standard Batch Normalization, which only centers and scales the data without enforcing semantic alignment.

Training a Concept Whitening module involves a multi-task loss function that balances two competing objectives:

• Classification Loss: The standard cross-entropy loss ensures the model maintains high accuracy on the primary task.
• Concept Alignment Loss: A secondary loss term measures how well the axes of the whitened space separate the data points belonging to different concepts. This is typically implemented as a softmax regression on the latent representation, predicting the presence or absence of each concept. This joint optimization ensures the model learns a representation that is both useful for prediction and inherently interpretable.

Unlike post-hoc explanation methods like TCAV or SHAP, Concept Whitening provides intrinsic interpretability. After training, the value of a single neuron in the whitened layer directly quantifies the presence of a specific concept. For example, in a medical imaging model, one axis might represent 'tissue density' while another represents 'lesion margin irregularity.' A doctor can trace a prediction back to these explicit, named axes without needing to run a separate explanation algorithm. This makes the model's reasoning process transparent and auditable in real-time.

The orthogonal matrix learned by Concept Whitening defines a concept subspace. The transformation can be decomposed into:

• Concept-Aligned Components: The projection of an activation vector onto the learned concept axes.
• Residual Components: The information orthogonal to the concept subspace, which captures task-relevant variance not explained by the predefined concepts. This decomposition allows for precise concept intervention. A practitioner can zero out the residual component to force the model to rely solely on known concepts, or edit the concept-aligned values to perform controlled counterfactual experiments, such as 'What would the prediction be if this concept were absent?'

Concept Whitening is often compared to Concept Bottleneck Models (CBMs), but they differ fundamentally in architecture:

• CBM: Creates a hard bottleneck by first predicting concepts explicitly and then using only those concept scores for the final task. This can limit accuracy if the predefined concepts are insufficient.
• Concept Whitening: Imposes a soft, geometric constraint on the latent space. It aligns axes with concepts but does not restrict information flow to only those axes. The residual subspace can still carry non-concept information, allowing the model to maintain higher accuracy when the concept set is incomplete, while still providing a clear, interpretable signal on the aligned axes.

To train a Concept Whitening module, the dataset must include dense concept annotations alongside the primary labels. For each input sample, a binary or continuous label must be provided for every concept of interest. For instance, a bird classifier would need annotations for 'wing color,' 'beak shape,' and 'size' for each image. This requirement is the primary limitation of the method, as acquiring such richly annotated data is expensive. However, the payoff is a model where the latent space is guaranteed to be organized by these human-specified attributes, enabling direct manipulation and verification.

The standard layer that Concept Whitening replaces. Batch Normalization standardizes activations to have zero mean and unit variance across a mini-batch, then applies a learned affine transformation.

Key differences from Concept Whitening:

• No decorrelation: Batch Norm operates on each channel independently, leaving axes correlated
• No semantic alignment: Axes remain arbitrary latent directions with no human meaning
• Training dependency: Relies on batch statistics, causing train/inference mismatch

Concept Whitening extends this by enforcing full whitening (decorrelation) and rotating axes to align with predefined concepts.

A linear algebraic operation that transforms a set of correlated variables into uncorrelated variables with unit variance. The output has an identity covariance matrix.

Mathematical process:

• Compute the covariance matrix Σ of the activations
• Perform eigendecomposition: Σ = UΛUᵀ
• Apply whitening matrix: W = UΛ⁻¹ᐟ²Uᵀ
• Transformed activations: z = W(x - μ)

In Concept Whitening, an additional orthogonal rotation is applied after whitening to align axes with concept directions, ensuring the latent space is both decorrelated and interpretable.

A latent representation where individual dimensions correspond to independent, semantically meaningful factors of variation in the data.

Properties of a disentangled representation:

• Modularity: Each latent dimension encodes exactly one generative factor
• Compactness: Each factor is encoded by a small number of dimensions
• Explicitness: Factor values are easily decodable from the representation

Concept Whitening produces a supervised disentanglement by forcing axes to align with known concepts, unlike unsupervised methods such as β-VAE which discover factors without semantic labels.

A square matrix R satisfying RᵀR = I, representing a rigid rotation in high-dimensional space that preserves distances and the whitened covariance structure.

Role in Concept Whitening:

• After whitening, the latent space has identity covariance but arbitrary orientation
• The rotation matrix R is optimized to align specific axes with concept activation vectors
• The objective maximizes the cosine similarity between axis j and the CAV for concept j
• Because R is orthogonal, the decorrelation property is preserved exactly

This two-step process—whiten then rotate—is what distinguishes Concept Whitening from simpler approaches like adding a concept loss term.

A direction in a neural network's activation space that represents a human-understandable concept. CAVs are the target alignment vectors for Concept Whitening's axes.

How CAVs are constructed:

• Collect positive examples of the concept (e.g., images with stripes)
• Collect negative or random counterexamples
• Train a linear classifier to separate them in activation space
• The vector orthogonal to the decision boundary is the CAV

In Concept Whitening, CAVs are pre-computed and used as optimization targets during training, so the learned axes directly correspond to these concept directions.

An alternative interpretable architecture that predicts concepts from inputs first, then makes final predictions using only those concept scores.

Comparison with Concept Whitening:

• CBM: Forces a hard bottleneck—the final classifier sees only concept probabilities, not raw features
• Concept Whitening: Aligns latent axes with concepts but allows the downstream layers to use all information
• Trade-off: CBM guarantees interpretability but may sacrifice accuracy; Concept Whitening preserves capacity while adding interpretability

Both methods aim to make concept-level reasoning explicit, but they impose different structural constraints on the model.