Chroma-VAE: Disentangled Color and Debiasing

• The paper demonstrates how Chroma-VAE leverages vector-quantized and partitioned VAE designs to disentangle chromatic representations and isolate shortcut features.
• It employs analytical color space conversions (e.g., RGB to CIE L*a*b*) and staged classifier training to enhance reconstruction quality and debias performance.
• Empirical evaluations on datasets like ImageNet, CelebA, and Chest X-Ray show improvements in color accuracy, segmentation IOU, and classification robustness.

Chroma-VAE encompasses two distinct lines of research in generative modeling: (1) a vector-quantized VAE trained on quasi-unsupervised color space conversion for disentangled, opponent-style chromatic representations (Akbarinia et al., 2020), and (2) a partitioned VAE architecture designed to isolate and mitigate shortcut learning in classifiers (Yang et al., 2022). Both approaches exploit the structure and interpretability of latent spaces, but target different challenges: structured color representation and robust debiasing, respectively.

1. Foundations of Chroma-VAE in Generative Modeling

The term “Chroma-VAE” first emerged in the context of color representation learning, specifically within vector-quantized variational autoencoders (VQ-VAEs) tasked with mapping images across analytically defined color spaces (e.g., RGB to CIE L*a*b*) (Akbarinia et al., 2020). The core VQ-VAE framework comprises an encoder $E$ mapping the input image $x \in \mathbb{R}^{H \times W \times 3}$ to a latent map $z_e(x)$, a discrete codebook $E = \{e_k\}$ enabling quantization, and a decoder $D$ reconstructing the output. The total training objective is:

$$L_{\text{VQ-VAE}}(x) = L_{\text{rec}} + L_{\text{codebook}} + L_{\text{commit}}$$

with $L_{\text{rec}}$ as the reconstruction loss (typically $\ell_2$ over possibly different color spaces), $L_{\text{codebook}}$ the codebook “pull,” and $L_{\text{commit}}$ the encoder “push” for latent stabilization.

A separate but related line of work characterizes Chroma-VAE as a mitigation approach for shortcut learning, leveraging a partitioned VAE latent space and staged classifier training (Yang et al., 2022). This structure assigns a small latent subspace $x \in \mathbb{R}^{H \times W \times 3}$0 to absorb cheap predictive cues (shortcuts), while a complementary subspace $x \in \mathbb{R}^{H \times W \times 3}$1 encodes the remaining information.

2. Chroma-VAE for Decorrelated Color Representation

The color-centric Chroma-VAE modifies standard VQ-VAE by training on cross-color-space conversion, e.g., mapping from RGB to opponent color spaces such as CIE L*a*b* or DKL. The conversion is explicit and relies on analytic formulas:

• RGB $x \in \mathbb{R}^{H \times W \times 3}$2 XYZ: linear transformation after gamma correction.
• XYZ $x \in \mathbb{R}^{H \times W \times 3}$3 CIE L*a*b*: nonlinear mapping disentangling luminance from chromaticity.

By requiring the model to reconstruct in a decorrelated target space, the codebook is pressured to organize along opponent axes, enforcing interpretable disentanglement between luminance and chromatic channels. Post hoc analysis of the codebook via usage statistics and lesion experiments reveals that this training paradigm results in distinct clusters of code vectors: some encoding mainly luminance shifts, others encoding chromatic variations (red–green or yellow–blue). Linear transformations fitted to lesion effects demonstrate high singularity indices for chromatic codes and low indices for luminance, corroborating the emergent structure (Akbarinia et al., 2020).

Quantitative evaluation on ImageNet, COCO, and CelebA indicates that Chroma-VAE achieves 1–2 lower ΔE (CIE ΔE₀₀) compared to standard RGB→RGB autoencoding, and boosts downstream ResNet-50 classification accuracy by 5–7%, as well as FPN-based scene segmentation IOU by 5–10%.

3. Chroma-VAE for Shortcut Isolation in Discriminative Learning

In a separate development, Chroma-VAE refers to a two-stage VAE-classifier architecture that systematically mitigates shortcut learning (Yang et al., 2022). Shortcuts—spurious, easily compressible predictive cues—dominate under empirical risk minimization due to their low information cost. Contrary to intuition, naive integration of generative modeling does not suffice to eliminate shortcut reliance.

Chroma-VAE addresses this by partitioning the latent code $x \in \mathbb{R}^{H \times W \times 3}$4, allocating a small dimensionality $x \in \mathbb{R}^{H \times W \times 3}$5 to $x \in \mathbb{R}^{H \times W \times 3}$6 (shortcut candidate) and the complement $x \in \mathbb{R}^{H \times W \times 3}$7 to $x \in \mathbb{R}^{H \times W \times 3}$8 (semantic candidate). Training proceeds in two phases:

• Stage 1: Jointly train encoder, decoder, and classifier using ELBO (VAE) and cross-entropy (classification) losses. The classifier receives only $x \in \mathbb{R}^{H \times W \times 3}$9.
• Stage 2: Freeze encoder/decoder and train a new classifier exclusively on $z_e(x)$0.

This design exploits the tendency for shortcuts to be encoded in the smallest, most discriminative space under classification gradient. The secondary classifier on $z_e(x)$1 thus operates on a representation with shortcut information effectively sequestered.

Empirical validation on ColoredMNIST, CelebA (blond/gender, attractive/smiling), MF-Dominoes, and Chest X-Ray datasets consistently shows dramatic worst-group accuracy improvements (e.g., from 0.6% to 73.8% on MF-Dominoes, from 10.3% to 57.8% on Chest X-Ray) over baseline ERM and VAE-classifiers.

4. Latent Structure, Interpretability, and Analytic Tools

Both variants of Chroma-VAE yield interpretability by design. In the color space approach, lesioning specific codebook vectors enables attribution of functional roles (luminance, red–green, yellow–blue). “Partial reconstructions” and blockwise sampling confirm that fixed subcomponents of the latent space control distinct generative factors.

In shortcut debiasing, Grad-CAM visualizations show that $z_e(x)$2 (shortcut block) attends to spurious cues (e.g., color patches or machine tokens), while $z_e(x)$3 (semantic block) attends to object- or pathology-relevant regions. Manipulating individual subspaces verifies that semantic content and shortcut features are largely segregated, enabling exploratory analysis and robust classification.

5. Evaluation, Limitations, and Extensions

Chroma-VAE’s efficacy is first, quantitative: lower color difference (ΔE) and significantly increased downstream classification and segmentation accuracy in decorrelated color space settings (Akbarinia et al., 2020); robust minority-group generalization in shortcut mitigation (Yang et al., 2022). Sensitivity analyses reveal that too small a total latent dimensionality impairs both reconstruction and separation, while excessive shortcut subspace size (partition fraction $z_e(x)$4) allows leakage of non-shortcut semantic features into $z_e(x)$5. Notably, increasing KL penalty ($z_e(x)$6) in joint supervision degrades performance in the shortcut-isolation context.

Recognized limitations include the dependency on high-quality generative modeling (VAE reconstruction), the challenge of isolating “rich” shortcuts that are not strictly low-dimensional, and reliance on proper architectural hyperparameters for successful separation. Future directions proposed include extending the approach to other generative models (normalizing flows, diffusion), introducing structured priors for disentangled subspaces, and explicitly targeting multiple spurious factors.

6. Applications and Broader Implications

The Chroma-VAE methodology yields multiple downstream applications:

• Color representation: Interactive editing (hue/lightness adjustment by code-vector manipulation), color-constancy and relighting through latent/conditional interventions, and style-transfer exploiting disentangled attributes (Akbarinia et al., 2020).
• Robust classification: Systematic shortcut-minimization in biomedical imaging and real-world tasks subject to spurious correlation, with broad implications for fairness and distributional robustness (Yang et al., 2022).
• Exploratory analysis: Latent manipulation (“chromatography”) as a tool for interpreting model behavior with respect to shortcut vs. semantic features, applicable for domain diagnostics and debiasing.

Taken together, Chroma-VAE constitutes a dual-pronged advance in generative modeling: structuring latent spaces to reflect interpretable (color or semantic) factors and providing robust debiasing strategies by explicit architectural partitioning and staged training. These principles generalize to a range of vision and classification tasks where disentanglement and reliable feature learning are essential.