The Trade-off Between Covariate Dependence and Latent Structure in Representation Learning

Małgorzata Łazęcka; Ewa Szczurek

cs.LGAug 17, 2026

The Trade-off Between Covariate Dependence and Latent Structure in Representation Learning

Authors: Małgorzata Łazęcka, Ewa Szczurek

Organizations: Faculty of Mathematics, Informatics, and Mechanics, University of Warsaw Warsaw, Poland · Institute of AI for Health, Helmholtz Munich Munich, Germany

Abstract

Disentangled representation learning seeks latent representations whose indicidual dimensions each align with a distinct covariate. Unsupervised approaches typically target latent dimension independence, yet this gives no guarantee that the resulting dimensions align with semantically meaningful covariates. Supervised approaches structure the latent space using observed covariates, but under correlated covariates they cannot simultaneously control one-to-one latent-covariate alignment and latent independence. We introduce a unified, supervised framework that couples latent dimension-covariate dependence with constraints on the latent structure. Within this framework, we show an inherent trade-off, where enforcing latent independence or exclusive one-to-one latent-covariate dependence comes at a provable cost in latent-covariate alignment. We prove that the resulting disentanglement regimes are ordered by the strength of that alignment. Each regime admits a closed-form transformation of the latent space. We apply these transformations post-hoc to realign the representations of pretrained models such as CLIP, DINOv2, and ViT, and we fold them into the inference of informed factor analysis (iFA), a probabilistic model with covariate-informed factors. On simulated and real multi-omics data, we show that both post-hoc alignment and iFA enable controllability of structured latent representations.

Figures & tables

Figure 1. Overview of the proposed framework. Correlated covariates ( $\Sigma_{Y}$ ) shape the latent structure ( $\Sigma_{Z}$ ). The transformations impose different constraints on the latent space, trading off structural constraints (exclusive latent-covariate dependence, latent inter-independence) against the strength of pairwise latent-covariate dependence ( $\Sigma_{Y,\tilde{Z}}$ ).

Figure 2. Post-hoc latent space transformations of pretrained representations (CLIP, DINOv2, and ViT): dimensions with max corr (latent dimensions most strongly correlated with the covariates) and Procrustes orthogonal rotation against the proposed transformation applied after Procrustes rotation.

Figure 3. Trade-off between covariate dependence and factor independence for simulation scenarios 1 (PN) and 4 (N). Outlined markers denote means over $20$ repetitions.

Figure 4. Disentanglement metrics for the latent space transformed using $\,T_{\lambda}$ family and the exclusive transform $T_{\mathrm{ex}}$ as a function of covariate-correlation strength $\alpha$ in simulation scenario 1 (PN).

Figure 5. Results on the breast TCGA multi-omics data. Outlined markers denote means over $20$ repetitions.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Transformation $T$	(a) $T$ formula (general $Z$ )	(b) $d_{p}$ $(D=\diag(d_{p}))$	(c) $J_{T,p}$
$T_{0}^{*}$	$DM^{-1}\,\Sigma_{ZY}^{\prime}\Sigma_{Z}^{-1}$	$1$	$M_{pp}$
$T_{1}^{*}$	$D\,\Sigma_{ZY}^{\prime}\Sigma_{Z}^{-1}$	$\big((M^{2})_{pp}\big)^{-1/2}$	$\sqrt{(M^{2})_{pp}}=\\|M_{\cdot p}\\|$
$T_{\lambda}^{*}$	$D\big((1-\lambda)M^{-1}+\lambda I\big)\Sigma_{ZY}^{\prime}\Sigma_{Z}^{-1}$	$\big(((1-\lambda)I+\lambda M)^{2}_{pp}\big)^{-1/2}$	$\dfrac{(1-\lambda)M_{pp}+\lambda(M^{2})_{pp}}{\sqrt{(1-\lambda)^{2}+2\lambda(1-\lambda)M_{pp}+\lambda^{2}(M^{2})_{pp}}}$
$T_{\mathrm{ex}}^{*}$	$D\,\Sigma_{ZY}^{-1}$	$\big((M^{-2})_{pp}\big)^{-1/2}$	$\big((M^{-2})_{pp}\big)^{-1/2}$

Appendix

Table S1. Summary of the transformations, (a) their formulas, (b) diagonal normalizers, and (c) per-dimension correlations with covariates.

Panel of Figure 7 — Figure S6. Graphical model for iFA (uni-modal case).

Scenario 1	Scenario 2	Scenario 3	Scenario 4
Positive & negative (PN)	Autoregressive (AR)	Positive (P)	Negative (N)
-	$(\Sigma_{Y})_{ij}=(0.6)^{\|i-j\|}$	$(\Sigma_{Y})_{ij}=\begin{cases}1&\text{if }i=j\\ 0.25&\text{if }i\neq j\end{cases}$	$(\Sigma_{Y})_{ij}=\begin{cases}1&\text{if }i=j\\ -0.25&\text{if }i\neq j\end{cases}$

Unstructured dependencies	Time/space dependences	Scale effects	Dummy-coded categories

Appendix

Table S2. Simulation scenarios with different covariate dependency structures

Parameter	Description	Value
$N$	number of observations	$500$
D	number of features	100
$K$	number of factors	10
$\beta$	a vector of coefficients	$(0.9,0.75,0.6,0.45,0.3)$
$\theta$	noise to signal ratio	$0.1$

Appendix

Table S3. Simulation parameter settings

Parameter	Description	Values
$\alpha\in[0,1]$	Controls the strength of covariate dependences. The covariance matrix is transformed as ${\Sigma_{Y}^{(\alpha)}=\alpha\Sigma_{Y}+(1-\alpha)I}$ , with smaller values yielding weaker correlations ( $0$ - no correlations, $1$ - strongest correlations).	$0,\,0.25,\,0.5,\,0.75,\,1$
$b\in[0,1]$	Controls the strength of latent-covariate dependence by scaling the coefficient vector: $\beta^{(b)}=b\beta$ ; the smallest $b$ is selected so that the lowest $\beta_{p}=0.1$ , then $b$ are spread equilly ( $1/3$ - smallest factor-covariate dependence $0.2$ on average, $1$ strongest factor covariate dependence $0.6$ on average)	$1/3,\,1/2,\,2/3,\,3/4,\,1$

Appendix

Table S4. Simulation parameters varied across experiments.

Name	Description	Implementation details
PCA	Principal Component Analysis. An unsupervised linear method extracting orthogonal directions that maximise variance in $X$ . ( Pearson, 1901 ; Hotelling, 1933 )	PCA(n_components=10/20) - sklearn 1.8.0 ( Pedregosa et al., 2011 )
MOFA	Multi-Omicss Factor Analysis A probabilistic multi-view factor analysis model with a mean-field variational approximation that approximates factor independence. ( Argelaguet et al., 2018 )	mofapy2 0.7.4 model options: factors=10/20, spikeslab_weights=True, ard_weights=True, ard_factors=True , train options: convergence_mode="medium", iter=2000/5000
PLS	Partial Least Squares. Extracts latent directions that maximise the covariance between $X$ and $Y$ , subject to orthogonality of successive components. ( Wold, 1975 )	PLSRegression(n_components=10/20) - sklearn 1.8.0 ( Pedregosa et al., 2011 )
SPCA (Bair 2006)	Screening-based supervised PCA. A two-step method that first screens features of $X$ by their univariate association with $Y$ (e.g., absolute correlation coefficient above a threshold), then applies standard PCA to the surviving submatrix. The factors are PCA directions of a covariate-relevant subset of features.	Our implementation: features filtered by absolute correlation above the $0.5$ empirical quantile treshold, then PCA(n_components=10/20) - sklearn 1.8.0 ( Pedregosa et al., 2011 )
SPCA (Barshan 2011)	HSIC-based supervised PCA. Finds directions in $X$ maximising the Hilbert-Schmidt Independence Criterion with $Y$ . The objective reduces to a generalised eigenproblem on $X^{\prime}HLHX$ , where $L$ is the $Y$ -kernel matrix and $H=I-\frac{1}{n}\mathbf{1}\mathbf{1}^{\prime}$ is the centring matrix.	Our implementation: RBF kernel on $Y$ with median bandwidth, linear kernel on $X$ , scipy 1.17.1 and numpy 2.4.6
iVAE	Identifiable VAE. A variational autoencoder where the prior over $Z$ is conditioned on the auxiliary variables $Y$ via a factorised exponential family, $p(z\mid y)=\prod_{p}p(z_{p}\mid y)$ . Trained by maximising the conditional ELBO with both an encoder $q(z\mid x,y)$ and a decoder $p(x\mid z)$ . ( Khemakhem et al., 2019 )	Our implementation: PyTorch 1.13.1+cu117 , Gaussian variant encoder, prior and decoder MLPs with $2\times 64$ hidden units, LeakyReLU $(0.1)$ ; Adam, lr $=10^{-3}$ , batch $=256$ , epochs= $2000$ / $5000$ .

Appendix

Table S5. Comparison methods, with a brief description and implementation details.

Figure S7. Covariance structure of the pretrained image representations under each transformation, for A. CLIP, B. ViT, and C. DINOv2. Each panel shows latent-covariate (top) and latent-latent (bottom) covariances for the two baselines (dimensions with max correlation, Procrustes rotation) and the proposed $T_{\mathrm{ex}},T_{0},T_{1}$ .

Figure S8. Scenario 1 (PN) results. A. Covariance matrices $\Sigma_{YZ}$ and $\Sigma_{Z}$ in the data-generating setting. B. Covariance matrices $\Sigma_{Y\tilde{Z}}$ and $\Sigma_{\tilde{Z}}$ after transformations $T_{\textrm{ex}}$ , $T_{0}$ , $T_{1/2}$ , and $T_{1}$ . C. Frontier after transformations $T_{\lambda}$ and $T_{\textrm{ex}}$ . D. Empirical results for methods under two parameter settings: baseline and reduced $\alpha$ .

Figure S9. Scenario 2 (AR) results.A. Covariance matrices $\Sigma_{YZ}$ and $\Sigma_{Z}$ in the data-generating setting. B. Covariance matrices $\Sigma_{Y\tilde{Z}}$ and $\Sigma_{\tilde{Z}}$ after transformations $T_{\textrm{ex}}$ , $T_{0}$ , $T_{1/2}$ , and $T_{1}$ . C. Frontier after transformations $T_{\lambda}$ and $T_{\textrm{ex}}$ . D. Empirical results for methods under two parameter settings: baseline and reduced $\alpha$ .

Figure S10. Scenario 3 (P) results. A. Covariance matrices $\Sigma_{YZ}$ and $\Sigma_{Z}$ in the data-generating setting. B. Covariance matrices $\Sigma_{Y\tilde{Z}}$ and $\Sigma_{\tilde{Z}}$ after transformations $T_{\textrm{ex}}$ , $T_{0}$ , $T_{1/2}$ , and $T_{1}$ . C. Frontier after transformations $T_{\lambda}$ and $T_{\textrm{ex}}$ . D. Empirical results for methods under two parameter settings: baseline and reduced $b$ .

Figure S11. Scenario 4 (N) results. A. Covariance matrices $\Sigma_{YZ}$ and $\Sigma_{Z}$ in the data-generating setting. B. Covariance matrices $\Sigma_{Y\tilde{Z}}$ and $\Sigma_{\tilde{Z}}$ after transformations $T_{0}$ , $T_{1/2}$ , and $T_{1}$ (no $T_{\textrm{ex}}$ is reported, as in Scenario 4 (N) $\Sigma_{ZY}$ is not invertible). C. Frontier of transformations $T_{\lambda}$ . D. Empirical results for methods under two parameter settings: baseline and reduced $b$ .

Figure S12. Performance in Scenario 1 (PN), measured by three proposed metrics (average factor–covariate correlations, factor independence measured by the squared Frobenius distance of the correlation matrix of $Z$ from the identity, and average variance explained per factor) and five disentanglement metrics (disentanglement, completeness, informativeness, SAP, and normalised SAP). A. Varying the strength of correlations among covariates. B. Varying the strength of factor–covariate associations.

Figure S13. Performance in Scenario 2 (AR), measured by three proposed metrics (average factor–covariate correlations, factor independence measured by the squared Frobenius distance of the correlation matrix of $Z$ from the identity, and average variance explained per factor) and five disentanglement metrics (disentanglement, completeness, informativeness, SAP, and normalised SAP). A. Varying the strength of correlations among covariates. B. Varying the strength of factor–covariate associations

Figure S14. Performance in Scenario 3 (P), measured by three proposed metrics (average factor–covariate correlations, factor independence measured by the squared Frobenius distance of the correlation matrix of $Z$ from the identity, and average variance explained per factor) and five disentanglement metrics (disentanglement, completeness, informativeness, SAP, and normalised SAP). A. Varying the strength of correlations among covariates. B. Varying the strength of factor–covariate associations

Figure S15. Performance in Scenario 4 (N), measured by three proposed metrics (average factor–covariate correlations, factor independence measured by the squared Frobenius distance of the correlation matrix of $Z$ from the identity, and average variance explained per factor) and five disentanglement metrics (disentanglement, completeness, informativeness, SAP, and normalised SAP). A. Varying the strength of correlations among covariates. B. Varying the strength of factor–covariate associations

Figure S16. Consistency of iFA (no $T$ ) as the number of observations $N$ grows, Scenario 1 (PN). A. Recovery of factor-covariate coefficients. B. Mean squared error of the estimated $\hat{\beta}_{p}$ averaged over $p$ . C. cosine similarity between recovered and true informed factors. Shaded bands show $\pm 1$ s.d. over repetitions.

Figure S17. Runtime of all methods on the one simulation setup.

Disentangled Representation Learning

Disentangled representation learning separates distinct factors of variation so models can interpret, control, and generalize their predictions.

See all 231 papers

XFACTORS: Disentangled Information Bottleneck via Contrastive Supervision
Jan 29, 2026Alexandre Myara, Nicolas Bourriez, Thomas Boyer +3Variational Autoencoders Disentangled Representation Learning
ENS Paris
Variational Learning of Disentangled Representations
Jun 20, 2025Yuli Slavutsky, Ozgur Beker, David Blei +1Variational Autoencoders Disentangled Representation Learning
Department of Statistics, Columbia University, New York City, USA · Irving Institute for Cancer Dynamics, Columbia University, New York City, USA · Department of Computer Science, Columbia University, New York City, USA +1
Sculpting the Latent Space: Disentanglement with WAE-MMD and Programmable Priors
Oct 13, 2025Quentin Fruytier, Akshay Malhotra, Shahab Hamidi-Rad +3Disentangled Representation Learning Semantic Representation Learning
The University of Texas at Austin · InterDigital Communications, AI Lab

Jan 29, 2026J/K move · Enter open · S save

XFACTORS: Disentangled Information Bottleneck via Contrastive Supervision

Alexandre Myara, Nicolas Bourriez, Thomas Boyer, Thomas Lemercier+2

ENS Paris

Variational Autoencoders Disentangled Representation Learning Information Bottleneck Variational Information Bottleneck

Open paper AskPDF arXiv

Disentangled representation learning aims to map independent factors of variation to independent representation components. On one hand, purely unsupervised approaches have proven successful on fully disentangled synthetic data, but fail to recover semantic factors from real data without strong inductive biases. On the other hand, supervised approaches are unstable and hard to scale to large attribute sets because they rely on adversarial objectives or auxiliary classifiers. We introduce XFactors, a weakly-supervised VAE framework that disentangles and provides explicit control over a chosen set of factors. Building on the Disentangled Information Bottleneck perspective, we decompose the representation into a residual subspace

\mathcal{S}

and factor-specific subspaces

\mathcal{T}_1,\ldots,\mathcal{T}_K

and a residual subspace

\mathcal{S}

. Each target factor is encoded in its assigned

\mathcal{T}_i

through contrastive supervision: an InfoNCE loss pulls together latents sharing the same factor value and pushes apart mismatched pairs. In parallel, KL regularization imposes a Gaussian structure on both

\mathcal{S}

and the aggregated factor subspaces, organizing the geometry without additional supervision for non-targeted factors and avoiding adversarial training and classifiers. Across multiple datasets, with constant hyperparameters, XFactors achieves state-of-the-art disentanglement scores and yields consistent qualitative factor alignment in the corresponding subspaces, enabling controlled factor swapping via latent replacement. We further demonstrate that our method scales correctly with increasing latent capacity and evaluate it on the real-world dataset CelebA. Our code is available at github.com/ICML26-anon/XFactors.

Abstract

Disentangled Representation Learning

Explore similar work

XFACTORS: Disentangled Information Bottleneck via Contrastive Supervision

Variational Learning of Disentangled Representations

Sculpting the Latent Space: Disentanglement with WAE-MMD and Programmable Priors