Leveraging Task Structures for Improved Identifiability in Neural Network Representations
Authors: Wenlin Chen
Julien Horwood
Juyeon Heo
José Miguel Hernández-Lobato
Venue: Accepted by TMLR ()

Abstract: This work extends the theory of identifiability in supervised learning by considering the consequences of having access to a distribution of tasks. In such cases, we show that linear identifiability is achievable in the general multi-task regression setting. Furthermore, we show that the existence of a task distribution which defines a conditional prior over latent factors reduces the equivalence class for identifiability to permutations and scaling of the true latent factors, a stronger and more useful result than linear identifiability. Crucially, when we further assume a causal structure over these tasks, our approach enables simple maximum marginal likelihood optimization, and suggests potential downstream applications to causal representation learning. Empirically, we find that this straightforward optimization procedure enables our model to outperform more general unsupervised models in recovering canonical representations for both synthetic data and real-world molecular data.

Published in Transactions on Machine Learning Research (08/2024)
Leveraging Task Structures for Improved Identifiability in Neural Network Representations
Wenlin Chen"
vc337@com.ac.
University of Cambridge, Cambridge, United Kingdom Max Planck Institute for Intelligent Systems, Tübingen, Germany
Julien Horwood"
jdh950cam.ccr
University of Cambridge, Cambridge, United Kingdom
Juyeon Heo
jh2324 @com.ac.....mmmm.. uk
University of Cambridge, Cambridge, Unitted Kingdom
José Miguel Hernández-Lobato
jmh233 uk
University of Cambridge, Cambridge, Unvtted Kingdom
Reviewed on OpenReview: https: openreview. net/ net/ forum? id- W.cPrg6puO
Abstract
This work extends the theory of identifiability in supervised learning by considering the consequences of having access to a distribution of tasks. In such cases, we show that linear identifiability achievable in the general multi-task regression setting. Furthermore, we show that the existence of task distribution which defines conditional prior over latent factors reduces the equivalence class for identifiability to permutations and scaling of the Irue latent factors, stronger and more useful result than linear identifiability. Crucially, when we further assume a causal structure over these tasks, our approach enables simple maximum marginal likelihood optimization, and suggests potential downstream applications to causal representation learning. Empirically, we find that this straightforward optimization procedure enables our model to outperform more general unsupervised models in recovering canonical representations for both synthetic data and weal-world molecular data.
1 Introduction
Multi-task regression is common problem in machine learning, which naturally arises in many scientific applications such as molecular property prediction (Stanley et al.] 2021 Chen al.] 2023) and machine learning force fields (Jacobson al. 2023). Despite this, most deep learning approaches to this problem attempt to model the relationships between tasks through hcuristic approaches, such as fitting a shared neural network in an attempt to capture the joint structures between tasks. Beyond lacking a principled approach to modeling task relationships, these approaches fail to account for how we may expect the latent factors for related tasks to change. In this work, we show that by leveraging assumptions about the relationships between the latent factors of the data across tasks, in particular that they vary their causal and spurious relationships with the target variables, can achieve identifiability of the latent factors up to permutations and scaling.
A common assumption in the causal representation learning literature, known as the sparse mechanism shift hypothesis [Schölkopl 2019 Schölkopf et al. 2021 Perry et al. 2022). states that changes across tasks arise from sparse changes in the underlying cansal mochanisms. While we do not operate directly on structural causal models, our result. arises by similarly considering the implications of sparse changes in the causal graph
"Equal contribution.

---

Published in Transactions on Machine Learning Research (08/2024)
defining multi-task learning setting. We accomplish this by first extending the theory of identifiability in supervised learning to the multi-task regression setting for identifiability up to linear transformations (i.e.; weak identifiability). We then propose a new approach to identifying neural network representations up to permutations and scaling (i.c., strictly strong identifiability), by loveraging the causal structures of the underlying latent. factors for each task. We empirically validate our model's ability to recover the ground- truth latent structure of the data both in simulated settings where data generated from our model and for real-world molecular data. This contrasts with current state-of-the-ar approaches such as (Khemakhem) et al. 2020a Lu et al. 2022 whose assumptions also fit our assumed data generating process but which are difficult to train effectively and only identifiable up to block permutations and scaling of the sufficient statistics of their exponential family priors. We summarize our contributions in the following section.
1.1 Contributions
Our work extends the current identifiability literature in the following key aspects.
In contrast to prior work [Lachapelle et al. 2023 Fumero et al. 2023 which relates meta/ cotta/mutti- lask learning to identifiability via explicit sparsity constraints, this work expands these conceptual connections beyond sparsity constroints by considering the shared causal structure between tasks, This significantly reduces the number of tasks needed to recover the true representations.
2. Our method extends previous identifiability results by resolving the point-wise indeterminacies of prior work (Khemakhem et al. 2020a Lu et al. 2022).
3. Our model exlends the applicability of conditional prior models to discriminative settings at test time as our identifiability result does not require conditioning the target values during inference.
4. To our knowledge, our approach is the first to propose a conditionally factorized prior model which can achieve identifiability via optimizing the ezact marginal likelihood. This leads to improved empirical results in our experiments despite the probabilistic assumptions of our model.
5. While many works have shown that spurious correlations are a failure case of deep learning and focus
on eliminating them (Rojas-Carulaa Carulla et al. 2018 Arjovsky et al. 2019 Krueger et al. 2021 Eastwood
et al. 2022 Lu et al. 2022 Kirichenko et al.
2023), we leverage spurious features to improve the
robustness learned representations in the multi-task setting through our identifiability results.
1.2 Organization of the Paper
Section reviews prior works that are closely or broadly related to identifiable and disentangled represen- tation learning. Section 3 describes the proposed identifiable multi-task representation learning method and discusses model assumptions, thcorctical identifiability guarantecs, and potential limitations. Scction empirically evaluates the proposed method on both synthetic datasets and real-world molecular datasets. Section 3 summarizes the paper and its potential impact.
2
Related Work
This section discusses prior work in the identifiable and disentangled representation learning literature.
2.1
Disentanglement and ICA
The notion of optimizing for disentangled representations gained traction in the recent unsupervised deep learning literature when it was proposed that this objective may be sufficient to improve desirable altributes such as interpretability, robustness, and generalization (Bengio et al., 2013 Higgins et al.) 2017 Chen let al., 2016). However, the notion of disentanglement alone is not intrinsically well-defined, as there may be many disentangled representations of the data which are scemingly equally valid. Thus is not clear priori that. this criterion is sufficient. to achieve the above desiderata (Locatello et al. 2019) Tn the
2

---

Published in Transactions on Machine Learning Research (08/2024)
identifiable representation learning literature, the correct disentangled representation is assumed to be the one corresponding to the ground-truth data generating process. Thus, what is required is an identifiable representation, which must be equivalent to the causal one for sufficiently expressive model classes. In the linear case, identifiability results cxist in the classical literature for ICA, which requires non- -on-Gaussiantty assumptions on the sources for the data (Herault & Jutten 1986
Comon 1994].
2.2 Conditional Prior Models for Non-Linear ICA
Many extensions of ICA to the non-linear case have been proposed, together with significant theoretical advances. In particular, [Hyvarinen et al.] (2019) extend this by assuming conditionally factorized prior over the latent factors given some observed auxiliary variables, and propose a contrastive learning objective for recovering the inverse of the function which generated the observations. iVAE (Khemakhem et al. 2020a) further extends this to the setting of noisy observations, drawing connections with variational autoencoders Kingmaa & Welling, 2013) and enabling direct optimization via variational objective. Lachapelle et al. (2022) demonstrate that strong identifiability results remain achievable under weaker conditions on the sufficient statistics of the prior if the data generating process implics that the latent factors are governed by sparse mechanism shifts. iCaRL (Lu et al. 2022) derives analogous results for the case where the prior over the latent factors is a more general non-factorized exponential family distribution. However, the complex nature of the non-factorized prior in iCaRL requircs score matching, which is difficult to optimize in practice. Khemakhem et al. (2020b) explore general conditions for identifiability in energy -bsed models, and introduce the notion of linear identifiability. This is expanded upon in the context of classification models in Roeder itt al. (2021), showing that the representations obtained via the final hidden layer of a neural network may be identifiable up to linear transformations when conditioning on the label set. The works of Hälvä et al. (2021b;
95 Morioka et al. (202T) both obtain strong identifiability results by exploiting specific temporal or spatial structure in the encoded latents and modelling the joint distributions as dynamical systems, however their models do not translate well to the static setting, and their identifiability results remain restricted to non-linear coordinate- wise transformations of the latent variables. While Hyvärinen & Pajunen (1999):
Khemakhem et al.] (2020a) show that identifiability is not achievable without any form of conditioning in the prior, Willetts & Paige (2021); Kivva et al. (2022) recently extended the results in unsupervised
generative models to the case of models with mixture model priors. This can be seen as providing analogous identifiability results to prior methods using conditionally factorized priors, without assumptions on the obscrvability or the dimensionality of the conditioning variable. Nonetheless, these results do not apply the exact likelihood, and it remains unclear to what exlent the practical consistency and identifiability is achievable when optimizing variational surrogate objective.
2.3 Structural Approaches to Identifiability
In contrast, Brady et al. (0023) discuss identifiability results which arisc from assumptions on the structure of the mixing function, specifically targeting dual objectives of compositionality with respect to partitions of the latent factors and invertibility of the mixing function. Thus, no distributional assumptions are made on the prior. While this approach has similarities with our proposal by introducing assumptions on how partitions of the latent space evolve with respect to well-defined objects, we propose a general setting which is not restricted to representation learning in visual scenes. Furthermore, by formalizing these assumptions within our probabilistic model, we climinate the need for auxiliary terms in our optimization objective.
Recent work (Lachapelle et al.) 2023 Fumero et al. 2023 has expanded this area of research to consider the multi-task and meta-learning settings, and thus investigate the connections between identifiability and identifiable
the structure of the lcarning problem itself. However, their approach to achieving representations ions relies on introducing heuristic sparsity constraints, such as entropy and L2-norm regularizers, within a bi-level optimization objective, which turns out to be difficult to solve both in theory and in practice (Sinha et al. 2017). In addition, their approaches are less applicable in practice since a huge number of tasks are required (more than 108 in their experiments). This contrasts with the straightforward, principled and task-efficient optimization objective arising from our probabilistic model.
3

---

Published in Transactions on Machine Learning Research (08/2024)
Unstoarved groand-truth lakett
dats generation
Oloarsed input variable
Resovera iduntifiale
Recavond parmutaian-dtnifaak factse Stage
Stage
Shared
Feature
O O
Extractor
We.
Tasks
licntiffcatiin up to permotattonn
up
cossal latent ficter
liscar
trandurmation
spaziaas
observed varibble
Matti- Task Regression Network (MERN)
Malti- Tink Inear Cancal Model (MILLUM)
The workflow of our proposed method. Shapes are used to track the positions of the ground-truth and recovered latent factors. Colors are used to differentiate between causal and spurious latent factors. We assumee that the observed variable is obtained by transforming the ground-truth latent factors with some mixing function. We show that a multi-task regression network (MTRN) can recover the ground latent factors (i.c., data representations) up to linear transformation and further propose multi-task lincar causal model (MTLCM) to reduce the equivalence class for identifiability to permutations and scaling.
Figure 2:
3
Identifiable Multi-task Representation Learning
We propose a novel method that leverages task structures in the multi-task regression setting identify the ground-truth data representations up to permutations and scaling.
3.1 Problem Formulation
The assumptions of the ground-truti data generating process considered in this paper are encapsulated in the causal graph shown in Figure [ where the input variable x E X C IR,, the target variable y E R and the task index variable (1.....,. Ni are observed variables, and the latent factors z E Rd (d S n) are unobserved variables. We assume that is generated by transforming some (unobserved) ground-truth latent factors z" with some unknown injective mixing function f. Rd + X, i.e,, x fz(z). To incorporate the sparse mechanism shift hypothesis across tasks, we further assume that each task has its own partition of the ground-truth latent factors z" = z U z into a set of causal latent factors z and a set of spurious latent factors z and such partitions polentially vary across lasks. The larget variable assumed to be weighted sum of the causal latent factors, i.e., y = (w?)"zz, where w E gd are the ground-truth regression weights for task which assign zero weights for the spurious latent factors. Note that there may be latent factors that are uncorrelated with in some tasks, which can be included within zt but with zero regression weights. The spurious latent factors are assumed to be generated from the target variable with a different linear correlation function in each task t. Our goal is to recover the unobserved ground truth latent factors z² given an empirical task distribution p(t) over N₂ training tasks and an empirical data distribution p(x, ylt) for each task E (1,...,N,).
Ze
X
y
Z
Figure Assumed data generating process.
Overall, our proposed method consists of two stages as illustrated in Figure In the first stage (yellow), train multi-task regression network (MTRN) with a feature extractor shared across tasks and N, task- specific linear heads using maximum likelihood estimation (MLF). We show that upon convergence, the reprosentations learned by the fcature extractor are identifiable up to some invertible lincar transformation (Corollary 3.3). In the second slage (green), we use the assumed causal structure across tasks to define a conditional prior over the underlying independent latent factors. We show that this multi-task linear causal model (MTLCM) enables simple maximum marginal likelihood learning for recovering the lincar transformation in the representations obtained in the first stage, which reduces the identifiability class to permutations and scaling (Theorem 3.8 and automatically disentangles and identifies the causes and effects of the target variable from the learned representations.
Without loss of generality, we assume that R(y) This can achieved by standardizing practice.
4

---

Published in Transactions on Machine Learning Research (08/2024)
3.2 Stage 1: Multi- Regression Network
In the first stage, we train a multi-task regression network (MTRN) to recover the ground-truth latent. factors up to some invertible linear transformation.
Let wo. (x) = w: (x) be the output of an MTRN for task t, where WN E IR are the rogression weights in the linear head for lask i, and (x) E Rd the data represenlation produced by the feature extractor hg shared across all tasks with learnable parameters Φ. As in typical non-linear regression settings, the likelihood assumed to be Gaussian po (y) f) (o (x),or, with mean modeled by an MTRN and variance fixed to some constant o7 where 0 == (Ø. w WN.) denotes all parameters in the MTRN.
We first define lincarly identifiable (or weakly identifiable) representations in the multi-task setting. Definition 3.1 (Multi-task weak identifiability). Let and 0' be any two sets of parameters. Then, the data representations are linearly identifiable if there exists an invertible matrix A E Rdxd such that
p)( (yl t), xt
hø(x) Aha((x).
(1)
We show that data representations of MTRN are linearly identifiable if we have access to set of sufficiently diverse tasks measured by the linear dependencies among their regression weights.
Theorem 3.2. Let 0 = (Φ, (we, and 0' = ( (w', Di be any two sets of porameters such that po(u|xx()= = por (); Vl,x,y-
(2)
Assume that Span(Imi [m(hb)) Rd, iee. the vectors in the image of the feature eatractor he span the whole Rd. Suppose that there exist tasks [ti]s S ff.....,N such that at least one set of regression weights (i.e., eilher {wr /d or (wi, Jd are linearly independent. Then, the dala represenlations of the MTRN
are linearly identifiable
The proof of Theorem 3.2 can be found in Appendix A Following standard practice, we train the MTRN via maximum likelihood estimation (MLE):
0' arg max boo ()]-
Using Theorem 3.2 it is straightforward to show that MTRN trained with maximum likelihood estimation (MLE) can recover the ground-truth data representations up to some invertible linear transformation.
Corollary 3.3. Let h. : X + Rdd be the ground-trull, mapping from inpul variables to the ground-Irull latent factors, i.e., z' = (x). Assume that (mon)) = Rd. Suppose that there exist d tasks {ts) S {1,..... .N.) such that the set of ground-truth regression weights {wt,)e= are linearly independent. Assume that (5) has unique solution. Suppose that the optimization procedure for (3) converges to the optimal predictive likelihood under standard regularity conditions for MLF estimators [Gurland, 1954), ie.,
Po (y(x,, t) Px t) (())(w?)) (x), Vt,
(4)
Then, the feature extractor h& is guaranteed to recover the ground-trun truth latent factors up to some invertible linear transformation A., i.e., h (x)
Corollary 3.3 essentially states that the effective number of tasks defined by the number of independent ground-truth linear heads at least needs to be the same the number of latent dimensions to guarantee multi-task linear identifiability.
Remark 3.4. While Lachapelle et al. 2023 Proposition 2.2] prove similar proposition on MLE invariance to linear feature transformations, their proposition is built upon their Assumption 2.1 that the learned feature extractor ba is linearly equivalent to the ground truth feature extractor h. However, they do not specily under what conditions this assumption will hold for the MLE objective; they only specify conditions for the bi-level objective with a sparsity regularizer in their Scction 3. In contrast, our Corollary such conditions for MLE, i.e., Span(Im(h,) = Rd and the existence of d lincarly independent ground-tree fask-specific regression weight vectors (w])]
Corollary3.33oro 3.3 explicitly reveals
5

---

Published in Transactions on Machine Learning Research (08/2024)
3.3 Stage 2: Multi- Linear Causal Model
In the second stage, we freeze the feature extractor h& learned in the first stage and denote its representations by h = hx Corollary 3.3 suggests that h = A,z" for some invertible matrix A.. We propose a multi- task linear causal model (MTLCM) to recover the ground-truth latent factors up permutations and scaling from based on our assumed causal graph in Figure | The core idea of the MTLCM is to model the change in the causal and spurious latent factors across tasks with learnable task-specific parameters.
Let T(t) - } be collection of task-specific variables associated with task 6, which are free parameters to be learned from data, where C₂ E {0, 1] are the causal indicator variables which determine the partition of = 2, U 2, for the given task (i.e., Ctai = 1 indicates that 2j is a causal latent factor in task t and Cr, = 0 indicates that is a spurious latent factor in task t), and 7/2 are the coefficients used to generate the spurious latents from y for task t.
3.3.1 Conditionally Factorized Prior Given Task and Target Variables
Following the standard setting of generative models, the prior distribution over causal latent factors z, are assumed to be a standard Gaussian distribution:
pr (z, (z,)0,, 0,1).
which depends on the task variable t since the causal indicator variable that determines which subset of latent factors are causal varies across lasks.
According to the assumed data generating process, the target y a linear function of the latent data reprosentations z. Following the common setting of the last layer in a non-lincar regression noural network, we assume that y is generated from 2e via linear Gaussian model with the regression weights w, masked by the causal indicators C₁:
PT t) = (yl((w, c,) 2, o,)),
(6)
and that the spurious latent factors 2, generated from y via another linear Gaussian model:
pr(z.lu, = -N2..
(7)
The structured conditional prior over all latent factors given and y can then be obtained by Bayes' Rule:
pT pz i) py (y) /ec t)pr (Zaly,
pr (r(y,y, = J (z. P) (y Ze, t)pT (z, ly, dz,
Since no prior knowledge of regression weights w. is assumed, we marginalize out w, from pr pylze, 1) under uninformative prior (i.e., infiiitevvariance Gaussian prior). This makes the structured conditional prior factorize over all latent factors (see Appendix C for a derivation):
Pr t) PT
(Vzza,, Ar),
(9)
where the mean and covariance A, can be compactly expressed as:
a = v o (1- A. = diag(o; ct) + c.).
(10)
3.3.2 Linear Gaussian Likelihood
Since the data representation learned in the first stage is equivalent to zz up to some linear transformation, we assume a linear Gaussian likclihood with invertible lincar transformation A, similar to the likelihood function in a probabilistie PCA model (Tipping Bishop 1999:
PA (blz) - N oD),
(11)
where A to be learned from data, which aims to recover the ground-truth linear transformation A, for the linearly identifiable representation h.

---

Published in Transactions on Machine Learning Research (08/2024)
3.3.3 Exact Maximum Marginal Likelihood Learning
Let 4b 7) denote all parameters in an MTLCM, including the linear transformation and the task- specific parameters T()) = (e), for all tasks The marginal likelihood for MTLCM is given by
Pvp (bly,t) = 1 (z(bl t) = (b E.),
(12)
where the mean he and covariance ΣS have closed- form expressions (see Appendix D for derivation):
= yA e (1 - ct)), Σ = Adiag(o - ct) + A²) + +
(13)
Remark 3.5. The conditional prior p(z.y) y) over the latent factors is typically non- non-factorized according to the data generating process described in Section since the causal latent. factors be are parents of the target variable y, which become correlated when conditioning on y. In order to guarantee strong identifiability, iCaRL (Lu al.) 2022) paramcterizes such non-factorized conditional priors using ReLU activated energy- based models optimized by variational inference and score matching, which turns out to be difficult. to train in practice due to variational overpruning (Trippe & Turner 2018) and high computational complexity [Hyvärinen & Dayan 2005). In contrast, our proposed structured conditional prior (9) factorize over all latent factors, which, together with the linear Gaussian likelihood (11) allows us to use exact maximum marginal likelihood learning for (12) to recover the ground-truth latent factors z up to permutations and scaling from the linearly identifiable data representations
hæ(x) learned in the first stage:
arg max (p(t)(xx Po P (h
(14)
+
Remark 3.6. It is worth noting that our method has greater applicability for supervised learning than the methods that rely on learned probabilistic inverse (z/u) to extract identifiable latent factors from data such iVAE (Khemakhem et al. 2020a) and iCaRL (Lu et al. 2022). While these approaches theoretically could apply to learned representations atioss in discriminative settings by letting in such contexts since qu (z)x, y) depends on the target variable which is generally unknown at test time. In contrast, our method does not depend on y at inference time, since the identifiable latent factors can be obtained by applying the inverse linear transformation learned by the MTLCM to the linearly identifiable data representations produced by the MTRN, i.e., A A -h, This enables our model to be applicable discriminative settings at test time.
y), they are impractica
3.3.4 Identifiability Theory
We first define the concept of strictly strong identifiability in the multi-task setting. Definition 3.7 (Strictly strong identifiability) Let vb and w/ be any two sets of parameters. The latent factors are identifiable up to permutations and scaling if there exists a permutation and scaling matrix P E [Rdxd such that
Per (bly.t (hly, vh,t,
h- Ph.
(15)
We show that the latent factors of MTLCM are strictly strongly identifiable if there are sufficient variations of causal/spurions latent factors across tasks measured by the linear dependencies among the natural parameters of their conditional priors.
Theorem 3.8. Let u = lv, t] denote the conditioning variable and 2d. Assume that the learned and ground-trulh linear transformations A and A, are invertūble. Suppose that there cxist points Ug u. ug such that
L = (m(u1) - (70)), muu)) (
(16)
is invertible, where nuu)= = A, di) E R are the natural parameters of pT (z)u). Assume that 14 has unique solution. Suppose that the optimization procedure for 14 converges to the optimal marginal

---

Published in Transactions on Machine Learning Research (08/2024)
likelhhood under standard regularity conditions for mazimum marginal likelihood estimators Gurland 1954). i.e., for all h, y.
Per (bly. t) (n-chlu.t)=N(hhpia E;),
(17)
where pi and Ex are defined by Equation IS) but with the ground-truth linear transformation A., ground- truth causal indicators or and ground-truth spurions coefficients Ti- Then, the latent factors recovered by MTLCM are guoraateed to be strictly strongly identifiable.
Remark 3.9. The proof of Theorem 3.8 can be found in Appendix B The first part of the proof adapts the proof technique from Khemakhem et al. (2020a) to show identifiability up to block permutations and scaling. The second part of the proof is novel, which leverages the properties of the linear likelihood as shown in Equation (11p to further reduce the block-identiffable equivalence class to permutations and scaling of the actual ground-truth latent factors. This resolves the point-wise indeterminacies of Khhmakhhe et al.) )020aa);
Lu et al. (2022 as they are only identifiable up to block transformations.
3.4 Discussion of Model Assumptions
This section discusses some of the main assumptions underlying our model and their implications.
Regarding causal associations, our model proposes that the correlations between latent factors and the regression targets for each task be modelled as a partitioning of causal and spurious influences. We provide real-world motivating examples to justify this assumption in Appendix F One could in principle consider other cases; one where there is no correlation between a latent variable and the target, or one where the correlation between a latent variable and the target arises from a confounding variable. We note that the former case could be handled by the model by treating it as a causal variable with a regression weight of zero. In the latter case, this confounding variable would then itself be latent variable with a causal association to the target, and thus would not be unobserved. These possibilities are depicted graphically in appendíx E While it may be interesting for future work to consider the potential pairwise structure between latent variables, the simplicity of our model's optimization arises from its conditionally factorized structure.
Regarding probabilistic assumptions, while our model requires certain Gaussianity assumptions, we note that the final latent representation obtained by our model is a simple transformation only of the arbitrarily non-lincar latent representation obtained from the MTRN described in Section 3.2 Thus, the conditional Gaussian form of Equation may be viewed standard prior as in VA (Kingma & Welling) 2013). The Gaussian assumption in Equation [1] follows the standard lincar PCA model (Tipping & Bishop 1999). which is a natural choice given the linear identifiability result arising from the MTRN in Stage 1. The linear Gaussian regression model of Equation (6) is analogous to the standard predictive distribution for the final layer of regression neural networks.
While the causal and probabilistic assumptions of our approach do not constitute the most general con- ceivable case, we note that there is an inherent tradeoff between full generality and tractability. Indeed, prior work which may theoretically allow for more general causal or probabilistic models typically entail an approximation in the optimization. Further, the empirical results real-world data of Section 4 suggest that our approach may indeed be robust to moderate mis-specification.
4 Experiments
This section empirically validates our model's ability to recover canonical representations up to permutations and scaling for both synthetic and real-world datæ We contrast our model with the more general identifiable models of iVAE [Khemakhem et al. (020a) and iCaRL (Lu et al. 2022). For fair comparison, we also consider the multi-task extensions of iVAE and iCaRL, namely MT-iVAE and MT-iCaRL, which include the task variable t in the conditioning variables u in their conditional priors (z|u), with the task-spocific parameter T(I) = }v.} to be learned from data, which is the counterpart to T(1) = (c,, 7:3 in our MTLCM but has no explicit interpretations with respect to a causal graph. We note that while the works of Fumero ct al. 2023): Lachapelle et al. 2023) also consider methods for idcntifiability arising from learning across
2Qur code available at attps://githhb com/ jch.rvoo/ooo//nccn
8

---

Published in Transactions on Machine Learning Research (08/2024)
Table 1: Identifiability performance for recovering the linearly transformed syuthetic latent factors measured by strong MCC (%).
#Causal
2
# Balkent/Ohherveed /Ohervved
3/3
5/5
10/10
20/20 20
50/50
5//5
10/10
20/20
50,50
iVAE
57.75_5.00
78.0220.73
81.3610.57
6..30
81.9620.07
$1.67::2.97
74.29_0.30
77.57
79.79100110
iCaRL
75.22-6.40
74.55.12.09
72.37 1.22
79.43 0.52
80.0011..0
60.98.11.38
66.00 3.00
71.54 1.65
1..67
MT-IVAE
91.78=8.12
90.14=5.001
99.89-1004
97.90=1.51
90.56=3.118
76.0927.69
76.36=2.32
98. 088888
94.5322.49
MT-iCaRL
81.09-3.37
71.12.1..9
76.13.0..33
79.26_1.00
81.3010.844
61.55.11.26
64.04_1.08
72.7911.92
79.54.10.99
MTLCM
99.95.10.01
99.96.10.0
90.7740.116
99.70-10.16
98.97 0.555
99.71 99.95
99.51 10.36
99.14:10.27
Table 2: Identifiability performance for recovering the non-linearly transformed synthetic latent factors mcasured by strong MCC (%). The weak MCC (%) for MTRN also reported.
#Causol
12
#1 alteett// Coaontt/Oherror /Ohserved
20/50
20/100
20/200
20/50
20/100
20/200
20/50
20/100
20/200
WAA
73.11-1.113
77.42.1..20
76.95
65.18 65.18_1.49 1.49
68.66.10.14
69.05
58.70
50.85.10.31 60.33..0.27
ICaRL
56.70=3.49
13.29:1.26
58.442288
57.09=2.41
60.6622.71
61.0222.43
52.93=2.13
58.8021881
54.10:2.54
MT--VAE
71.78-1.45
80.14:0.37
73.8912.96
63.44_1.60
09.3110.35
68.5610.34
55.79-1.61
60.5610.23
50.6110.30
MT-ICaRL
67.57-1997
7..2613.22
65.52
63.37-0.84
0.84
63.75.12.19
61.61:1..22
57.13-1.007
6..56.00155
58, 58.10.11.06
MTL.CM
93.31+1.10
97.9440.71
97.44+0.68
95.67-0.16 16
98.12+0.70
89.05-+0.97
95.75-+0.11
96.28- + 1.20
84.28-1 1.27
MTRN (weak)
59.38=0.71
96.15=0.91
96. 19.1970.87
93.96=0.22
97. 07.63=0.79
87.75=0.99
95, 96.14=0.17
96.12=1.22 .27
53. 83.70=1.22
tasks, their approaches effectively implement meta-learning setting (i.e., require that the support and query sets be disjoint in the bi-level optimization process). The assumption on task support variability for the latter also requires a much larger number of tasks (more than 10" tasks as in their experiments) than what we consider here. These methods are thus not particularly well suited to comparison with our approach and we omit them from our baselines. Detailed model configurations can be found in Appendix G Each experiment run until convergence and repeated across 5 random sceds to guarantce reproducibility.
4.1 Synthetic Data
We first validate our approach in the situation when the data generating process agrees with the assumptions of our models. For each task, we first sample the causal indicator variables cz. The causal latent factors x are then sampled from a standard Gaussian prior. These are then linearly combined according to random weights wF to produce observed targets y with task-dependent noise corruption. Finally, the spurious variables z are generated via different weightings YÝ of the target y. This mirrors the causal data generating process described in Section 3 For the linear case, we generate observed data using random linear transformations of the ground-truth latent factors. For the non-linear case, we extend this to non-linear transformations parameterized by randomly initialized neural networks and demonstrate that our approach can be combinod with the multi-task identifiability result up to linear transformations to recover permulations and scaling of the ground-truth. The detailed experimental setup can be found in Appendix H
4.1.1 Linear Case
We study the ability of our proposed multi-task linear causal model (MTLCM) to recover the latent factors up to permutations and scaling via the Mean Correlation Coefficient (MCC) as in Khemakhem et al. (2020a) The synthetic data is generated by sampling 200 tasks of 100 samples cach. Each task varies in its cansal indicator variables et, causal weights w;, and spurious coefficients 7²- We then transform the ground-truth latent factors z" with random invertible matrix A, shared across all tasks to obtain linearly identifiable representations h. Identifiability in this setting is assessed by directly computing the MCC score between the representations obtained from our MTLCM and the ground-truth latent factors, which is referred to as strong MCC. Since the data is known to be linearly identifiable, we use linear likelihoods for the baselines.
In Table [ we show that MTLCM manages to recover the ground-truth latent factors from h up to per- mutations and scaling, and the result is scalable as the number of latent factors and the number of causal factors increase. In contrast, iVAE, iCaRL and their multi-task extensions underperform our model by
9

---

Published in Transactions on Machine Learning Research (08/2024)
Table 3: Identifiability performance for the latent factors learned on the superconductivity dataset measured by strong MCC (%). The weak MCC (%) for MTRN is also reported. "_" indicates divergence of optimiza- tion during training.
Latent dim
10
20
80
IVAE
32.87+1.16
33.21=1.04
30.68+00.39
37.41+0.84
45.52++0.81
iCaRL
32.23=0.61
35.62=+0.40
32.58+2.16
32.19+2.45
MT-iVAE MT-iCaRL MTLCM
3..58
0.80
0..68
35.144
44.49:10.96
42.26+2233
98.90=0.033
96.93+0.12
84.56=1.11
46.31
48.9442.16
MTRN (weak)
98.85+00033
97.17=0.04
93.23=0008
78.58+0.09
52.02+00.19
large margin in most cases. We find that for all tasks, the learned causal indicator variables also exactly match the ground-truth and the results from conditional independence testing (Chen 2021 Lu et al.) 2022) Ablation study for the effects of the learnable paramcters and the type of lincar transformation can be found in Appendix !
4.1.2 Non-Linear Case
A more general analysis of the identifiability of our proposed approach to consider the extension of the linear experiments to the setting of arbitrary transformations of the latent factors. For this, consider the case where random (non-lincar) MLPs are used to transform z" into higher dimensional observations x. By Corollary 3.3 it is possible to recover linearly identifiable representations h of the data by training standard multi-task regression networks (MTRNs). Identifiability in this setting is assessed by first perform- ing a Canonical Correlation Analysis (CCA) as in Roeder et al. (2021). which linearly maps the obtained representations such that they maximize the covariance with the ground-truth latent factors. The resulting mapped representations can thus be compared with the ground-truth latent factors via the MCC score. This is referred to as weak MCC, which quantifies the linear identifiability of the learned representations from MTRNs. We further train our MTLCM on the lincarly identifiable representations h obtained from the MTRN to obtain idiali representations up to permulations and scaling. Identifiability in this setting is assessed by directly computing the MCC score between the representations obtained from our MTLCM and the ground-truth latent factors as in the linear case (i.e., strong MCC). We assess this for various di- mensionalities of the observed data and for different settings the causal variables, where we generate 500 tasks of 200 samples each to improve convergence of the multitask model. The MTRN and the likelihoods in the baselines are parameterized by one-hiddeen-layr MLPs.
In Table 2 we find that the strong MCC for our MTLCM is able to match the weak MCC for the MTRN. In contrast, the strong MCC for iVAF, iCaRL and their multi-task extensions significantly underperform MTLCM. Again, we find that for all tasks, the learned causal indicator variables exactly match the ground- truth and the results from conditional independence lesting (Chen 2021 Lu et al. 2022
4.2 Rerld Data
We further evaluate our model on two real-world molecular datasots. We assume that the data x generated by transforming some unknown ground-truth latent factors z" with some unknown non-linear mixing func- tion. Since z² are unknown to us, identifiability in this setting is assessed by first training model 5 times with different random seeds, then computing the MCC score between the data representations recovered by each pair of those 5 models, as in Khemakhem et al. (2020b). As in Section 4.1 we employ the weak MCC score to assess the linear identifiability of the represent ations learned by the MTRN and the strong MCC score to assess the strictly strong identifiability of the latent factors z recovered by each method. Given that the truc latent dimension is unknown, we assess the identifiability of cach model gradually increasing latent dimensions. In practical scenarios, the latent dimension would be selected based on a similar model selection exercise. While the MCC threshold is likely be dependent the particular use case, the results
10

---

Published in Transactions on Machine Learning Research (08/2024)
0.95
0.90
0.85
I
MCC
0.80
0.75
0.70
Method
0.65
MTRN (weak)
0.60
MTLCM
5
10
Latent Dimension
Figure 3: Identifiability performance for the latent factors learned the QM9 dataset.
in Table g and Figure 3 suggest there is a relatively rapid shift from strong identifiability above 0.9 to much lower identifiability scores.
4.2.1 Superconductivity Dataset
The superconductivity dataset (Hamidich, 2018 consists of 21, 263 superconductors. We consider the tasks of regressing 80 readily computed target. features such as mean atomic mass, thermal conductivity and valence of the superconductors from their chemical formulae, represented as discrete counts of the atomss present in the molecule. The MTRN and the likelihoods in the baselines are parameterized by MLPs.
In Table 3. we find that the strong MCC for our MTLCM is greater than 0.96 and is able to match the weak MCC for the MTRN when the dimensions of the latent reprosentations are and 10, showing that our method manages to recover canonical latent representations for the superconductors. Interestingly, the strong MCC score for the MTLCM decreases suggesting that there are at most 10-20 independent tasks out of the 80 targets used for this data. In sharp contrast, all baseline models fail to recover identifiable latent factors for the superconductors in all cases as their strong MCC scores do not exceed 0.4. There are several settings where optimization diverged during training, since VAE-based models are generally difficult to train on discrete inputs chemical formulae.
increase the number of latent factors in the model,
4.2.2 QM9 Dataset
The QM9 dataset (Ruddigkeit al. 2012 Ramakrishnan al. 2014 is a popular benchmark for molecular prediction tasks consisting of 134, 000 enumerated organic molecules of up nine heavy atoms together with set of 12 calculated quantum chemical properties. In contrast the more artificial supercondnetivity dataset, the QM9 dataset enables us to assess the feasibility of achieving identifiable representations in the context of highly non-trivial quantum chemical propertics which are highly relevant to their pharmacological profile. Accurately modeling this dataset requires us to capture potential three-dimensional alomic interactions, allowing us to assess the translation of our theoretical results to more complex equivariant graph neural network architectures. For this reason, we use an equivariant graph neural network (EGNN) (Satorras et al. 2021) as the feature extractor for the MTRN. This enables the model incorporate positional features of each atom while exhibiting equivariance to their rotation, translation reflection. Given that the graph
11

---

Published in Transactions on Machine Learning Research (08/2024)
autoencoders proposed in Satorras al. (2021) and prior works (Kipf & Welling 2016 Simonovsky & Komodakis 2018 Liu et al. 2019) do not provide a means of jointly decoding the feature and adjacency matrices, we do not consider the iVAE and iCaRL baselines for this dataset.
In Figure 3 the weak identifiability achieved from the MTRN implics that identifiability achievable up to eight latent features, after which there a sharp decline in MCC. The implication of this observation is that there exist some redundancies between tasks (i.e., the number of effective tasks is less than the total number of tasks), which limit the maximal identifiable latent dimension. This is clearly the case for certain tasks. For example, prediction of the HOMO-LUMO gap can be directly obtained result of the difference between HOMO (highest occupied molecular orbital energy) and LUMO (lowest unoccupied molecular orbital energy) values. Nonetheless, the MTLCM is able to closely approximate the weak MCC score up to eight latent factors, always surpassing a score of 0.9, demonstrating its ability to recover permutation identifiable representations the context of realistic molecular datasets.
5 Conclusion
We have proposed a novel perspective on the problem of identifiable representations by exploring the impli- cations of explicitly modeling task structures. We have shown that this implics new identifiability results for linear equivalence classes in the general case of multi-task regression. Furthermore, while spurious corre- lations have been shown to be failure case of deep learning in many recent works, we have demonstrated that such latent, spurious signals may in fact be leveraged to improve the ability of a model to recover more robust disentangled reprosentations (i.c., point-wisc identifiability. In particular, when the latent space is explicitly represented as consisting of partitioning of causal and spurious features per lask, the linear identifiability result of the multi-task setting may be reduced identifiability up simple permutations and scaling under sufficient variability conditions. We have thoroughly discussed the assumptions underlying our proposed model and their implications. Empirically, we have confirmed that the theoretical results hold for both linear and non-linear synthetic data and for two real-world molecular datasets of superconductors and organic small molecules. We anticipate that this may reveal new research directions for the study of both causal representations and synergies with multi-task methods, and hope that these methods will enable robust generalization across lasks.
Acknowledgments
We thank Jon Paul Janet and Dino Oglic for helpful discussions. WC acknowledges funding via a Cam- bridge Trust Scholarship (supported by the Cambridge Trust) and a Cambridge University Engineering Department Studentship (under grant G105682 NMZR//05 supported by Huawei R&D UK). Julien ac- knowledges funding via a Cambridge Center for AI in Medicine Studentship (supported by AstraZeneca). JMHL acknowledges support from a Turing AT Fellowship under grant EP BP/VV23756/1.
Part of this work was performed using resources provided by the Cambridge Service for Data Driven Discovery (CSD3) operated by the University Cambridge Research Computing Service (www. csd3, . cam ac. uk), provided by Dell EMC and Intel using Tier-2 funding from the Engineering and Physical Sciences Research Council (capital grant EP EP/T022159/1) and DiRAC funding from the Science and Technology Facilities Council www.wwiracc.. www.ddirrccoomm uk.
References
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint XXXt:1907.028933 2019.
Yoshua. Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspec- tives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798 1828, 2013.
Jack Brady, Roland Zimmermann, Yash Sharma, Bernhard Schölkopf, Julius von Kügelgen, and Wieland Brendel. Provably learning object-centric representations. arXiv preprint ar Xix:22305. 14229, 2023.
12

---

Published in Transactions on Machine Learning Research (08/2024)
Wenlin Chen. Causal representation learning for latent space optimization. MPhil thesis, University of Cambridge, 2021.
Wenlin Chen, Austin Tripp, and José Miguel Hernández- Meta-learning adaptive deep kernel gaus- sian processes for molecular property prediction. In The Fleventh International Conference on Learning Representations, 2023.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Nlya Sutskever, and Picter Abbeel. Infogan: In- terpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
Pierre Comon. Independent component analysis, a new concept? Signal Processing, 36(3):287 314, Apr 1994. ISSN 0165-1684. doi: 10.1016/0167, (0.1016/0165-11655)))9000222-020 1684(94)9902299..
Cian Eastwood, Alexander Robey, Shashank Singh, Julius Von Kügelgen, Hamed Hassani, George J Pappas, and Bernhard Schölkopf. Probable domain generalization via quantile risk minimization. Advances in Newrul Information Processing Systems, 35:17340-17335 2022.
Marco Fumero, Florian Wenzel, Luca Zancato, Alessandro Achille, Emanuele Rodolå, Stefano Soatto, Bern- hard Schölkopf, and Francesco Locatello. Leveraging sparso and shared feature activations for disentangled representation learning. arXie preprinl urXiv:2344.0933) 2023.
John Gurland. On regularity conditions for maximum likelihood estimators. Scandinavian Actuarial Journal, 1954(1):71-706 1954.
Hermanni Hälvä, Sylvain Le Corff, Luc Lehéricy, Jonathan So, Yongjie Zhu, Elisabeth Gassiat, and Aapo Hyvarinen. Disentangling identifialle features from noisy data with structured nonlinear ica. Advances in Newral Information Processing Systems, 34:1624 1633, 2021.
Kam Hamidieh. Superconductivty Data. UCI Machine Learning Repository, 2018. doi: https://doi.orr//10 24432/C53P47.
J. Herault and C. Jutten. Space or time adaptive signal processing by neural network models. ATP Conference Proceedings, (11(1):206 211, Aug 1986. ISSN 0094-243X. doi:
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mo- hamed, and Alexander Lerchner. Beta-vae: Learning basic visual concepts with a constrained variational framework. International Conference on Learning Representations, 2017.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Aapo Hyvarinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ica using auxiliary variables and generalized contrastive learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 859 868. PMLR, 2019.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural Networks, 12(3):429-439, Apr 1999. ISSN 0893-6080 doi:
Leif Jacobson, James Stevenson, Farhad Ramezanghorbani, Steven Dajnowicz, and Karl Leswing, Leveraging multitask learning to improve the transferability of machine learned force fields. ChemRriv, 2023.
lless Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: unifying framework. In Proceedings of the Twenty Thind International Conference on Artificial Intelligence and Statistics, pp. 2207 2217. PMLR, Jun 2020a.
Ilyes Khemakhem, Ricardo Monti, Diederik Kingmaa, and Aapo Hyvarinen. Ice-beem: Identifiable conditional energy-based deep models based on nonlinear ica. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 22768-12778 Curran Associates, Inc., 2020b.
13

---

Published in Transactions on Machine Learning Research (08/2024)
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
Thomas N Kipf and Max Welling. Variational graph auto-encoders. arXiv preprint arXiv:1611 10100030,,,,,,, 2016.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for ro- bustness to spurious correlations. In The Eleventh International Conference on Learning Representations, 2023.
Bohdan Kivva, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Identifiability of deep gen- erative models without auxiliary information. Advances in Neural Information Processing Systems, 35: 15687 15701, 2022.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. duttof-distributton generalization via risk extrapolation (rex). International Conference on Machine Learning, pp. 5815 5826. PMLR, 2021.
Sébastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity rogularization: A new principle for nonlinear ica. In Conference on Causal Learning and Reasoning, pp. 428 484. PMLR, 2022.
Sébastien Lachapelle, Tristan Deleu, Divyat Mahajan, Ioannis Mitliagkas, Yoshua Bengio, Simon Lacoste- Julien, and Quentin Bertrand. Synergies between disentanglement and sparsity: Generalization and iden- tifiability in multi-task learning. In International Conference on Machine Learning, pp. 18171 18206. PMLR, 2023.
Jenny Liu, Aviral Kumar, Jimmy Ba, Jamie Kiros, and Kevin Swersky. Graph normalizing flows. Advances in Neural Information Processing Systems, 32, 2019.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachemn Challenging common assumptions in the unsupervised learning of disentangled represen- tations. In Proceedings of the 36th International Conference on Machine Learning, pp. 4114-4124 PMLR,
May 2019.
Chaochao Lu, Yuhuai Wu, Jose Miguel Hernández-Lobtto, and Bernhard Schölkopf. Invariant causal repre- sentation lcarning for out-of-distribution generalization. International Conference on Learning Represen- lations, pp. 32, 2022.
Hiroshi Morioka, Hermanni Hälvä, and Aapo Hyvarinen. Independent innovation analysis for nonlinear vector autoregressive process. In International Conference on Artificial Intelligence and Statistics, pp. 1549 1557. PMLR, 2021.
Ronan Perry, Julius Von Kügelgen, and Bernhard Schölkopf. Causal discovery in heterogeneous environments under the sparse mechanism shift hypothesis. Advances in Neural Information Processing Systems, 35: 10904-10917, 2022.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chem- istry structures and properties of 134 kilo molecules. Scientific Data, 2014.
Geoffrey Roeder, Luke Metz, and Durk Kingma. On lincar identifiability of learned representations. In Proceedings of the 38th International Conference Machine Learning, pp.
PMLR, Jul 2021.
Mateo Mojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research, 19(1):1309-138) 2018.
Lars Ruddigkeit, Rund Van Deursen, Lorenz C. Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of Chemical Information and Modeling, 52(11) ::6664 22864-28775 2875, Nov 2012. ISSN 1549-9596, 1549-960X. doi: 10.1021 2.1021/ci3044154
14

---

Published in Transactions on Machine Learning Research (08/2024)
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (u) equivariant graph neural networks. In International conference on machine learning, pp. 9323 9032.332222 PMTR, 2021.
Bernhard Schölkopf. Causality for machine learning. arXiv preprint 1rXiv:1911 10500, 2019.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the TEEE, 109(5): 612 634, 2021.
Martin Simonovsky and Nikos Komodakis. Graphvae: Towards generation of small graphs using variational autoencoders. In Artificial Neural Networks and Machine Learning ICANN 2018: 27th International Conference on Artificial Newral Nelworks, Rhodes, Greece, Oclober 4-7, 2018, Proceediiggs Part I 27, pp. 412-422. Springer, 2018.
Ankur Sinha, Pekka Malo, and Kalyanmoy Deb. A review bilevel optimization: From classical to evo- lutionary approaches and applications. IEEE Transactions on Evolutionary Computation, 22(2):276 295, 2017.
Megan Stanley, John F Bronskill, Krzysztof Maziarz, Hubert Misztcla, Jessica Lanini, Marwin Segler, Nadine Schneider, and Mare Brockschmidl. Fs-mol: A few-shot learning dataset of molecules. In Thirly-fii Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Society Series B: Statistical Methodology, (1)):6111-62) 1999.
Brian Trippe and Richard Turner. Overpruning in variational bayesian neural networks. orXiv preprint arXiv:1801.06230, 2018.
Julius von Kügelgen, Luigi Gresele, and Bernhard Schölkopf. Simpson's 's paradox in covid-19 case fatality rates: a mediation analysis of age-related causal effects. IEEE transactions on artificial intelligence, 2(1): 18 18-27, 2021.
Cuiyan Wang, Agata Chudzicka-Czuppe Michael L Tee, María Inmaculada López Núñez, Connor Tripp, Mohammad A Fardin, Hina A Habib, Bach X Tran, Katarzyna Adamus, Joseph Anlacan, et al. A chain mediation model on covid-19 symptoms and mental health outcomes in americans, asians and europeans. Scientific reports, 11(1):6481, 2021.
Matthew Willetts and Brooks Paige. I don't need u: Identifiable non-linear ica without side information. Xixiv preprinl Xiv:21 05238, 2021.
15

---

Published in Transactions on Machine Learning Research (08/2024)
A Proof of Theorem 3.2
Proof. By the assumption that the predictive likelihoods for the two sets of parameters O' and O are equal, we have
(y/x, pe(y/x,t), Vt,x, v, of (. (x), v, N ba (x) W', o7,) = -Nuylb (x)"
(18)
(19)
" "e, o,.),
Vi, N,
(20)
This implies that the means of the two Gaussian likelihoods on both sides are identical:
h (x) (x)TW( Ww/ ho(x)" vt,
(21)
By the assumption that Span(Im(ho)) = Rd, there exist d inputs X X2 such that the matrix H = [xi),x), [xa]] Radd invertible. By the assumption that there exist d lasks (4.)/11 such that the set of regression weights (wr are linearly independent, we construct an invertible ma- trix W [Wa.".. w.yl Raxd For has, we similarly define H' [ha (x1)) , (xa) e Rdxd and W w w e Raxd
Xy Now, we evaluate Equation 21 at the d inputs x Xx and d tasks 1 ta defined above, which gives us the following linear equation:
H 'W'
W' = HW.
(22)
Since both H and W are invertible by assumption and the weight matrices W and W do not depend the input x, the matrix W must be invertible.
Now, evaluating Equation
21 at the lasks 4,, ,d, we have
(W)Thar (x) = WT)s(x), Vx
(23)
he (x) (W')-W"H (x),, Vx
(24)
hx(()) (b)(x), Vx.
(25)
Note that we have shown that A = (W") TWT invertible. This completes the proof.
B Proof of Theorem 3.8
Proof. Let & = 2d and u = [v. We first rewrite the density of the conditional prior in the exponential family form:
Pr(z|u) = Z(u) exp (T(z)"n(u),
(26)
where (uu) = (2x)4/2|A,1" 0.5 exp (- aaf A,a;) the normalizing constant, T(z) [ sufficient statistics, and 17( = PA (h|z) using the noise distribution Pe. (c.)
g+ are the
A, 'u
E R* are the natural parameters. We also rewrite the likelihood
diigg L A:)
N(
(h/z)) Nh o = N(h Az(0, Pe. (h Az).
(27)
Let A, be the ground-truth transformation matrix such that A and T+(t) = (c; 78 the ground- truth task-specific variables associated with each task t. The proof starts off by using the fact that we have maximized the marginal likelihood [12 of A and T for all tasks. This means that the marginal likelihoods of the two models are identical:
(h/u( (...T((lu)) Vh, u.
(28)
16

---

Published in Transactions on Machine Learning Research (08/2024)
The goal is to show that the latent factors hh recovered by our model and the ground-truth latent factor z" h are identical up to permutations and scaling for all h.
Starting from the equality of the two marginal likelihoods
128 we have
PA.T(hlu) = PA. (h/u)
(29)
/ Pa. (2l2))
pa(hiz)pr((/uddd
(30)
/ Per (h (z)jr(2//jdz
/ Pe. (h - A.Z)PT, (z/u)dz
(31)
+ f Pe, Pe. (h det))
/ ((hh-)pr.ae (A,
ch
(32)
- / Ph. - a1 (h)dh
/ Ph. - TI
(33)
- (Pe, * * (Pe, *FA..T.,u)())
(34)
Flpe.1(W)F(P. F'A. uu((()
(35)
F(PA,T (w) F(ūA..T.,ul(w))
(36)
PA,T. (h) PA..T..u(h)
(37)
PT (A
h u) del A) -1 PT. (A,' del(A.)
(38)
T(A T(A-'h) h) (u) - log ((u) - log (A) = T h) T 77 (u) log Z. (u) - log de.(
(39)
where
Equation
32
follows by the definition h = Az, h == A.z.
Equation
33
follows by the definition 7A.T (h) == pT (A -
Equation
344
denotes the convolution operator,
F in Equation
35
denotes the Fourier transform operator,
Equation 36 follows since the characteristic function F(P of the Gaussian noise e, is nonzero
almost eherywhere.
Now evaluate Equation and subtract the first equation from the remaining k equations to obtain the following h: equations:
33 at uo u.. Un from our assumption to obtain k 1 such equations
Z(uo)
= ((A,"))"( - (uo)) + log Z.
T(A (m)())( - (7m))) + log Z(u)
(40)
Z.(w))
where 1 1, k. Putting those & equat ions in the matrix-vector form gives
L'T(A-"h) = L. +
(41)
where qi = log Z. (wn) ((u)) Z. ((u)) L is the invertible matrix defined in the assumption, and L, is similarly defined for the sccond model. Since L is invertible, we can left multiply Equation (41 by L TT to obtain
L-T
T(A 2h) MT(A, h)) +r,
(42)
where M = L-TLT and r = L-Tq We note that our assumption only says L is invertible and tells us nothing about L,. Therefore, we need to show that M is invertible. Let h == Azi 0,...,.. We evaluate Equation ) these & + points obtain & + such equations, and subtract the first equation from the remaining k equations. This gives us
(Z))) T(zo). ... T(za) T(zo) = M(T(A,"), T(A,'ho), ... T(A, ha)) T(A-'ho) (43)
17

---

Published in Transactions on Machine Learning Research (08/2024)
We denote Equation 143 by R MR.. We need to show that for any given 2o, there exist k points 21:-. Zt such that the columns of R are linearly independent. Suppose, for contradiction, that the columns of R would never be linearly independent for any Z1,-. 222- Then the function g(z) == T(z) - T(zo) would live in a - or lower dimensional subspace, and therefore we would be able to find a non-zero vector À E IR* orthogonal to that subspace. This would imply that. (T(z) T² = and thus T(z) = T(za)T>
const, Yz, which contradicts the fact that our conditionally factorized multivariate Gaussian prior PT (z) u) is strongly exponential (see Khemakhem et al. (2020a) for the definition). This shows that there exist & points 21. Ze such that the columns of R are linearly independent for any given Zo- Therefore, R is invertible. Since R = MR, and M is not a function of 2, this tells us that M must be invertible.
Now that we have shown that M is invertible, the next step is to show that M is a block transformation matrix. We define a linear function l(z) = A, Az Now, Equation (2 becomes
T(z) = MT(())) +r.
(44)
We first show that the linear function I is a point-wise function. We differentiate both sides of the above equation w.rtt 2g and 2t (Vs # t) to obtain:
OT(z)
d OT(()))
01(2)
M
(45)
Oz.
(z)
Oz.-
1=1
TT)
8 TT((()) (1) (z) (z) (z) d2, Oz.
TT((2)) (1.(z)
M
M
(46)
02,02t
(z) 0z,02t
1=1j=1
Since the prior (z|u)) is conditionally factorized, the second-order cross derivatives of the sufficient statistics are zeros. Therefore, the second equation above can be simplified as follows:
TT)(z)
(47)
Oz,02e
M PT(((2) 01((z) 01,(z)
TT((Z)) Fl((z)
+ M
(48)
(2) Oz 0z,
(l) 0z,0z1
MT (z)1k (z) + (z)
(49)
MTT (z) (z),
(50)
where
PT(((z))
OTT(((Z))
(z)
<Rd
(51)
01 (
0lg(z)²
(z) O4 (z)
(la(z)
Old(z) Oz,
1. (x)=
E ERd
(52)
O2 0z,
Oz,
OT((())
(T(())))
(z)
E gkxxd
(53)
Øl (z)
0l,(z) T
(11
(La(zz)
1" (z)
E ERd
(54)
02,021"
02,02,
(z) (z), T'(z)] E Raxk
(55)
!" (x) = (2) (2), (2) E R
(56)
By Lemma in Khemakhem et al. (2220a) and the fact that & = 2d, we have that the rank of T (z) is 2d and thus it is invertible for all z. Since M is also invertible, we have that MT" (z) invertible. Since MT" (z) o. must be that 1" (z) 0. Vz. In particular, this means that [z_(()) 0, vs # for all a, which shows that the linear function 1(x) A, Az is point-wise linear function.
18

---

Published in Transactions on Machine Learning Research (08/2024)
Now, are ready to show that M is block transformation matrix. Without loss of generality, assumme that the permutation in the point-wise linear function 1 is the identity. That is, 1(z) = 11, ... [aza] for some linear univariate scalars 4, , E R. Since A and A. are invertible, we have that 1 (z) =
PT 21. .1 2d] Define
= ( +
(57)
and plug it into Equation
44
gives:
T(z) (T)(
(58)
We then apply 1- to z at both sides of the Equation 58 to obtain
T(I (z).
(59)
Since 1 point-wise function, for a given q (1, .k we have that
0 OT(I (z))g
(T()), for any such that # and q # 2s.
Mgd
(60)
Oz.
Ozg
J=1
Since the entries in T(z) are linearly independent, it must be that Mg. -0 for any j such that 0T()), 8s, + O. This includes the entries j in T(z) which depend on 2s (i.e., J = 8 and j = 2s). Note that this holds true for any such that A and q o 2s. Therefore, when is the index of an entry in the sufficient statistics T that corresponds to 2, (i.e., = or q = 2i, and i * s), the only possible non-zero Mqg for j are the ones that map between (24 and (2+)) where T, are the factors in T that depend on 2 and T, are similarly defined. This shows that M block transformation matrix for each block [zi, a7] with scaling factor 4. That is, the only possible nonzero element in M are Mi,i M120 M21, and M2222 for all i (1,..,d))
Furthermore, for any i E f1, dd} we have that
1- OT(I-
K Σ
oT(z);
Mi.j
Mi,5 + 2M,
(61)
02
dz,
1=1
(
0T(z);
21, 2,
Σ M2i.j
M2rs + 222
(62)
Øz;
Ozi
This implies that M22 = 0 and M21, = 0 for any i E [1.....dd d), and M,, 47 for i E {1, k}, which reduces M from a block transformation matrix to a permutation and scaling matrix. In particular, this mcans that the latent factors 2, are identifiable up to permutations and scaling, with the transformation matrix E Iaxd defined by the first rows and d columns of M:
A A h PA hr
h = AP +
(63)
Since h is lincarly identifiable by assumption, it must be that Ar by Definition 3.1 Since is invertible by assumption, it must be that = 0. Therefore, we have
A h h
(64)
This completes the proof.
C Derivation of the Factorized Structured Conditional Prior
Since no prior knowledge is assumed for the task-specific regression weights w, E Rd, we put an uninformative prior over W. d for all tasks t:
p(w.) oc 1.
(65)
19

---

Published in Transactions on Machine Learning Research (08/2024)
This uninformative prior can be thought of as Gaussian prior with infinite variance:
p(w)) = lim (-()),, T
(66)
where (w()) = (( We marginalize out we from py (ylze,t) = Noll(w, o 22)) 02) under our uninformative prior over W5, which makes the marginal uninformative:
(y / pr (rllue ()p(ws)dw;
(67)
/ pT 1) lim or (we)dw, 70000
(68)
/ prilize
lim T+00
(69)
lim (z)l0, og
(70)
~-000
(71)
This is known improper uniform distribution since it does not necessarily integrate to one. However, it is worth noting that the posterior pr () is still well-defined even p'7 (ullee () is inproper. To see this, we denote the improper uniform distribution by pT (ulz,, t) = C. for some constant C. Then, we have
PT (zel p) (ylz Zes O)T ()
(z(y
(72)
J pT (zel 1) f) (1) Ze, t)p (2, y, u,tt) f)d
PT Zel PT (2al21 (y,))
(73)
J PT () 1) 22 w. dz
PT (2((2))) (2.ly.))
(74)
Since (zell) factorizes over the causal latent factors and py pz () factorizes over the spurious latent factors, the structured conditional prior (rly), factorizes all latent factors
Furthermore, we verify that the compact expressions for the mean and variance of (zly) t) in Equation 10p are correct. Recall that Equation 5 tells us that
pT (zelt) NZ 0,1).
(75)
and Equation 7 tells that
pT( t) = Nizzlum. 1I..
(76)
Recall that the compact expressions given by Equation 110 are
a, = yn: o (1 (1- c.), A, = diag c1) c.).
(77)
For any causal latent variable 2,, we have Ct; = and therefore at,i and At,, = For any spurious latent variable Zj, have Ctjj and therefore arjj = Way and Ar,j = T.. This verifies that Equation 110 is correct.
D Derivation of the Marginal Likelihood for MTLCM
The marginal likelihood for MTRN given by Equation 12 is
/ 1::00<pAAhh
Pe (bly,t)
(Nhuu, E.),
(78)
where PA (h/z) (h/AAz o and (z 1) Nizla, Ar). Equivalently, we can rewrite the likelihood the following form:
h Az +
(79)
20

---

Published in Transactions on Machine Learning Research (08/2024)
where p(e) = o3)) Since both PA (h/z) and pr (zly, t) are linear Gaussians, derive closed-form expression for the mean μr and covariance Σj using moment matching:
E Bpr AB p = Aa, ( (1-c)),
(80)
Σg
(x = A Aa [2] AMAA Acdia((o[11c)) (81)
This verifies that Equation 12 is correct.
E Possible DAGs for the Generating Process the Target Variable
Possible structures for the latent factors generating the target y are given in Figure
24
2j
Zj
2k
2k
Zk
(a) Standard case
(b) Uncorrelated case
(c) Confounder case
Figure 4: Illustration causal relationships which are captured by our model (a, b) and not captured by our model (c) for the relationships between latent variables and observed target y for given task. The red arrow in (c) indicates the portion of the graph which is not captured by MTLCM. Note than in (b), the existence of learned regression weights encapsulates this case if the learned weight is zero the arrow 2, + This is depicted with the dashed green arrow.
F Motivating Examples for the Assumed Non-Causan Relationship
We acknowledge that indeed our assumed direct edge y Z, in Figure I does not in general capture all possible non-causal correlations between latent features and the target since the Reichenbach principle states that non-causal correlations can originate from either (1) common cause (i.e., confounders) or (2) an anti-causal/spurious relationship (as assumed in this work). However, we argue that there are many situations where the proposed model can be useful in practice, even does not explicitly model the confounders in full generality, since this anti-causal relationship (as our paper) well-documented real- world examples in epidemiology and drug discovery. We provide couple of real-world motivating examples to justify this anti-causal assumption below.
Epidemiology. See Figure
Wang al. (2021) treating perceived pandemic impact or IES-R score the regression target, or Figure (right) in von Kügelgen al. (2021), where testing status feature in estimating Case Fatality Rates, but there is likely to be causal
may well be included influence between overall case fatality rate and testing policy. Broadly, any form of selection bias may lead to similar cases, where the selection criterion may be included in the feature set.
Drug discovery. In most drug discovery campaigns, molecules be tested are selected based some structural similarities to originally promising molecule (based on the quantity to be estimated, e.g. drug potency). Structural molecule foatures are then likely to be spuriously correlated with the regression target due to their selection criteria, without actually being involved in the drug's mechanism of action.
21

---

Published in Transactions on Machine Learning Research (08/2024)
Algorithm Pseudocode for the data generating process in the synthetic data experiments
Require: the number of latent features d, the number of causal features Ne, the number of tasks Nr, the number of points per task N,, Ground-trthh transformation F (random invertible matrix random MLP) Set 0, 0.1 and o, 0.01 for each task do
Sample d binary causal feature indicators 11, 12 1.. , Sample weights wf. wh.... w/ (4(01)) 1)
Sample spurious coefficients 73 1) for all j such that 13 =
for each data point x in this task do
Sample causal features Zij ~N(00, ) for all such that Iÿ. =
Sample of ~ (2)
Obtain target = م 3.j + c o ~ NO, ((,)3)
Obtain spurious features žij = د + 6,iys 6,,i. ~ N(,,o?) for all j such that 1= = 0
Obtain observed features via the transformation x F(z)) Eo. cor NOO, I
end for end for
G Model Configurations
In Stage 1, the learnable parameters of multi-task regression network (MTRN) are the feature extractor parameters & and the task-specific regression weights w, for all tasks t. These model parameters are learned by maximum likelihood as defined in Equation (3)
In Stage 2, the learnable parameters of a multi-lask linear causal model (MTLCM) are the linear transfor- mation A, the causal indicators c, for all tasks and the spurious coefficients 7/ for all tasks These are free parameters learned by maximum marginal likelihood defined in Equation (14h. The binary causal indicators C2 are parameterized gradient update of Cą2 we do not binarize the output of the sigmoid function during training; instead, we use a soft version è₁ E [0, 1jd during training. In practice, we find that this works well and all learned values for G.1 are very close to either 0 or 1. In the synthetic data setting, the learned causal indicators match the ground-truth values. In practice, we fix the spurious noise variance 0, and the observational noise variance do to some constants using cross validation.
free parameters squashed to [0, 1] by the sigmoid function. To allow for
For fair comparison, we also consider the multi-task extensions of iVAE and iCaRL, MT-IVAE and MT- iCaRL, which include the task variable in the conditioning variables u in their conditional priors (z/u), with the task-specific parameter T(t) {v,} to be learned from data, which is the counterpart to T(t) {cr, (c) in our MTLCM but has no explicit interpretations with respect to a causal graph. We set dim(v.) = dim(c;) + dim( to ensure the same degree of flexibility as our MTLCM. The task-specific parameters V, are free parameters learned together with other parameters in these models by optimizing their variational/score matching objective.
H Experiment Settings for the Synthetic Data
This section details the precise process for the data generation of the synthetic data for both the lincar and non-linear experiments in Section 4.1. Algorithm I details the full data generation process, Table 4 details the experiment hyperparameters used in the linear setting and Table 5 details the hyperparameters used in the non-linear setting. The transformation in the linear experiments corresponds to cither the identity, an orthogonal or a random matrix of size d, while in the non-linear experiments it corresponds a randomly initialized neural network with the specified hidden dimensions and relu activations.
22

---

Published in Transactions on Machine Learning Research (08/2024)
Ablation Study for the Linear Synthetic Data
In Figure 5 we contrast the effect of training only the linear transformation matrix A in our MTLCM when the ground-truth task variables C₁2 7 are known to the model, with the more general setting of learning all parameters jointly via maximum marginal likelihood. We assess the convergence of our multi- task linear causal model across random seeds for increasingly complex linear transformations (identity, orthogonal, random) for data consisting of 10 latent factors with two causal features. Rather than inhibiting convergence, we find that training all parameters jointly leads to improved performance, possibly due to additional flexibility in the parameterizations of the model. For all types of linear transformations, model succeeds in recovering the ground-truth latent factors. In additional, we find that standardizing the features accelerates convergence.
23

---

Published in Transactions on Machine Learning Research (08/2024)
Table 4: Experimental Settings for the Linear Synthetic Data
Latent Dim
3, 5, 10, 20, 50
Observation Dim
Latent Dim
Causal Seed
2, 4
1, 3, 4, random
Matrix Type
Table Experimental Scttings for the Non-Lincar Synthetic Data
Observation Dim
50, 100, 200
Encoder Network Num Hidden Layer
Encoder Network Hidden Dim
2 Observation dim
Latent Dim
20
Num Causal
4, 8, 12
Seed
1, 2, 3, 4,
Matrix type Identity
Matrix type = otthogonal
Matrix type random
0.50
MCC
0.85
0.80
Ground truth scaled No ground truth scaled Ground truth not acaled
ground truth
1000
2000
3000
4000
5000
6000
1000
2000
3000
4000
5000
8000
1000
2000
3000 Epoch
4000
5000
6000
Epoch
Epoch
Figure 5: Convergence of the model in the case of transformations of the latent factors for identity, orthogonal and arbitrary lincar transformations. Scaled means standardizing the features.
24


Reviews:
[{"review_id": "nsJhBsveQM", "reviewer": "Reviewer_jUzE", "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.\tOne more interesting identifiability result (up to permutation and scaling) for multi-task setup is proposed.\n\n2.\tExperiments results are sufficient to show the variability of the proposed theorem and method. Linear and nonlinear synthetic experiments clearly validate the theoretical results in two stages, respectively. Also, the experiments on real datasets validate the theoretical results.\n\n3.\tThe overall explanation is clear and easy to follow.\n\n# Weakness:\n\n1.\tSome assumptions appear to be restrictive. The linear relationship between $z$ and $y$ may not hold in real-world scenarios. Additionally, assuming Gaussianity in the latent space limits the applicability of the theoretical results.\n\n2.\tThe model structure presented in the problem formulation section lacks the intuitiveness and is not accompanied by a motivating example. The directed edge from $y$ to $z_s$  is particularly difficult to justify, as its purpose and rationale are unclear.\n\n3.\tStill in the problem formulation section, where ground truth latent factors are assumed to be independent there. This seems to contradict the model structure where $z_c$ and $z_s$ are dependent via the path passing $y$.\n\n4.\tFor the synthetic dataset, ground truth causal latent variables $z_c$ are generated from a standard Gaussian, which implies there are no causal relations inside $z_c$. This setup is not typical in causal representation learning, where allowing causal relations or statistical dependencies within $z_c$  is crucial. If the theorem does not apply to cases with such dependencies, its relevance to causal representation learning is unclear. Conversely, if the theorem accommodates dependencies, it would be beneficial to demonstrate this in the synthetic dataset.", "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 work focuses on the identifiability results of multi-task learning. It proposes two theoretical results: the first result demonstrates identifiability up to a linear transformation, while the second extends this to identifiability up to permutation and scaling. Both results are based on a novel model setup with Gaussian assumption. Experimental validation on both synthetic and real-world datasets confirms the effectiveness of the proposed theorems.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "My main concerns are the final three points in the weakness part, especially the loss of explanation for the existence of the directed edge from $y$ to $z_s$. If one or more motivating examples are provided, the whole setting would be more convincing.", "broader_impact_concerns": "I do not see any potential negative ethical impact of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oUCIhWj6KT", "reviewer": "Reviewer_YUaf", "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:\nThis work is well-written, provides a thorough examination of the prior literature, and introduces interesting results about their proposed model. These are supported by a good variety of experiments.\n\nWeaknesses:\nI do have some suggested areas for improvement. This work could be made stronger if both (1) the linearity assumptions and (2) causal structure (depicted in Figure 1) in Section 3.1 were better justified. The decomposition into causal and spurious latent factors is interesting but there are other possibilities here, for instance $t \\to \\{z_c, z_s\\} \\to x \\to y$ (as an example). This is discussed somewhat in Section 3.4 but Section 3.4 could be improved by using the mathematical notation you have conveniently introduced (for instance, perhaps introduce alternative DAGs in the Appendix for the reader to refer to). More generally, it would help the reader if these assumptions were discussed earlier before 3.4. I was left wondering until 3.4 what the justification was for the modeling assumptions (linearity and DAG).\n\nSome of the theorems could benefit from more precise statements. E.g., Corollary 3.3 vaguely refers to \"standard regularity conditions for consistency of MLE estimators.\" Definition 3.7 Strictly Strong Identifiability is introduced but isn't actually referred to in Theorem 3.8 (instead, you say 'up to permutations and scaling' but you might as well use your Definition 3.7 there).\n\nI quickly looked at Appendix D, and introducing an infinite variance Gaussian prior $p(w) \\propto 1$ seems suspicious. Is this valid? It seems as though you need to be more careful here, and prove a statement as the limit of this variance goes to infinity, instead of simply assumption $\\propto 1$.\n\nAt some points it seems as though $h$ and $z$ both refer to \"latent factors\". You should stick with consistent notation for $h$ vs $z$, or make it clear which \"latent factors\" you are actually referring to. This term seems to get overloaded frequently", "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 work explores multi-task regression through the lens of a latent variable model. Specifically, it introduces a causal structure relating tasks to regressors and response values. First, this model is used to show linear identifiability of data ($x$) representations (Theorem 3.2). This is used to show that MLE recovers (up to a linear transformation) the true data representation. Then, it is shown that the true latent factors $z$ are recoverable up to scaling and permutation (Theorem 3.8). It is shown for both synthetic and real-world experiments (superconductivity, molecular prediction) that this approach shows stronger correlation (identifiability measure) than baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I request that all of the weaknesses above be addressed, but these are not critical to securing my recommendation for acceptance. They would simply strengthen the work.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dfAj13ksis", "reviewer": "Reviewer_A9UQ", "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\n•\tThe paper is well organized and clearly written, and the topic is of potential interest to the general audience of TMLR.\n\n•\tThe model structure and the form of the likelihood function make sense along with the proposed data generation process and assumptions on the task structure, and the resulting conclusion on identifiability is theoretically sound.\n\n•\tThe authors discuss the assumptions on the data generation process and model parameters, including potential limitations and future research directions in Section 3.4.\n\n•\tThe proposed model consistently outperforms iVAE and iCaRL benchmarks on both synthetic and real-world datasets in terms of MCC between learned and ground truth latent factors (or learned latent factors across different initializations for real-world dataset).\n\nWeaknesses:\n\n•\tThe linear identifiability claim in Theorem 3.2 assumes there exist $d$ tasks such that at least one set of regression weights are linearly independent so that the resulting regression weight matrix $\\textbf{W} = [\\textbf{w}\\_{t_1}, …, \\textbf{w}\\_{t_d}]$ is invertible. However, this assumption appears to lack adequate substantiation or explanation within the context of real-world experiments, particularly those conducted on the superconductivity and QM9 datasets. Specifically, the relationship between task independence and the correlation between target features and the dimension of the latent factor $\\textbf{z}$ is not clear. In practice, how should we select the dimension of the latent factor to ensure identifiability?\n\n•\tIn Remark 3.6, the authors mention that methods such as iVAE and iCaRL depend on target variable $y$ to extract identifiable latent factors. However, based on my understanding, iVAE in general is not specifically designed for causal representation learning and the conditional prior on $\\textbf{z}$ depends on an auxiliary variable $\\textbf{u}$ which is not necessarily the target variable. It would be beneficial to further clarify this point.", "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 introduces a method for weakly identifying the ground truth data representations by leveraging task structures in the multi-task regression setting. The authors first define a causal graph for data generation where the ground truth latent representations can be separated into causal and spurious factors for each task. They then propose a framework to recover latent representations based on data and task distributions, which consists of a multi-task regression network (MTRN) and a multi-task linear causal model (MTLCM). Both models can be learned via maximum likelihood. The authors theoretically demonstrate that, given the data generation process and the specified likelihood functions, the latent representations learned by their method are identifiable up to scaling and permutation. Furthermore, they validate their model's effectiveness through experiments on both synthetic and real-world datasets, using the mean correlation coefficient (MCC) to quantify the accuracy of the recovered representations in comparison to the true latent data.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "•\tPlease kindly refer to the bullet points in the weaknesses section above and provide corresponding clarifications in Remark 3.6 and Section 4.2.\n\n•\tOn the last line of page 3, should the target generation process be $y = (\\textbf{w}_t^*)^T \\textbf{z}_c^*$ instead of $y = (\\textbf{w}_t^*)^T \\textbf{z}^*$?\n\n•\tTo make the illustration clearer, it might be desirable to show the detailed architecture of MTRN in Figure 2, which consists of a shared feature extractor network $\\textbf{h}_{\\boldsymbol{\\phi}}(\\cdot)$ and $N_t$ linear heads for each task. \n\n•\tIt might be better to show how the forms of mean and covariance term in Equations 10 and 13 are derived in the appendix to make the methodology section more self-contained.", "broader_impact_concerns": "I do not see immediate ethical concerns of the work presented, and I think the authors have appropriately addressed potential limitations on the generality and applicability of the method.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a latent variable to model the relations between tasks. Theoretical results on the identifiability of the latent variable are provided. Also, the experimental studies are sufficient to show the correctness of the proposed theory. All reviewers are positive about this paper, though there are some concerns on the linear relations in part of the model. I would recommend acceptance of this paper due to novel technical contributions, solid theory, and a variety of empirical studies.
Decision: Accept (poster)