Sparsity regularization via tree-structured environments for disentangled representations
Authors: Elliot Layne
Jason Hartford
Sebastien Lachapelle
Mathieu Blanchette
Dhanya Sridhar
Venue: Accepted by TMLR ()

Abstract: Many causal systems such as biological processes in cells can only be observed indirectly via measurements, such as gene expression. Causal representation learning---the task of correctly mapping low-level observations to latent causal variables---could advance scientific understanding by enabling inference of latent variables such as pathway activation. In this paper, we develop methods for inferring latent variables from multiple related datasets (environments) and tasks. As a running example, we consider the task of predicting a phenotype from gene expression, where we often collect data from multiple cell types or organisms that are related in known ways. The key insight is that the mapping from latent variables driven by gene expression to the phenotype of interest changes sparsely across closely related environments. To model sparse changes, we introduce Tree-Based Regularization (TBR), an objective that minimizes both prediction error and regularizes closely related environments to learn similar predictors. We prove that under assumptions about the degree of sparse changes, TBR identifies the true latent variables up to some simple transformations. We evaluate the theory empirically with both simulations and ground-truth gene expression data. We find that TBR recovers the latent causal variables better than related methods across these settings, even under settings that violate some assumptions of the theory.

Published in Transactions on Machine Learning Research (07 (07/2025)
Sparsity regularization via tree-structured environments for disentangled representations
Elliot Layne
elliot. Ulliotlloyne@mmll.megg
McGill University, Mila
Jason Hartford
josonnwvalencellbbs.oo
Valence Labs
Sébastien Lachapelle
lachaseb Ammia. quebec
Université de Montréal, Mile
Samsung SAIT Al Lob, Montreal
Mathieu Blanchette
mabhies.Baancheeteeenggg..l.
McGill Universily, Mila
Dhanya Sridhar
dhangy..ridhhrrr@mila..guu quebec
Université de Montréal, Mile
Reviewed on OpenReview: https: // openreview. net/ form? id= 2=U20j0200
Abstract
Many causal systems such biological processes in cells can only be observed indirectly via measurements, such as gene expression. Causal representation learning-the the task of correctly mapping low-level observations to latent causal citiablle-collle advance scientific understanding by enabling inference of latent variables such as pathway activation. In this paper, we develop methods for inferring latent variables from multiple related datasets (environments) and tasks. As running example, we consider the task of predicting phenotype from gene expression, where often collect data from multiple cell types organisms that are related in known ways. The key insight is that the mapping from latent variables driven by gene expression to the phenotype of interest changes sparsely across closely related environments. To model sparse changes, we introduce Tree-Based Regularization (TBR), an objective that minimizes both prediction crror and regularizes closely related environments to learn similar predictors. We prove that under assumptions about the degree of sparse changes, TBR identifies the true latent variables up some simple transformations. We evaluate the theory empirically with both simulations and ground-truth gene expression data. We find that TBR recovers the latent causal variables better than related methods across these settings, even under settings that violate some assumptions of the theory.
1 Introduction
Discovering new knowledge using machine learning is made ever possible with the growing amounts of unstructured data that we collect, from images, video, and text to experimental measurements from large-scale biological assays. However, scientific discovery from such low-level measurements requires repre- sentation learning, the task of mapping low-level observations to high-level feature space.
As running example, consider learning about drivers of discase from gene expression mcasurements, where genes drive protein abundance levels that mediate discase. If we want discover meaningful features that could be causally linked to outcomes, need methods to learn disentangled representation. This can be

---

Published in Transactions on Machine Learning Research (07/202) 2025)
x
Z
y
The causal directed acyelic graph (DAG) relating observations X, latents Z and target variable of interest Y. The marginal of X and Z + relationship is specific to environment e.
Figure 1:
done an encoder that maps observations to causally relevant latent variables, and not to transformations of these variables that change their meaning.
Disentanglement is impossible from independent and identically distributed (IID) data without assumptions
to constrain the solution space (Hyvärinen & Pajunen 1999 Locatello et al.) 2019 but by leveraging known structure such as non-stationarity (Hyvarinen Morioka 2017 Hyvarinen et al. 2019) sparsity (Brehmer et al. 2022: Ahuja al. 2022c Lachapelle et al., 2022b), , it is possible identily latent variables. The chal- lenge is in finding assumptions that are strong enough to rule out spurious solutions, while remaining flexible enough to fit the domain of interest. To this end, this paper proposes Tree-Based Regularization (TBR), a new disentangled representation learning method that is particularly well-suited to biological settings. TBR leverages non-stationarity from multiple datasets called environments that are hierarchically related via a known tree.
In our rumning biology example, the environments defined by cell types, and their relations come from the tree describing the cell-type differentiation process Enver al. (2005). We assume that the relationship between low-level observations and latent variables, represented by P(Z is constant across these envi- ronments as it is driven by the underlying physics of the cell. However, the mechanism driving the outcome (e.g., disease), (z)))) changes across environments, reflecting processes like variation across cell types or evolution.
The key assumption we make is that across closely related environments, the effects of only a sparse set of true latents will vary. That is, changes the conditional distributions Pe PY are sparse, such that environments that are siblings in the known tree share the same functional form up few paramcters, while more distant environments will have accunulated more changes. TBR enforces this structure via sparsity- inducing penalty. These constraints reflect biological settings, where it is common to observe data across multiple cell types or model organisms with hicrarchical relationship structures (e.g. a phylogenetic tree), and allow us to learn representations that identify the true latent variables up to irrelevant transformations.
We show both theoretically and empirically that by enforcing this sparsity assumption on changes of the conditional distributions, it is possible to identify the ground truth latent variables up a permutation and scaling factor, assuming the presence of sufficient number of environments such that the effects of each variable are varied. We further analyze the sensitivity of TBR to assumption violations and find that it remains robust up to a point. Then, we apply TBR to a dataset of gene expression measurements across different cell types, where we hold out some genes as latent variables and simulate phenotypes derived from these latents. We find that TBR recovers the true latent variables with higher fidelity than standard roprosentation lcarning approaches, which then leads to higher transfer learning performance, demonstrating concrete benefit of the TBR approach.
2 Problem Formulation
We consider an input vector X (c.g., gene expression) and a target variable of interest Y (c.g., discase phenotype). We observe n iid samples (xf, (f from E different environments where each e denoles the environment that produced the sample. Figure lillustrates the causal model the data generating process, which we note is distinct from most causal representation learning works as Z is causal descendant of X, rather than a causal ancestor. The inputs X drive latent features Z that mediate the effect on Y. P(Y\Z) varies across environments while P(Z[X) remains invariant. Specifically, we consider the following generative
2

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
process:
- - +1 Y" w. Z +
(1)
There exists an environmenttinvariaal continuous function relating X to Z, with noise term 7,, assumed to be within the space of functions approximable with a neural network (Cybenko 1989). Additionally, there exists an environment-specift function & that relates latent factors Z to Y In particular, case where & is a linear mapping with weights we E IR* where k [Z, and is an environmenn-nneeeendenn noise term. Linearity of the output map is clearly restrictive, but given that Φ(-) can be basis function expansion of Z, this process could still capture variety of nonlinear maps between latents and targets.
consider the
Critically, we posit that the environments have diverged according to some evolution-like process. We denote directed tree T = A), where environments are the nodes of the tree, and A C E² denotes the set of a between environments. The environment we will often observe samples only from leaf environments, and not from internal nodes in the tree (which represent historical environments, such as ancestral species or transient progenitor cells), but our framework is agnostic to this.
is assumed to be the root environment. In practice,
We assume that the environment-speciiie parameters we are subject lo mutations between parent-child environment pairs. In addition, we assume that these mutations are sparse, such that across each are (Cm, ex), there cxists a subset of indices C {1....k} such that Vi we.. * We, and the components of We.. and we,, are equal otherwise. We denote the vector of differences in parameters across are 6g Wen, Wem Note that the weights for any environment we be expressed as the sum of the root weights we and the relevant mutation vectors:
We W/
Σ
8a
@Epath(O.e)
where (0, e) denotes the set of ares forming the path from the root environment 0 to environment e. We denote E ARx as the matrix whose rows are the set of vectors &g for each are a E A.
Figure 2 illustrates an example data-generating process involving tree with 5 environments and leaves.
The proposed evolution-like process relating environments via sparse local changes causes the distributions (Y Z) to gradually change. This model of evolution aligns with many observed processes in nature where sparse genetic or epigenetic changes alter the process leading to a given phenotype, such as cell type differentiation and species evolution.
2.1 based regularization
We approximate y with an arbitrary deep neural network ve with parameters 6:
2 = (X(
(3)
We estimate a set weights úig E R², along with a matrix À E RLAXAA whose rows contains the estimates 8, for each a E A, permitting estimation of cach áē in the form of Eq.
We jointly optimize Vo, wh and Â to produce optimal predictions Y. across all environments. Additionally, regularize our estimate Å with sparsity-inducin norm || Ilu-
Assuming that we only observe data samples in subset of environments L C E, typically the leaves in T. this yields the following loss function:
Loss
"(w Tw y0)²+ + \ lb
(4)
eEC
where À > 0 is a scaling parameter. The first term in Eq. 4 is a standard prediction error term, minimizing the mean-squared error (MSE) for predictions Y. The second term regularizes the Lo norm of the matrix A.
1 An is directed edge between two nodes.
3

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
A
Ely, [X] = (%g By WCX)
Hp
B
BED
6p.
Fly [X] = (Ny + hap + 6p2) (P())
Sp.c
C
[Xy = (My BED + 4p.c (V())
Figure 2: An example of posited DGP and corresponding TBR parameterization for dataset with samples originating from three different environments (A, B, and C). A set of weights Wo is associated to the root. A trainable update Wa is associated to each edge The prediction function fu node will make predictions data samples originating from environment Adapted from Layne et al. (2020).
Since parameters in A the differences across each arc in A, by maximizing the sparsity of A, encourage parent-child environment pairs to learn similar predictors. This captures the essence of our motivation, that the mechanisms controlling P(Y Z) should only change rarely. We will denote our method, optimizing Eq. as free-based regularization (TBR).
3 Identifiability via Sparse Brsed Regularization
This section presents our main theoretical contribution on disentanglement. More formally, we will show that the learned Ż identifies the true latent variables Z up permutation and element-wise rescaling, i.e., 2 DPZ where D is invertible diagonal matrix and P is permulation. It is necessary that the dimension Ż Z). We discuss estimation of [Z), should it be unknown, in $
The problem of non-identifiability arises because the first term in the loss function in Eq. which penalizes prediction errors, admits many solutions for 2. Consider the generative process specific to environment A
in Figure
(wg + 5A) Tz +
(5)
We can see that the distribution Y remains unchanged if Z is subject to some lincar transformation L, so long as an inverse transformation is applied to wy and A:
y= (w) L S LZ +
(6)
However, the solution = LZ potentially linearly entangles the components of the true features Z. For tasks that require the learned features to be semantically faithful to the true features, egg, inferring causal effects of the features, this entanglement will lead to biased conchusions.
The sccond term in the loss in Eq. plays a key role selecting disentangled solutions among all solutions that achieve optimal prediction error. This is because the regularization term Å||Â|lo loo generally not invariant to linearly entangled solutions such LZ, because for most choices of L, ||Allo # ||LL 'lo. unless L = PD (when the representation is disentangled).
The core of the theory in this section lies in establishing assumptions under which the only linear transfor- mations that can be applied to A while maintaining the same level of sparsity, and thus the regularization penalty Eq. 4 are permutation and scaling operations. Combined with a result demonstrating the identi- fiability of up to linear transformation, this regularization constraint permits disentanglement of both A and the latent features Z.

---

Published in Transactions on Machine Learning Research (07/2025)
Our intuition focused on linearly entangled solutions ž LZ but in general, the hidden layer of neural network could produce arbit entangled solutions. As such, we prove the identifiability of features learned using tree- basdd regularization in two parts. First, we establish conditions under which the features learned with a nonlinear function, c.g., a neural network, are identified up to lincar transformations, bringing us back to the linearly entangled case. This result builds on proof techniques found in the literature (Khemakhem ct al. 2020b a Roeder et al. 2021 Ahuja et al. 2022a Lachapelle et al. 2022a).
Second, we show that with some assumptions, the sparsity constraint on |l permits further identification of up to permutation and scaling. Specifically, we show that any A with most one non-zero value per row cannot be entangled without resulting increase in regularization cost.
3.1 Linear Identification of
Assumption 3.1 (Data-generating process (DGP)). For each environment e, we assume the pairs (æ; y) are drawn i.i.d from distribution satisfying
E|Y X, = w. (X).
(7)
Additionally, assume the support of p( e), denoted by X' R is fixed across environments.
Assumption 3.2 (Sufficient task variability). We assume that there exist environments e1,-..; E & such that the matrix [Weys"ie W..]e is invertible.
Assumption 3.3 (Sufficient representation variability). We assume there exist x² ..... x² X such that the matrix (x') ).... (x*)] is invertible.
Taken logether, these assumptions establish our first requirement, by implying that the latents are identified up to linear transformation L,
Proposition 3.4. Suppose Assumptions 3.1
3.2 & 3.3 hold. Moreover, consider the learned parameters
Exemanoo 80 for all E. If for all X X and all Analogously to Eq. 2 define ih wo + an inverlible matrir L E kkk such that
do and (da) and the learned encoder function T (x)
have E|Y X, e] ae (X),, then, there exists
For all x E X, V(x)= (x);
2. For all e e E, : L = w._ and
For all a E A, 8 8n
The proof is included in Appendix A.1. Identification Up To Scaling and Permutation
The previous section shows that we can identify every 5, up common invertible linear transformation L. Recall A, Å glA[xk simply the concatenations of the 6, and 8,, respectively. This means write AL = Å. We now show that penalizing the Lo-norm of our estimate Á allows us to identify A up to a permutation and scaling.
Assumption 3.5 (1-sparse perturbations). Each 8, has at most one nonzero value.
Assumption 3.6 (Sufficient perturbations). For all i {l...., k}, there exists A such that da, # 0.
Proposition 3.7 (Disentanglement via 1-sparse perturbations). Suppose Assumptions 3.5 8 3.6 hold, let L an invertible matrix and let â = AL. If || < ||A||o, then is a permutation--eauiis matria.
Proof sketch To understand the crux of proving this result, recall the aspect of TBR that drives the result in the first place: the Lu norm prefers solutions Á that are sparse. For this regularization choose the disentangled solution, 2 = ZDP, we have to show that. Å Z L-1 the result. of linearly entangled solution,
5

---

Published in Transactions on Machine Learning Research (07 (07/2025)
cannot be sparser than À for the disentangled solution. To this end, our strategy will be to fix the amount of sparsity in the ground-truth matrix A. Then, we will show by contradiction that there exists a mapping between the columns of A and Á such that each column in A is at least as sparse as the corresponding column in A. Using this result, we then show that to meet our constraint on the sparsity of Å, must be the case that L permutatioonssclinng matrix.
While Assumption 3.5 which requires all rows in A have most 1 non-zero entry, is unlikely hold in many practical applications of interest, we do note that empirically consider violations of this assumption in 4 and find the results promising. Some further discussion about relaxation of the 1-sparse assumption is included in Appendix A.2.
3.2 Proof of Proposition 3.7
This result builds on arguments from both Freyaldenhoven (2020) and Lachapelle al. (2022b).
Proof. By Lomma A.1 there exists a permutation æ such that for each column í in L, there is an index tT)) where ( is a non-zero value a. Thus A is equal to some linear combination of columns in A where column (3) has coefficient a. Details are included in Appendix A.
We use operator to denote the "sparsity pattern" of a matrix. The sparsity pattern of is denoted as Sa. and equal the set of indices of non-zero elements A. Sg indicates the complement of the sparsity pattern: the set of indices corresponding to entries with value zero. Similarly, Sa.; denotes the sparsity pattern of A at column i.
Our proof builds upon series of column-wise comparisons. For each index i E {I....,k)) ,k,, we compare the Lo norm of column x()))) and column Å...
We define two sets of indices, Ω4 and I.,
2 = Sa 055 A..
(8)
r = S m()) n
(9)
Thus, 24 represents nonzero entries that were present in A,,)),, but not in Å. and I; represents the opposite, i.e. nonzero entries that were not present in A.,(()) but that are present in A. Thus, the following equality holds:
م|ر |A =0llo + [a - 12,1
(10)
We now show that Ω is empty by contradiction. Assume there exists Pr This means Aa, + 0 and d.s = La, = O. By the 1-sparse assumption, A,,-x(i) = This further implies that D. = (a.10) = 0, which is a contradiction. We thus conclude that Ω is empty.
Therefore, for arbitrary column i, it will always be the case that || 2 11A
We rewrite the original sparsity constraint:
| ||Aloo
(11)
i s IlA.o
(12)
> ll s Σ ||A,=
(13)
E(1 -IIA o Ilo) 0
(14)
The above combined with ,/ |1Alo 2 implics that, for all i, || 1|Aloolllo This means T3 is empty well.

---

Published in Transactions on Machine Learning Research (07 (07))
Because I, is empty for all i, we have that
Vi, a, (1, = = Se, = 0,
(15)
or, equivalently,
Vi, a, Da,i
a,)-1(0) = =0,
(16)
Choose arbitrary i. By Assumption 3.5 & 3.6 By Eq. 16 we have
there exists such that Aa,7()) # 0 and -a,))(1)
0.
(() =0 = A ((-((()) = 0
(17)
(18)
A L. =0
(19)
Aa Lx))-- 0,
(20)
which implies (x) (since (e,x(1) # 0)- This holds for all i, thus L is permutatioo--calling matrix.
4 Empirical Studies
The goal of the empirical studies is to demonstrate that TBR recovers disentangled latents, which further lead to accurate causal inference and transfer learning. To that end, we conduct studies with simulated data and gene expression data across cell types investigate the behaviour of TBR under various generative scttings, robustness to violations of Assumption 3.5 and potential downstream benefits to disentangled representations. For all these experiments, compare the represenlations produced by TBR to those from baseline method: a standard linear map w" fit to the target label Y on top of the learned nonlinear representation V yielding the following optimization objective:
LosSbaseline
Σ (w fo(')) yoy2 eEC
(21)
We note that throughout experimentation, the regularization of |lÂl approximated via regulariza- tion of |||]] following the practice of other sparsity- based disentanglement methods (Lachapelle et al, 2022a bj. This approximation introduces slight bias into model by shrinking the non-zero values of ||l- Additional implementation details are included Appendix A.5. All results are averaged over repetitions unless otherwise specified. Additionally, in the simulation study, we consider second baseline, which we refer the "Sparse Predictors" method, following the approach developed by Lachapelle et al.) (2022b). This method learns predictors for multiple tasks that share representation Z such that each task Y depends sparsely on the components of the true underlying latents 7. We implement this method by training enviromment-specife weights we with L1 regularization to induce sparsity. Although Lachapelle jet al. 2022b) show that this form of sparsity regularization suffices to recover disentangled features Z, the sparse dependence assumption does not hold in our settings, where mechanism changes- not mappings are sparse. As such, we hypothesize that the Sparse Predictors approach will not produce discntangled features in our studies.
Performance metric In Figure 3a we evaluate the disentanglement of Ż using the mean correlation coefficient (MCC). The MCC between Z and 2 is defined by:
1 b k
arg max TE
Σ
Pearson ((2,)))((22))))))))))))))))) (, Z,) ||.
(22)
where is the set of possible permutations of 2 and we calculate the Pearson correlation coefficient.
An MCC score of 1.0 indicates an estimated representation equal the ground truth up to permutation and scaling.

---

Published in Transactions on Machine Learning Research (07/2025)
Simulation details The general simulation procedure is: generate balanced binary tree T, of depth 7 with total of 128 leaves. In each non-root environment, we sample 3000 datapoints according to,
x E R 1)
2 E R ~ (.1)
(23)
((x) = tanh(w' x); w -~ (0,11
wo ~ 1)
Each are in T is assigned sparse 6, E R2, with non-zero Gaussian values (o²=0.25) at random subset of S indices, where S denotes the number of non-zero entries and is varied throughout experimentation. We is generated using Equation 2 and the target Y is sampled according to Equation || with e ~~ (O, 0.1). In Appendix A.5, we show results for an equivalent simulation where only leaf nodes are observed, more closely matching biological applications.
Disentanglement. To study the first question about validating our discntanglement thoory, we evaluated the MCC of the representations learned by TBR and the baseline methods while varying the level of sparsity in A, from S = 0, indicating that each 8 A is a null vector, to S = 5, resulting in fully dense The results, displayed in Figure 3a align with our theoretical findings. With S 0, A is completely sparse, and V is free to entangle Z without any increase in regularization cost. In this setting, TBR, "Sparse Predictors" and the baseline hots each learn entangled representations, exhibiting mean MCC of approximately 0.5.
In the 5=1 sctting, which uniquely satisfics Assumption 3.5 TBR clicits mean MCC scores of 0.98 0.01 indicating that the representations of Ž are almost fully disentangled. "Sparse Predictors" exhibits mean MCC scores of 0.81 + 0.12, and the baseline method achieves mean MCC of 0.47 + 0.06
Sensitivity to assumption violations. The settings of S > 2 test the robustness of our method to violations of Assumption 3.5 which require 1-sparse changes across environments. As S increases, we see that the MCC scores obtained by TBR gradually decline, a reasonable result given that the guarantees of no longer hold. However, TBR remains notably stronger than the baseline methods with S <= 4, indicating that tree-based regularization incentivizes partial disentanglement even in settings not covered by our theory. Some additional results exploring further variations on our simulation setting are included in Appendix A.5, showing robustness to all-linear DGPs, more complex X - Z relationships, varied distributions of X and more stochastic amounts of variation in A.
Predictive performance, We evaluate the MSE of Ý for each method while varying the sparsity level S for each ō E A. Figure 3b) highlights these results. When S = O, Y|Y |] is invariant across environments, matching the model fit by the standard baseline as well as TBR and Sparse Predictors. As expected, all methods perform well in this case. As increases, TBR and Sparse Predictors maintain low prediction error, while the performance of the baseline method quickly degrades. This aligus with expectations, as the baseline method lacks the capacity to fit the environment--peecfic effects of the latents Z.
Prediction error and estimating latent dimension. We further find evidence that varying the di- mension of 12] in order to minimize predictive error can enable estimation of the true dimension of |Z] if it is unknown, satisfying a key assumption of our method. We mcasured the MSE of Ý achieved by the baseline and TBR when varying dimensionality of 2 from to 8, while setting the average sparsity to half the dimensionality by setting cach entry in A to zero with probability 0.5. For both methods, and for four different values of the dimension [Z], setting [ < [z] resulted in elevated MSE. Notably, the MSE achieved is shown in Figure 4 While not
by TBR distinctly plateaued after reaching the correct dimensionality, conclusive, these results suggest that there is indeed potential in the analysis of prediction error with TBR as method of estimating the number of latent features present.
Causal effect estimation. We estimated the causal effects of the high-level latent variables Z on the target Y using the representation 2 produced by cach mothod as a substitute for the true latents. The goal is investigate whether disentangled representation learning is necessary for tasks such as causal inference where the semantics of the variables involved matter. Formally, we estimate the average treatment effect
8

---

Published in Transactions on Machine Learning Research (07/202) 2025)
TER Spare rredicars Sanel
Sparse Hediccrr Band
1.5
0.75
0.50
0.2
0.00
Setting
(S)
(a) Disentanglement performance (MCC)
(b) Prediction error (MSE).
Figure Assessment of model performance across various simulation settings. The x-axis indicates the number of non-zero entries in each Left: A comparison of the MCC between and the estimates 2 produced by TBR and the baseline methods. TBR achieves near perfect disentanglement when S 1. Right: Comparison of prediction error produced by TBR and the bascline methods when estimating Y.
Baseline - Latent Dim
Latent
Baseline Latent Dim
Latent
Baseline- Latent Dim
Latent
Squared
Baseline Latent Dim
0.2
0.0
9
Estimated Latent Dimension
Figure 4: An analysis on the effects of changing the dimension of 2 on MSE, across various dimensionalities of Z. For |Z| E (4,5,6), 6), TBR achieved minimal predictive loss when | IZ]. In contrast, the baseline method had no cases where the global minimum for MSE occurred when the estimated and true latent vectors were of the same dimension.
(ATE) of each latent ZE on Y,
ATF3 (P(Ydoo(Z) = 0)] :o = ==]]
which measures the change in the mean of Y when we intervene (indicated by the do operator) and increase the value of Z& by some constant 5. Since each ATE is not identified from observational data without assumptions, follow the necessary assumptions that: (i) there are no hidden confounding variables affect

---

Published in Transactions on Machine Learning Research (07/2025)
both Z and Y and (ii) for all values 2 of - < P(Z, =2|Z
2) < 1. Following Pearl et al. (2000)
with these assumptions, and from the graphical model in F Figure
each causal effect is identified by,
ATF2 = Ez F k, Z5 = 2 = 6 2 k; Zu = =]
As consequence of the linearity in the mapping from Z to Y, each ATE can be estimated by including all variables Z in a lincar regression of Y and reading off the fitted linear coofficients.
In our experiments, we estimate ATE& in ten randomly selected environments in the 1-sparse setting. We fit ordinary least squares regression to Y using the representations Ž learned using TBR and the baseline methods and compared the estimated effects to those estimated when using the ground truth Z. We nor- malized the represenlation values and calculated the mean squared error of the absolute coeffieient values, to ignore the effects of scaling. Table included in the appendix, shows the results from this study. We see that the disentangled representation obtained by TBR yields highly accurate effects (MSE=0.05 + 0.07) in this setting while the baseline method produces representations that severely introduce bias into the effect estimates (MSF=1.51 + 1.8). Further details are presented in Appendix A.6.
4.1 Single-cell RNA-seq experimentation
In addition to analyzing the behavior of TBR in simulation, we evaluated the ability of our method to learn disentangled representations of ground-truth gene expression data, with simulated downstream cell-state phenotypes.
In our sctup, both X and Z consisted of real gene expression data obscrved across different colls. We selected genes from within set of 11 broadly expressed keepenk genes (isennbey & Levanon 2013) to hold out as latent variables. For X, extracted all transcription factors (TFs) annotated as regulators of Z, in the hTFttaget database (Zhang et al. 2020). We removed from consideration two weakly regulated house- keeping genes, resulting in a set of nine candidate Z genes (see Appendix A.5), each regulated by 18 to 95 TFs within X. We performed all experimentation on the publicly available GTEx V8 aiRNA-See datasel, described by Eraslan et al. (0022). We consider the 43 epithelial cell types as the set of environments E.
Preprocessing All preprocessing of the GTEx data was performed using standard functions within the Scanpy library (Wolf et al. 2018). To denoise the data and adjust for dropout during sequencing, we first applied across the entire dataset the MAGIC function developed in Van Dijk et al. (2018). The resulting dataset was highly imbalanced, with a small number of cell types comprising the majority of examples. To prevent. these cell types from dominating our analysis, we randomly selected a subset of each cell-type population such that no individual environment had more than 1000 samples, resulting in a final dataset of 15256 cells.
We constructed the tree T relating the cell types by running the dendrogram method in the Scanpy library, such that the 43 epithelial cell types made up the leaves in T. We simulated phenotype Y using latent gene expression Z and cell type tree T via the simulation procedure outlined in |S 4
Results We repeatedly sampled subsets of 5 candidate genes to serve as Z and the corresponding TF expression X, and compared the performance TBR and our previously described bascline method across 30 different simulated phenotypes, for each sparsity level S.
The performance of the two methods in disentangling the latent gene expression values are depicted in Figure 5a TBR and the bascline method exhibit equivalent average MCC scores of 0.51 and 0.00 respectively in the S = 0 setting, when there is no variation between environments. Once variation is introduced, the performance of TBR increases, exhibiting an average MCC of 0.61 + 0.02 with both S = and S = 2, while the performance of the baseline method degrades rapidly to 0.31 + 0.02 as K increases to 2. Globally, MCC scores are lower than observed in simulation. We expect that much of this gap in performance can be attributed the high levels of endogenous noise within the X - relationship.
Generalization to unseen environments. Since parameter values across related environments only sparsely vary when we consider the true latents, but densely vary when we consider incorrect transformations
10

---

Published in Transactions on Machine Learning Research (07 (07/202) 2025)
2.0
0.5
0.4
MCC
MSR
Model
0.3
TBR
TBR
Baseline
0.5
0.1
0.0
0.0
0
Setting
Spersity Setting
(a) Disenlanglemennt performance (MCC)
(b) Prediction error (MSE) unseen Lest cell types.
Figure 5: Left panel: A comparison of the MCC between 2 and Z across various settings of sparsity in the simulation procedure, when X and Z are both ground-truth gene expression measurements. Results are averaged over 30 simulated phenotypes for each of 75 random choices of genes to serve as Z. Right panel: MSE exhibited when generalizing trained instances of TBR and the baseline method to unseen cell types with 1-Sparse generative parameters changes, while varying the S setting of training environments. Results are averaged over 10 simulated phenotypes each for 50 random choices of Z and held-out test cell types.
of these latents, we further expect that better disentanglement should lead to lower error when we transfer our learned predictor to unscen but similar environment. Thus, to demonstrate the value of learning partially disentangled representations of gene expression, evaluated the ability of our methods to predict phenotype in environments that were unobserved during training. We selected cell types to use as a test set, and completely excluded them from training.
We compared the performance between our baseline method, trained on the closest neighboring cell-type the test type, and TBR, using the parameters from the closest. environment the test. type in To maximize the potential performance from the baseline, we only considered test cell types whose closest neighbor had large population (pop. size > 1000). Additionally, restricted analysis to test types with population size > 100.
As in the previous experiment, we repeatedly sampled random choices of genes to make up Z and generated simulated phenotypes with varying settings of S. In all cases, the unobserved environments had Wtest ex- hibiting 1-sparse changes from the closest. observed environment T, resulting a non-trivial generalization challenge even when S for the training cnvironments. The results averaged across all choices of are presented in Figure 5b)
Notably, training TBR with S resulted in test MSE of 1.108 in the unseen cell types, significantly
higher error rate than the S = 1 and S = 2 scttings (respective MSE of 0.49 + 0.22 and 0.47 + 0.20). This improvement. in MSE when S > 0 is expected, as the 1-sparse changes in the unseen Wiest could become more widely distributed across the entangled estimated of 2. This reduction in generalization risk in the S >= setting illustrates the value of even partially discntangled representations. The bascline method exhibits significantly higher MSE all settings, which may be due the inherent limitations of training on only closely related cell types, restricting sample size.
11

---

Published in Transactions on Machine Learning Research (07/2025)
5
Related work
Several papers have shown that with samples from multiple environments, auxiliary labels, and assumptions on the data generating process, we can learn disentangled representations (Hyvarinen & Morioka, 2016
2017; Hyvarinen et al., 2019 Khemakhem et al., 2020a b Von Kugelgen et al.) 2021 [Shen et al., 2021
Klindt et al.]|2021 2021 Gresele al. al.]] 2019). Most of these previous works assume the latent variables are causal ancestors of the observed variables, whereas we consider latent variables which are causal descendants of the observations, enabling use of the target distribution as signal for latent identification, rather than a reconstruction objective. Within the research on disentanglement, this paper relates most closely to two lines of work that exploit sparsity to guarantee discntangled solutions to the representation learning task.
Sparse changes to latents. One line of work focuses on temporal observations (Lippe et al. 2022) or paired samples (Locatello et al. 2019 Ahuja et al.) 2022b Von Kugelgen et al. 2021 generated by vary- ing only a small number of generative factors. Similarly, Lachapelle et al. (2022b) 2024) leverage sparse interactions between latent factors across time and/or with auxiliary variables
learn disentangled repre- sentations. Instead of sparse latent differences between paired samples or units across time, recent works
Ahuja et al. (2023); Buchholz et al. 20223 Squires et al. (2023); Zhang et al. (2023); Varici et al. (2024);
von Küggeggen et al. (2023) consider unpaired data from multiple environments that sparsely differ due to
single-node interventions to the latent generative factors. In this paper, instead of leveraging sparse shifts in
the distribution over latent variables P(Z), we exploit sparse changes in P(Y (Z), the mechanism by which
the latent variables affect a target variable.
Sparse mappings. Another line of work exploits sparsity in the mappings from latent to observed variables to establish disentanglement. In the case of generative factors, Moran et al. (2022) Brady et al.
20223 Zheng
et al. (022) assune that each latent factor alfects a subset of the components of the observalion. In the latent features setting that we also study in this paper, Lachapelle et al. (2022a); Fumero et al. (2023) observe multiple targets of interest that depend sparsely on the latent features. In contrast, in this paper, we focus on sparse changes in the mapping from latent. features to the single larget variable of interest. Another key distinction is that we leverage a hierarchical structure over environments. This allows us to enforce that mechanism changes be sparse locally while still admitting changes that are dense between environments that have diverged significantly. Thus we do not need to constrain all pairs of environments to vary sparsely.
Tree based regularization Layne et al. 2020) employ regularization scheme similar to TBR, with the goal of improving generalization performance when modelling phylogenetically distributed datasets. Layne et al. (2020) leveraged the assumption distributional shifts of the features or labels, and applied Le regularization to changes in parameters across neighboring nodes in a phylogenetic tree in order to improve predictive performance. In contrast, in this work we explicitly assume that there are sparse changes in the effects of latent variables, and show how regularizing sparsity across neighboring environments can be used to learn disentangled representations of these latents.
6 Conclusion
The ability to learn disentangled representations is an important, prerequisite for the use of representation learning in a number of scientific tasks, such as causal effect estimation. There is a growing attention in the literature to the use of sparse variations between environments to achieve disentanglement. We make several novel contributions towards this line of pursuit. We demonstrate that disentanglement is possible when data originates from a set of evolving environments, such as those resulting from cellular differentiation or species evolution, where parameter shifts along a given tree branch are sparse. Tree-branch sparsity suffices to enable disentanglement, even if the resulting pair-wist wise differences between observable environments may be dense.
We present novel theoretical results, leveraging sparse tree-based regularization in parameter space to disen- tangle representations of learned features, downstream from the observable inputs in the causal graph. We validate the value of our approach through a scries of simulation studics, that demonstrate improvements in disentanglement scores, prediction error, and estimation of the causal effects of the latent features. We further investigate the ability of TBR to learn disentangled representations of ground truth gene expres-
12

---

Published in Transactions on Machine Learning Research (07/2025)
sion data, with simulated downstream cell phenotypes. While further improvements required to achieve the full potential of causal representation learning on scRNA-seq data, we observe promising trends. Our findings suggest that there is value in leveraging information about cell-type differentiation with TBR-style methodology, both to assist disentangling latent molccular mechanisms and to improve generalization performance.
There are multiple directions for future work, including better characterization of the number of environments required, further exploration into estimation of the number of latents, and expanding results into more flexible generative processes, such by relaxing the assumption that p(Z X) remains invariant.
A future direction of particular interest is the generalization of the linear mapping between latents and the larget, especially in the context of gene expression data. The regulatory activity of transcription factors on downstream genes can be non-lineat Methods developed by Behr et al. (2022 2024) identify interacting pairs of genes with the assumption that the underlying genetic regulation can be modelled as set of thresholded Boolean interactions belween gene expression values. Our experimentation in 1.. considered approximating this regulatory dynamic in X Z, where (xx) can approximate "spiky" regulatory functions due to use of Relu activations. However, generalizing our identification results to permit non-linear, non-smooth Z + Y would strengthen the use of TBR for the study of gene regulatory relationships.
Software and Data
Relevant code for replicating experimentation is released on Github Gene expression data can be retrieved via the GTEx portal website.
Acknowledgments
We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC).
Nous remercions le Conseil de recherches en sciences naturelles et en génie du Canada (CRSNG) de son soutien.
The authors would like to thank Paul Bertin and Alex Tong for helpful discussions.
References
K. Ahuja, D. Mahajan, V. Syrgkanis, and I. Mitliagkas. Towards efficient representation identification in supervised learning. In First Conference on Causal Learning and Reasoning, 2022a.
Kartik Ahuja, Jason Hartford, and Yoshua. Bengio. Propertics from mechanisms: an equivariance perspective on identifiable aepresentation learning. In International Conference on Learning Represenlations, 2022b. URL /ttpp:://openreviex htt.:ppwwwwaaai....eeeeoo/aaaaaa
Kartik Ahuja, Jason Hartford, and Yoshua Bengio. Weakly supervised representation learning with sparse perturbations. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in __Tt
Newral Information Processing Systems, 2022c. URL http:p////opaerreeeeeeeettrrrrra
Kartik Ahuja, Divyat Mahajan, Yixin Wang, and Yoshua Bengio. Interventional causal representation learning. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Merle Behr, Yu Wang, Xiao Li, and Bin Yu. Provable boolean interaction recovery from tree ensemble obtained via random forests. Proceedings of the National Academy of Sciences, 119122)111111861111 2022.
Merle Behr, Karl Kumbier, Aldo Pardova-Palomera, Matthew Aguirre, Omer Ronen, Chengzhong Ye, Euan Ashley, Atul J Butte, Rima Arnaout, Ben Brown, et al. Learning epistatic polygenic phenotypes with boolean interactions. Plos one, (902200029 2024.
13

---

Published in Transactions on Machine Learning Research (07/2025)
Jack Brady, Roland S Zimmermann, Yash Sharma, Bernhard Schölkopf, Julius von Kügelgen, and Wieland Brendel. Provably learning object-centric representations. Tn International Conference on Machine Learn- ing, pp. 3038 3062. PMLR, 2023.
Johann Brehmer, Pim de Haan, Phillip Lippe, and Taco Cohen. Weakly supervised causal representation learning, October 2022. URL
16437 arXiv:2203.1442 [cs, stat].
Simon Buchholz, Goutham Rajendran, Elan Rosenfeld, Bryon Aragam, Bernhard Schölkopf, and Pradeep Ravikumar. Learning linear causal representations from interventions under general nonlinear mixing. In Advances in Newral Information Processing Systems 36 (NewrIPS 2023), 2023.
George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4):303-33) 3-44, 1989.
Eli Eisenberg and Erez Y Levanon. Human housckoeping genes, revisited. TRENDS in Genetics, 29 569 574, 2013.
T Enver, S Soneji, C Joshi, J Brown, F Iborra, T Orntoft, T Thykjaer, E Maltby, K Smith, R Abu Dawud, Jones M, Matinm M, Gokhalem P, Draper, and PW. Andrews. Cellular differentiation hierarchies in normal and culture-adapted human embryonic stem cells. Human Molecular Genetics, 14:3129-38 -94-0, 11 2005.
Gökcen Eraslan, Eugene Drokhlyansky, Shankara Anand, Exgenij Fiskin, Ayshwarya Subramanian, Michal Slyper, Jiali Wang, Nicholas Van Wittenberghe, John M Rouhana, Julia Waldman, et al. Single-nucleus cross-tissue molecular reference maps toward understanding disease gene function. Science, 376(6594): cabl4290, 2022.
Simon Freyaldenhoven. Tdentification Through Sparsity in Factor Models. Working Papers 20-25, Federal Reserve Bank of Philadelphia, June 2020. URL http:p///.iaa...gee...oo/////p//
html
Marco Fumero, Florian Wenzel, Luca Zancato, Alessandro Achille, Emanuele Rodolå, Stefano Soatto, Bern- hard Schölkopf, and Francesco Locatello. Leveraging sparse and shared feature activations for disentangled representation learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
URL. https::t//// htttps:///opnrevieer
L. Gresele, P. K. Rubenstein, A. Mehrjou, F. Locatello, and B. Scholkopf. The incomplete rosetta stone problem identifiability results for multi-view nonlinear ICA. In Conference Uncertainty in Artificial Intelligence, 2019.
Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. Advonces in Neural Information Processing Systems, 29, 2016.
Aapo Hyvarinen and Hiroshi Morioka. Noulinear 1CA of Temporally Dependent Stationary Sources. The 20th International Conference Artificial Intelligence and Statistics, 2017.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Newral networks, 12(3):429-439, 1999.
Aapo Hyvarinen, Hiroaki Sasaki, and Richard E Turner. Nonlincar ICA Using Auxiliary Variables and Gen- eralized Contrastive Learning. The 22nd International Conference on Artificial Intelligence and Statistics, 2019.
I. Khemakhem, R. Monti, D. Kingma, and A. Hyvärinen. Ice-beem identifiable conditional energy-based deep models based on nonlinear ica. In Advances in Neural Information Processing Systems, 2020a.
Nyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 2207-2217. PMLR, 26-28 28 Aug 2020b. URL https: https::/prooeeedngg. mlr, wwtpppsss//.////aaaaaaaaaaaeeeea
html
14

---

Published in Transactions on Machine Learning Research (07/2025)
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Represenlations, 12 2014.
David A. Klindt, Lukas Scholt, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding. In International Conference on Learning Representations, 2021. URL https://openrevies..nee. htt.e//iiiw..tt/ooouuuddd
EbIDjBynYJ8
Sébastien Lachapelle, Tristan Deleu, Divyat Mahajan, Ioannis Mitliagkas, Yoshua Bengio, Simon Lacoste- Julien, and Quentin Bertrand. Synergies between disentanglement and sparsity: a multi-task learning perspective. arXiv preprint orXiu:2211 14666, 2022a.
Sebastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi LE PRIOL, Alexandre La- coste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ICA. In First Conference on Causal Learning and Reasoning, 2022b. URL
Sébastien Lachapelle, Pau Rodríguez López, Yash Sharma, Katie Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Nouparametric partial disentanglement via mechanism sparsity: Sparse actions, interventions and sparse lemporal dependencies, 2024.
Elliot Layne, Erika N Dort, Richard Hamelin, Yue Li, and Mathieu Blanchette. Supervised learning on phylogenetically distributed data. Bioinformatics, 36(Supplement 2):i895- 12 2020. ISSN 1367-4803.
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M. Asano, Taco Cohen, and Efstratios Gavves. CITRIS:
Causal Identifiability from Temporal Intervened Sequences, June 2022. URL http://arxiv.org/
03169 arXiv:2202.03169 [cs, stat].
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Ractsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled represen- tations. In international conference on machine learning, pp. 4114 4124. PMLR, 2019.
Gemma E Moran, Dhanya Sridhar, Yixin Wang, and David M Blei. Identifiable variational autoencoders via sparse decoding. Transactions of Machine Learning Research, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, pigh-performance deep learning library. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d' d'Alch--Bue, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Judea Pearl al. Models, reasoning and inference. Cambridge, UK: Cambridge UniversityPress, 19(2). 2000.
G. Roeder, L. Metz, and D. P. Kingma. On lincar identifiability of learned representations. In Proceedings of the 38th International Conference on Machine Learning, 2021.
X. Shen, F. Liu, H. Dong, Q Lian, Z. Chen, and T. Zhang. Disentangled generative causal representation learning, 2021.
Chandler Squires, Anna Seigal, Salil Bhate, and Caroline Uhler. Linear causal disentanglement via interven- tions. In Proceedings of the 40th International Conference on Machine Learning, 2023.
David Van Dijk, Roshan Sharma, Juozas Nainys, Kristina Yim, Pooja Kathail, Ambrose J Carr, Cassandra Burdziak, Kevin R Moon, Christine L Chaffer, Diwakar Pattabiraman, et al. Recovering gene interactions from single-cell data using data diffusion. Cell, 174(3):716-72) 2018.
Burak Varici, Emre Acartürk, Karthikeyan Shanmugam, Abhishek Kumar, and Ali Tajer, Score-based causal representation learning: Linear and general transformations, 2024.
15

---

Published in Transactions on Machine Learning Research (07 (025))
J. Von Kugelgen, Y. Sharma, L. Grescle, W. Brendel, B. Scholkopf, M. Besserve, and F. Locatello. Self- supervised learning with data augmentations provably isolates content from style. In Advances in Neural Information Processing Systems, 2021.
Julius Kügelgen, Michel Besserve, Liang Wendong, Luigi Gresele, Armin Kekić, Elias Bareinboim, David M. Blei, and Bernhard Schölkopf. Nonparametric identifiability of causal representations from unknown interventions. In Advances in Neural Information Processing Systems 36 (NeurIPS 2023), 2023.
F Alexander Wolf, Philipp Angerer, and Fabian J Thcis. Scanpy: large-scale single-cell gene expression data analysis. Genome biology, 19:1 5, 2018.
Jiaqi Zhang, Chandler Squires, Kristjan Greenewald, Akash Srivastava, Karthikeyan Shanmugam, and Car- oline Uhler. Identifiability guarantees for causal disentanglement from soft interventions. In Advances in Newral Information Processing Systems 36 (NeurIPS 2023), 2023.
Qiong Zhang, Wei Liu, Hong-Mei Zhang, Gui-Yan Xie, Ya-Ru Miao, Mengxuan Xia, and YuaYuan Guo. htftarget: comprehensive database for regulations of human transcription factors and their targets. Genomics, proteomics & bioinformatics, 18(2):120 128, 2020.
Y. Zheng, 1. Ng, and K. Zhang. On the identifiability of nonlinear ICA sparsity and beyond. In Advances in Neural Information Processing Systems, 2022.
A Appendix
A.1 Proof of Proposition 4.4
Proof. By Assumption 3.1 and the conditions on Ý outlined in Proposition 4.4, we have
EE X, e] u (X) w V (X) = . (X)
(24)
(25)
Let el be the environments given by Assumption 3.2
We can define matrices
We
U and U =
W
By Assumption 3.2 we know U is invertible, which allows write
UV(2) Úv(2)
(26)
(xx) = U~ÚÚ
(27)
((x) = LQ(x))
(28)
where L = U Ú
By Assumption 3.3 we can create an invertible matrix Q = [V(()))..., ((+))]. We can subsequently denote
Q = (xxx) (x) }] allowing us to write Q = -QQ. Since Q is invertible, both L and Q must also be
invertible.
Combining Eq. and Eq. 28 vields
w LV(x)) T (x)
(29)
LQ 10_ Q
(30)
w T L 1.
(31)
16

---

Published in Transactions on Machine Learning Research (07 (025)
where the last equation holds because Q is invertible.
Choose some A and assume a = (eo, c)). Since
Eq. 31
holds for all it holds in particular for and ej:
6g w
(32)
W 6
T
JT
(33)
By substracting both equations above, obtain
(woy - Wen
L (106) Wen
(34)
80 L 82
(35)
where the last step above follows from the definition of we in Eq. 2 This concludes the proof.
A.2 Useful Lemmas
The argument for proving the following Lemma is taken from Lachapelle et al. (202bb)-
Lemma A.1 (Sparsity pattern of an invertible matrix contains a permutation). Let L E Ruxm be an invertible matriz. Then, there erists a permutation π such that Li,(()) * O for all i.
Proof. Since the matrix L is invertible, its determinant is non-zero, i.c.
det(L) == Σ
sign (T) II
Li,(t)) # 0,
(36)
TE
1=1
where 6, is the of m-permutationss This equation implics that at least one term of the sum is non-zero, meaning there exists π 6m such that for all e [m], Lix(()) # 0.
A.3 Remarks on S
A.3.1 On relaxing the 1-sparse assumption:
Assumption 3.5 is particularly stringent, and many problem settings will not adhere to in practice. We note some of the expected results of relaxing this assumption, as they relate to our theory.
Central to our proof are the two sets I, and Ω4 Ω2 is the set of indices where entanglement by L decreases sparsity for column i, and I, the set of indices where sparsity is increased. We reformulate the definition of Ω
S (A)))(()) L. L-0((),
(37)
One can see that the condition for index to in Q equivalent to: Aj.(()) Lm(i), 0. Thus, the number of indices in Ω2 = a for some 2 2, implies that there is a subset of rows are lincarly dependent. Under the assumption that non-zero values of A are generated independently, it seems intuitive that the likelihood of finding a rows with linearly dependent columns decreases rapidly a increases, and the results in Figure 3a give some evidence towards this possibility. We leave for future work the possibility for formalizing a probabilistic bound on the size of a. We also note that Lachapelle et al.) find an identification result for latent variables that holds in a setting with infinitely many environments but more general assumptions about sparsity, which may have some connections to the probabilistic line of reasoning we outline.
where the columns
An additional possible direction for future work towards generalizing our result investigate the use of 'anchor environments' used by Moran et al. disentangle latent generative factors.
17

---

Published in Transactions on Machine Learning Research (07/2025)
Table 1: Summary of hyper-parameter settings.
Experiment
/ Dataset
Parameter name
Setting or range
Simulated
0.001
Simulated
Learning rate
0.001
Gene Expression
(0.0, 0.1, 0.01, 0.001, 0.0001) (0.001, 0.0001)
Gene Expression
Learning rate
A.4 Additional experimental details
Implementation We instantiated & as a neural network with two hidden layers (256 units, 64 units) and LeakyRelu activations. Performance was compared across methods of estimating linear map 4. First, TBR implementation was used to directly estimate dhy and A in order to predict y. Regularization of |Âllo was approximated by instead regularizing || We compare against a baseline of training single estimation of u without accounting for variation between environments.
In each experiment, I and the relevant estimation of & were jointly optimized with the Adam optimizer Kingma & Ba 2014). We trained models on 50% of the data. validation set of 25% the data was used tune hyper-parameters (learning ratc, for TBR models), and the remaining 25% was held out as a test set for evaluation. All models were implemented in Py Paszke al. 2019).
Hyper-parameters
Hyper-pane parameters are summarized in Table Gene expression experiments were tuned individually per choice of Z and phenotype from the range listed.
List of house-keeping genes
Sets of genes for inclusion Z described in s 4.1 were selected from the list in Table
Description of error estimates
All results presented with a mean value accompanied by + error estimate indicate standard deviation, computed with a call to numpy's sp. command. Box plots, such as Figure 3a include quantiles and outliers calculated by scaborn's sns. catplot () function with default scttings. Error bars on bar plots, such as Figure 5a indicate the 0.95CI and are computed by the errorbar function included in sns. ens. catplot with default settings.
Compute requirements
Experimentation on the fully simulated dataset was performed on MacBook Pro with 18GB memory and an M3 Pro CPU, Time to generate a simulated phenotype with K = and fit an instance of TBR was estimated at 42.98 + 13.40 over 10 repetitions.
Experimentation on gene-expression data described in $ 4.1 was performed on shared compute server with total of 755GB of ram and 96 CPU cores. Training leveraged shared use of Quadro RIX 6000 GPU, with average GPU memory usage of approximately 1GB. Run-time to fit instance of TBR for random choice of Z was estimated at 25.37 + 2.53 seconds over 10 repetitions.
A.5 Additional experiments and results
18

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Table 2: House House-keeping genes
Gene Name
Clorf48
CHMP2A
EMC7
PSMB2
PSMB4
REEP5
SNRPD3
VCP
VPS29
10
69
6.-
1.50
6.-
WCC
6.6-
0.7
65
0.2
0.00
(a) Disentanglement performance (MCC).
(b) Prediction error (MSE).
Figure 6: Assessment of model performance across various simulation settings, with observations available only that the leaf nodes.
Simulation with observations exclusively at leaves
Complementary to the experimentation depicted in Figure 3 we performed an equivalent set of experiments with the simulation altered such that only leaf nodes had observations, which more closely matches our motivating biological applications. We increased the depth of the tree from to 8, keeping the number of observable environments similar. The results display the same trends as discussed in the main text, with an MCC approaching 1.0 exhibited by TBR the observed by both methods when K 0.
sctting, and significant increase entanglement
Simulation with more complex X + Z relationship
We further tested the robustness of the TBR method when increasing the complexity of the function relating X to Z, by increasing the number tanh layers from to 2. MCC scores were slightly reduced across both methods for all settings, but MCC scores remained above 0.9 for TBR in the setting, Results are depicted in Figure
Simulation with skewed distributions for X
19

---

Published in Transactions on Machine Learning Research (07 (07/202) /025)
1.
TSA Soschiee
Mood
1.00
MEC
0.50
0.00
Searcht Setting
Serving (S)
(a) Disentanglement performance (MCC)
(b) Prediction error (MSE).
Figure 7: Assessment of model performance across various simulation settings with a more complex X + Z lunction.
TSA
1.c
Susciree
Mosi
MDC
02
Searsky
(5)
(a) Disentanglement performance (MCC)
(b) Prediction error (MSE).
Figure 8: Assessment of model performance across various simulation settings with non- non-Gaussian distribu-
tions of X.
We further tested the effects of considering alternative distribution for X in our simulation studies, by distributing each X; as Gaussian with a skew value drawn uniformly from (-1.0, 1.0). Our simulation results appcared robust this modification. Results are depicted in Figure
Causal effect estimation in simulation experiments.
20

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Table 3: Mean-squared error of effect size estimates calculated using Z, in comparison to estimates calculated using ground-truth Z, computed over 10 trials per environment. TBR achieves a much smaller mean error across all environments, showing that disentangled estimates of 2 allow for increased accuracy in determining the effect sizes of latent features.
Environment
TBR MSE
Baseline MSE
Env 1
0.05 + 0.07
0.87+0.76
Env
0.05 + 0.07
1.08#0.91
Env 3
0.03 + 0.04
1.29+1.67
Env 4
0.06 1 0.09
1.68.11.72
Euv 5
0.05 1 0.08
1.96.13.58
Env 6
0.06 + 0.09
2.55+2114
Env
0.05 + 0.08
1.3+0094
Env 8
0.05 + 0.07
1.85
Env
0.04 + 0.06
1.19
Env 10
0,04 + 0.06
1.311.29
Pooled
0.05+0.07
1.511188
The results from estimating causal effects from the estimated Ż using either TBR or the baseline method, described in are summarized in Table 3
All-linear data-generative process
We tested the effect of an all-linear DGP, where (X) was replaced with a linear map with random weights. We maintained the noural network architecture of V. We noticed minimal change in the MCC performance of TBR, results are included in Figure
Stochastic Deltas
Additionally, tested the effects of randomly selecting the number of non-zero values in cach 8 by drawing from a Bernoulli distribution, while varying the parameter in order to control the expected level of sparsity. Results are displayed in Figure 10
21

---

Published in Transactions on Machine Learning Research (07/202) 2025)
0.9
0.7
Model
TBR Baseline
0.5
0.4
0.3
0.2
K-Sparse Setting
Figure 9: An equivalent plot to Figure 6a obtained by using linear map in place of vp (X).
1.0
0.8
0.7
TBR
Baseline
0.4
0.3
0.0
0.1
0.2
0.3
0.4
0.6
0.7
0.8
0.9
Delta Probability
Figure 10: An equivalent plot to Figure 6a. obtained by using Bernoulli distribution of non-zero entries in A in place of a fixed k-sparse setting.
22


Reviews:
[{"review_id": "HLngdriRZs", "reviewer": "Reviewer_dfRA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\n1. The key assumption of TBR is a theoretically valid assumption. It assumes that only a sparse set of true latents will change across closely related environments. Therefore, TBR applies a sparsity-regularization term in the loss function to induce sparsity and recover true latent variables. \n\n2. This paper proves that the regularization on the matrix $\\hat{\\Delta}$ ensures the identification of true latent variables up to permutation and element-wise rescaling, as only linear transformations can be applied to $\\Delta$ while maintaining the same level of sparsity. Thus, TBR prevents the entanglement of the components of the true features. \n\n3. This paper performed experiments both on simulated data and real data to validate the effectiveness of the proposed TBR. On simulated data, this paper proves the capability for disentanglement and causal effect estimation of TBR. The experiments on real data also proved better disentanglement than baseline approaches. \n\nWeaknesses: \n1. This paper subsampled the dataset of real data during experiments to ensure balanced cell types. While this experiment demonstrates the effectiveness of TBR on balanced data, real-world data are highly imbalanced. It is unclear whether TBR would be robust even on highly imbalanced data. \n\n2. The authors could consider including more baseline approaches for empirical studies, such as the ones mentioned in the related work section. The current baseline approach might be too simple and TBR would almost certainly demonstrate improved performance. \n\n3. The proofs of TBR are valid. But more feedback from reviewers with ML background is certainly welcome.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel disentangled representation learning method for biological settings called Tree-Based Regularization (TBR). It infers latent variables from low-level observations across different environments (e.g., gene expressions from different cell types), and the learned latent causal variables can give valuable scientific insights. The proposed TBR method is validated by theoretical proofs and simulated experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Include the performance comparison of TBR and baseline approach on a real-world imbalanced dataset in the empirical studies. \n\n2. Include more competitive baseline approaches for empirical studies in addition to the current baseline approach.", "broader_impact_concerns": "After a review of the submitted work, no concerns regarding ethical implications have been identified.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "BfxKRPij25", "reviewer": "Reviewer_FWx2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**: Reasonable theoretical contributions, and well-performing algorithms. \n\n**Weaknesses**: the algorithm is not very connected to the theory.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors proposed a tree-based penalization to recover disentangled representations. The authors established identifiability results, and showed sparse penalization lead to disentangled representations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- I think the author(s) had to put some assumptions on the true $\\Psi$ and the neural net to show the *entire* procedure to identify and it cannot be a general neural net (e.g., a constant function is a weird neural net). I would encourage the authors to spell that out. It can even take the form that $\\hat{\\Psi}$ can approximate $\\Psi$ well in some clearly stated standard. Similarly, what is the assumption on $\\lambda$?\n\n- It is a bit hidden in the appendix that instead of optimizing the $L_0$ penalized objective (4), the author instead optimized \n\\begin{equation}\n\\sum_{e\\in \\mathcal{L}}(w_e^\\top \\hat{\\Psi}_{\\theta}(\\mathbf{X}^{e})-Y^e)^2+\\lambda||\\hat{\\Delta}||_1\n\\end{equation}\nThis is a significant relaxation that must be transparent in the main text. I encourage the authors to make some comments on this relaxation. It will not only encourage sparsity but also shrink values of other non-0 elements in $\\hat{\\Delta}$ towards 0, a behavior that does not exist when penalizing with $||\\hat{\\Delta}||_0$. Also, optimizing this is not that trivial: e.g., one cannot use naive gradient descent since the gradient will not exist at the sparse solution, like LASSO. In the paper, the authors just stated they implemented this in `pytorch` without details. \n\n- Choice of the penalization strength $\\lambda$ is not quite transparent, in appendix A4 the authors only stated they are tuned from the list but no details on how. Maybe I missed the detail. \n\n- In assumptions 3.1-3.2, does the invertibility also imply some dimension assumption?", "broader_impact_concerns": "N/A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xuBAU22LDl", "reviewer": "Reviewer_4mqX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Data-generation process and objectives.**\n\nThe paper considers the following data generation process:\n$$\\mathbf{Z} = \\Psi(\\mathbf{X}) + \\eta \\ ; \\quad Y^e = w_e^{\\top}\\mathbf{Z}^e + \\epsilon \\ , $$\n\nwhere $\\mathbf{X}$: observed variables, $\\mathbf{Z}$: latent variables, and $Y$: target variable of interest.\nThis formulation (observed variables *cause* the latents) deviates from the standard CRL approaches. Though this is not necessarily a weakness, perhaps even can be considered as a strength if it’s able to explain the biological processes better (along with the hierarchical environment structure).\n\nHowever, I want to emphasize two concerns.\n1.  The linear mapping from latent $\\mathbf{Z}$ to target variable $Y$ can be quite restrictive (the authors acknowledge this on page 2, with some motivation for nonlinear maps). That being said, I understand that studying linear maps is a reasonable starting point.\n2. Latent variables $\\mathbf{Z}$ are repeatedly referred as causal variables — since they *cause* target variable $Y$. However, reversing the data generation process ($X \\rightarrow Z$, as opposed to $Z \\rightarrow X$) also results in losing further causal interpretation of $Z$, e.g., there is no SCM interpretation to explain and identify the relationships among $Z_k$ variables. \n\n**Sparsity regularization**: The proposed sparsity regularization is quite simple yet seems to be effective. I like this simplicity as a strength.\n\n**1-sparse perturbations assumption**: Assumption 3.5 states that all changes in the hierarchical evolution of the latent processes are 1-sparse perturbations. The authors acknowledge this limitation.  Despite this weakness (limitation), simulations show that the proposed tree-based regularization is still helpful when the assumption is violated. I find this observation as one one the main strengths of the proposed method.\n\n**Simulations**: While observations of the simulations are strong (for instance, the partial robustness against assumption violations as mentioned above and applicability to gene expression data), the synthetic data setting is quite limited. See my further comments in the requested changes.\n\n**Insufficient explanations and organization of the paper**: Overall, the paper’s writing is good enough. However, at times, it leaves the reader yearning for more explanations (e.g., discussion of assumptions). This is not a major problem, but I have a number of suggestions/requests below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies learning disentangled representations using data from multiple environments and tasks. The focus is especially on the biological processes, for which multiple datasets (environments) are related by sparse changes and can be expressed in a hierarchical tree structure. For this setting, under a number of assumptions (most importantly, 1-sparse perturbations and linear mapping from latents to target), the proposed tree-based regularization is shown to help identify the latent variables. In simulations, the proposed method is tested on both synthetic data (which is a bit toy-ish) and gene expression data. Overall, the empirical results show promise that the proposed regularization technique can be helpful in practice.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### **Organization and writing related points**\n\n- Page 1. \"*The challenge is in finding assumptions that are strong enough to rule out spurious solutions, while remaining flexible enough to fit the domain of interest*\"\n\tI think this sentence should be more comprehensive. Currently, it only explains the *model* side, that is assumptions on the data generating processes. However, there is also *data* side (non i.i.d data, additional structures/information on the data etc.) as the paper leverages via sparse changes. I guess the authors mean to include everything when they say ``finding assumptions’’, but it’s not clear enough.\n\n- Page 2. \"\"*These constraints reflect biological settings, where it is common to observe data across multiple cell types or model organisms with known relationships..*\"\"\n\tCould you clarify what you refer to by known relationships, is it the hierarchical structure?\n- Section 4 first paragraph: When describing the baseline method, why not make it more explicit (representations learned with and without regularization, with eq. x vs. eq .y, entangled vs. disentangled etc.)?\n\n### **Discussion of assumptions and results**\n- It seems like latent dimension $k$ needs to be known. Please clarify it.\n- **Assumption 3.2 (sufficient task variability)**: This assumption is intuitive. However, I have two questions. First, is it possible to provide a counter-example to show that violating this assumption renders identifiability impossible? Second, while the assumption is intuitive, I think it should be made clearer earlier in the paper, at least via some hints. If I’m not mistaken, while the intro discusses the results, there is no mention of the assumption, and the reader has no idea until this point how identifiability will be possible (under tree-based regularization).\n- **Assumption 3.3 (sufficient representation variability)**: Why not have a comment/explanation about the role of this assumption? I guess it’s a quite mild assumption, e.g., it is weaker than the full latent support assumption used in many recent CRL works.\n- **Assumption 3.5 (1-sparse perturbations)**: Again, not sufficient interpretation on the role of this assumption? For instance, it seems almost analogous to single-node interventions used in recent multi-environment (interventional) CRL works (and even before that, sparse perturbations in Brehmer et al., etc.), where the perturbations are happening at a different level here.\n- **Assumption 3.6 (sufficient perturbations)**: Again, it seems analogous to ``each node should be intervened/perturbed’’ condition. This totally makes sense, but the broader audience may not see the implications and necessity.\n- **Proof of Proposition 3.7**: If the proof is given in the main body, I suggest moving it to the right after the proof sketch. Before giving the proof, I also would like to see a more detailed discussion of Proposition 3.7, which is the main result of the paper.\n- Could you provide more intuition on whether the proposed tree-based regularization can be used for nonlinear mapping from latents to target variable? It does not have to be an affirmative answer, though understanding why or why not it can be extended would be useful. \n\n### **Simulations**\n- **Synthetic data setting**:  The considered setting is a bit too limited — only one setting of latent and observed dimension, only a simple nonlinear transform ($\\Phi(x) = \\tanh(w^T x)$). I suggest demonstrating the performance on at least some of the following settings: i) increasing dimension of $X$, ii) making the data less normalized, instead of using $N(0,1)$ for all $X$ variables, iii) considering a more complex nonlinear mapping from $X$ to $Z$, perhaps multi-layers of tanh.\n- At the end of page 10, generalization to unseen environments is discussed. Can you explain this a bit more with respect to your modeling? For instance, is there a connection between holding out an environment (making it unseen) and violating an assumption (perhaps assumption 3.2)?\n\n### **Related work**\n- I suggest adding a couple more recent references for learning (causal) disentangled representations, e.g. for multiple environments(interventions) and multi-view CRL papers.\n- “We consider latent variables which are causal descendants of the observations”. While this is clear in problem formulation, the distinction from the previous works should be made clearer earlier in the paper. \n- Sparse changes to latents paragraph: If I’m not mistaken, the references in this paragraph are more on *paired samples*. As I mentioned above, there are also more recent interventional CLR studies that consider paired *distributions* instead.\n- I’m not familiar with Layne et al. (2020). If the regularization scheme is indeed similar, I feel like it warrants a more detailed comparison.\n\nOverall, I think the related work section need to be expanded, and perhaps giving it earlier in the paper (after introduction, before problem formulation).\n\n### **Other minor comments**\n- Page 1. Perhaps consider breaking the last sentence of the second paragraph into two sentences to improve readability.\n- Figure 2. Might be better to give the reference just after Eq. (2) instead of at the end of Section 2. \n- Proposition 3.4. Typo: equation equation. \n- Please make sure to add appropriate hyperrefs, some are missing (for instance, Appendix A.1 after proposition 3.4)\n- Figure 3 should come before Figure 4, according to discussion order.\n- I suggest moving Table 3 to the main body, while still leaving the details in the appendix.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)