Paper: Leveraging Task Structures for Improved Identifiability in Neural Network Representations (ICLR-2024)

Reviewer comment (Theory):
In the assumed causal graph, $Y$ is a parent of the spurious features $z_{s}$, which is not a natural assumption. The correlation between $Y$ and $z_{s}$ is often due to latent confounders.

Author response:
Thank you for taking the time to review our work. Regarding the primary concern of the target relationship with the spurious variables, we can first begin by acknowledging that indeed this assumption does not in general capture all possible non-causal correlations between latent features and $Y$. The Reichenbach principle states that such correlations (when non-causal) must either originate from a common cause (confounders) or from an anti-causal relationship, as assumed in this work. 

When dealing with models over unobserved causal variables which are non-linear transformations of the input space, it can be difficult to provide intuition for real-world settings. However, this anti-causal relationship (as in the paper) is well-documented in real-world examples in epidemiology. See Figure 1 in [1], treating perceived pandemic impact or IES-R score as the regression target, or Figure 6 (right) in [2], where testing status may well be included as a feature in estimating Case Fatality Rates, but there is likely to be causal 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. A similar example can be imagined in most drug discovery campaigns, where molecules to be tested are selected based on some structural similarities to an originally promising molecule (based on the quantity to be estimated, e.g. drug potency). Structural molecule features 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.
Thus the first portion of our answer is that there are many situations where the proposed model may be useful in practice, even if it does not explicitly model the confounders in full generality.

The second portion of our response is that any confounding variable on the spurious variable $z_s$ would also be a causal variable for $Y$, and thus already be included in the set of latents being modelled. Since this work is concerned with identifying the correct latent variables as opposed to the full connectivity of the underlying causal graph (which is a more general but much more challenging problem), the existence of an additional causal relation from a causal latent variable to a “spurious” one is thus possible but does not factor into our model.

Reviewer comment (Theory):
What does ``MLE (1) converges" in Corollary 3.3 mean? $\theta'$ defined in (1) is a population estimator. There is a similar issue for (17) in Theorem 3.5. Besides, the uniqueness of the (population) MLE estimators (i.e., (1) and (17)) should be assumed explicitly.

Author response:
For point 2, we have added the explicit uniqueness assumption for the MLE estimator and reworded “convergence” to more precise terminology. Convergence here refers to both convergence of the optimization procedure together with standard regularity conditions for consistency of MLE estimators.

Reviewer comment (Writing):
What is $c_{t}$ below equation (4)?

Author response:
For point 3, $c_t$ refers to the indices of the causal variables for task t. This is explained in the paragraph preceding equation (4): 

> Let $\mathcal{T}(t)=\{\mathbf{c}_t,\boldsymbol{\gamma}_t\}$ be a collection of task-specific variables associated with task $t$, which are free parameters to be learned from data, where $\mathbf{c}_t\in\{0,1\}^d$ are the causal indicator variables which determine the partition of $\boldsymbol{z}=\boldsymbol{z}_c \cup \mathbf{z}_s$ for the given task $t$ (i.e $c_t^i=1$ indicates that $z_i$ is a causal latent factor in task $t$ and $c_t^i=0$ indicates that $z_i$ is a spurious latent factor in task $t$).

Reviewer comment (Theory):
The main concern is the causal relation between $Y$ and the spurious features $z_{s}$. If there are indeed many real-world settings where $Y$ is the parent for the spurious features, please cite related references to support this assumption. Otherwise, the identifiability for settings with the latent confounder should be studied as well.

Author response:
Thank you for taking the time to review our work. Regarding the primary concern of the target relationship with the spurious variables, we can first begin by acknowledging that indeed this assumption does not in general capture all possible non-causal correlations between latent features and $Y$. The Reichenbach principle states that such correlations (when non-causal) must either originate from a common cause (confounders) or from an anti-causal relationship, as assumed in this work. 

When dealing with models over unobserved causal variables which are non-linear transformations of the input space, it can be difficult to provide intuition for real-world settings. However, this anti-causal relationship (as in the paper) is well-documented in real-world examples in epidemiology. See Figure 1 in [1], treating perceived pandemic impact or IES-R score as the regression target, or Figure 6 (right) in [2], where testing status may well be included as a feature in estimating Case Fatality Rates, but there is likely to be causal 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. A similar example can be imagined in most drug discovery campaigns, where molecules to be tested are selected based on some structural similarities to an originally promising molecule (based on the quantity to be estimated, e.g. drug potency). Structural molecule features 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.
Thus the first portion of our answer is that there are many situations where the proposed model may be useful in practice, even if it does not explicitly model the confounders in full generality.

The second portion of our response is that any confounding variable on the spurious variable $z_s$ would also be a causal variable for $Y$, and thus already be included in the set of latents being modelled. Since this work is concerned with identifying the correct latent variables as opposed to the full connectivity of the underlying causal graph (which is a more general but much more challenging problem), the existence of an additional causal relation from a causal latent variable to a “spurious” one is thus possible but does not factor into our model.

Reviewer comment (Writing):
My main complaint is the clarity of the paper.

Author response:
Given that the reviewer’s only concern appears to be the clarity of the work, which has otherwise been well received, we would kindly urge them to reconsider their rating as we do not believe this warrants a rejection score. If there are still specific elements which are unclear that the reviewer would like to point out, we are also happy to clarify them in the manuscript.

Reviewer comment (Writing):
I am not sure I properly understood the setup or the contribution: no algorithm is clearly proposed.

Author response:
The presentation of 3.3.1 was chosen to make clear, with textual explanation, where each equation comes from. Together they define our model. Since most reviewers appreciated the presentation of our paper, we are hesitant to change this. Please see below our condensed explanation of the method and let us know if this is still unclear. We are happy to elaborate further. 

The overall method is a two stage algorithm. This algorithm applies to the setting of multi-task regression, i.e. when multiple continuous outputs are associated to a joint set of input points. This is often relevant for example in drug discovery, where one would like to jointly model various relevant properties of a drug candidate. The first stage of the algorithm is simply to train a standard multitask regression neural network (which has separate linear heads $w_t$ for each task and a shared feature extractor $f(x)$ across all tasks) jointly on all tasks. Theorem 3.2 says that the resulting last-layer representation of the neural network will be linearly identifiable. The second stage of our algorithm starts from this linearly identifiable representation and “disentangles” it via our probabilistic model. Here, we recover the linear transformation $A$ (see equation 10) by maximizing the likelihood obtained in equations 11 and 12. Note that this general procedure is summarized in the manuscript in the last paragraph of section 3.1, and illustrated in Figure 2. We are again hesitant to change the presentation here given the generally positive feedback in this regard, but are happy to clarify further.

Reviewer comment (Writing):
Abstract "In such cases, we show that identifiability is achievable even in the case of regression, extending prior work restricted to linear identifiability in the single-task classification case." I am not sure to understand this sentence since the authors refer to previous identification results in the multitasks setting (section 3.2 for instance).

Author response:
This sentence in the abstract refers to the extension of Roeder et al’s work on supervised learning identifiability. Given that our linear identifiability result most closely resembles this work, and that we view Lachapelle et al’s work as concurrent to ours, this phrasing was used. Nonetheless, we have reworded this in the abstract to avoid any confusion. The sentence now reads: “In such cases, we show that linear identifiability is achievable in the general multi-task regression setting. Furthermore, …”

Reviewer comment (Presentation):
Section 3.3.1 is a succession of 6 equations, is it possible to encapsulate the assumptions in a proper environment.

Author response:
The presentation of 3.3.1 was chosen to make clear, with textual explanation, where each equation comes from. Together they define our model. Since most reviewers appreciated the presentation of our paper, we are hesitant to change this. Please see below our condensed explanation of the method and let us know if this is still unclear. We are happy to elaborate further. 

The overall method is a two stage algorithm. This algorithm applies to the setting of multi-task regression, i.e. when multiple continuous outputs are associated to a joint set of input points. This is often relevant for example in drug discovery, where one would like to jointly model various relevant properties of a drug candidate. The first stage of the algorithm is simply to train a standard multitask regression neural network (which has separate linear heads $w_t$ for each task and a shared feature extractor $f(x)$ across all tasks) jointly on all tasks. Theorem 3.2 says that the resulting last-layer representation of the neural network will be linearly identifiable. The second stage of our algorithm starts from this linearly identifiable representation and “disentangles” it via our probabilistic model. Here, we recover the linear transformation $A$ (see equation 10) by maximizing the likelihood obtained in equations 11 and 12. Note that this general procedure is summarized in the manuscript in the last paragraph of section 3.1, and illustrated in Figure 2. We are again hesitant to change the presentation here given the generally positive feedback in this regard, but are happy to clarify further.

Reviewer comment (Writing):
What's the proposed algorithm? I was not able to parse it.

Author response:
The presentation of 3.3.1 was chosen to make clear, with textual explanation, where each equation comes from. Together they define our model. Since most reviewers appreciated the presentation of our paper, we are hesitant to change this. Please see below our condensed explanation of the method and let us know if this is still unclear. We are happy to elaborate further. 

The overall method is a two stage algorithm. This algorithm applies to the setting of multi-task regression, i.e. when multiple continuous outputs are associated to a joint set of input points. This is often relevant for example in drug discovery, where one would like to jointly model various relevant properties of a drug candidate. The first stage of the algorithm is simply to train a standard multitask regression neural network (which has separate linear heads $w_t$ for each task and a shared feature extractor $f(x)$ across all tasks) jointly on all tasks. Theorem 3.2 says that the resulting last-layer representation of the neural network will be linearly identifiable. The second stage of our algorithm starts from this linearly identifiable representation and “disentangles” it via our probabilistic model. Here, we recover the linear transformation $A$ (see equation 10) by maximizing the likelihood obtained in equations 11 and 12. Note that this general procedure is summarized in the manuscript in the last paragraph of section 3.1, and illustrated in Figure 2. We are again hesitant to change the presentation here given the generally positive feedback in this regard, but are happy to clarify further.

Reviewer comment (Theory):
My main concern is the Gaussianity assumptions made throughout the work. Such assumptions rarely hold in practice; and in contrast to standard machine learning tasks where most methods remain useful even if such assumptions are violated (e.g. typical regression procedures continue to estimate the conditional expectation), for causal representation learning the utility of the proposed method is far less clear.

Author response:
We thank the reviewer for taking the time to review our work. First, we appreciate the recognition that this work addresses an important problem, namely that of defining models which have more robust guarantees on the types of representations they learn. 

We would like to emphasize that the probabilistic assumptions formulated in our model not only make our model more computationally efficient, they also enable direct optimization of the underlying likelihood. This is in contrast to iVAE (which optimizes a variational objective) and iCaRL (which optimizes the score function), and the work of Lachapelle et al (2023), which proposes a bi-level stochastic optimization procedure. This is a key property which leads to the stability of our empirical results. We also believe that our formulation most plainly establishes the connections between identifiability and multi-task learning. We thus believe that the reviewer is undervaluing the contributions of our work.

The reviewer’s main concern lies in our Gaussianity assumptions. We feel that points 1 and 4 may be related in this regard:

- The feature space linearity (point 4) is accurate precisely because of the first stage of our procedure: namely this is not an assumption but a guarantee which follows from linear identifiability of the representations learned from the multitask regression network (theorem 3.2). More generally, any non-linear neural network broadly makes this assumption (where the conditional expectation is parametrized by a linear combination of the last layer activations), which is in general less well motivated than in this work given the lack of generic identifiability results. We would be happy to clarify this further if we are misunderstanding any of the reviewer’s concerns here.

- Because of the feature space linearity, the second stage procedure reduces to a linear problem. A gaussian likelihood is the widely accepted standard for linear regression models, and is thus a very natural choice. Given that the linearity only arises in the context of arbitrarily non-linear feature mappings $f(x)$ provided by the multitask neural network in the first stage, we do not believe this to be particularly restrictive. We would point to our empirical evaluation as evidence that our model can indeed be useful. The connection to our causal representation learning result, recovering the linear transformation to the latents, is simply a very appealing byproduct of maximizing the likelihood in equation (11); it does not inherently require different assumptions on the regression setting.

Reviewer comment (Experiments):
Regarding experiments, I think it would be more convincing if the authors could include comparisons to the recent works of Lachapelle et al (2023) and Fumero et al (2023), as they addressed the same problem and have been publicly available for 4-5 months.

Author response:
Regarding point 3, we agree that these are relevant papers to our work. Fumero et al primarily focuses on the meta-learning setting and does not provide identifiability results, however we agree that the work by Lachapelle et al would be a valuable baseline as it provides a more detailed discussion of identifiability. We will be incorporating this baseline into our experiments and update the paper accordingly in the coming days, however we urge the reviewer to reconsider our contributions even without these additional results.

Reviewer comment (Theory):
Finally, I find the feature-space linearity assumption \(E(y\mid x) = w_t^\top f(x)\) somewhat unsettling for the use in causal representation learning. While similar assumptions have appeared in stylized theoretical analyses for multi-task learning, designing a causal representation learning procedure based on such assumptions appears to be asking much more.

Author response:
The reviewer’s main concern lies in our Gaussianity assumptions. We feel that points 1 and 4 may be related in this regard:

- The feature space linearity (point 4) is accurate precisely because of the first stage of our procedure: namely this is not an assumption but a guarantee which follows from linear identifiability of the representations learned from the multitask regression network (theorem 3.2). More generally, any non-linear neural network broadly makes this assumption (where the conditional expectation is parametrized by a linear combination of the last layer activations), which is in general less well motivated than in this work given the lack of generic identifiability results. We would be happy to clarify this further if we are misunderstanding any of the reviewer’s concerns here.

Reviewer comment (Theory):
I am not convinced that the assumption of a *correctly specified* Gaussian likelihood is not a natural choice, as is also noted by reviewer 2cPA. Currently your proofs rely on the likelihood being correctly specified; the best way to address this concern is to rewrite them so that they apply to (certain) misspecified cases.

Author response:
Because of the feature space linearity, the second stage procedure reduces to a linear problem. A gaussian likelihood is the widely accepted standard for linear regression models, and is thus a very natural choice. Given that the linearity only arises in the context of arbitrarily non-linear feature mappings $f(x)$ provided by the multitask neural network in the first stage, we do not believe this to be particularly restrictive. We would point to our empirical evaluation as evidence that our model can indeed be useful. The connection to our causal representation learning result, recovering the linear transformation to the latents, is simply a very appealing byproduct of maximizing the likelihood in equation (11); it does not inherently require different assumptions on the regression setting.

Reviewer comment (Novelty):
**Some of the claims with respect to identifiability results feel exaggerated**: "*We emphasize that this is a stronger identifiability result than identifiability up to block permutations and scaling (i.e., strong identifiability) as in prior works*", "*We emphasize that this strictly strong identifiability class is stronger than the strong identifiability class as in prior works (Khemakhem et al., 2020a; Lu et al., 2022) which are only identifiable up to*" -- Similarly strong identifiability results have been shown in a general setting in Halva et al. (2021), please see their Theorem 2 in particular and its proof that share similar ideas. Theorem 3, 4 in Halva and Hyvarinen (2020) show strong identifiability in a more specific case, with the similarity here being the exploitation of linear independence and Gaussianity. The authors have cited these works elsewhere and should therefore be also aware of their identifiability results. Further, Morioka et al. (Independent Innovation Analysis, 2021) shows similar identifiability results and proof in their Theorem 2 (though for a different model). Theorem 3.2. and the MTRN identifiability results appear to be essentially the same as those of iVAE (Khemakhem et al.) so I am having hard time understanding whether it has any novel contribution in itself or whether it only sets up the 'second part' with stronger identifiability.

Author response:
Thank you for your response. We respectfully disagree about the statements that [1] and [2] provide equally strong identifiability results. There are multiple differences between these works and ours. In both of these cases, and with iVAE, the nature of the permutation identifiability is weaker than ours. Specifically, each of these prior works are permutation-identifiable either up to pointwise non-linearities of the latent components [1] and [2], or up to block-permutations of the sufficient statistics, not of the latent variables themselves (iVAE). See the specific passages in each prior work:

[1] : “We show that IIA guarantees the identifiability of the innovations with arbitrary nonlinearities, up to a permutation and componentwise invertible nonlinearities”

[2]: “$f^{−1}$ can be recovered up to permutation and coordinate-wise transformations from the distribution of $(f(s_{t_1} ), . . . ,f(s_{t_{m}}))$” (Theorem 2). “Coordinatewise” here has no restrictions and thus implies non-linearities as well.

iVAE: “. That is, we can recover the original latent variables up to a component-wise (pointwise) transformations $T_i^{*}$, $T_i$, which are defined as the sufficient statistics of exponential families.”

And

 “These two Theorems imply that in most cases $f^{-1} \circ f : Z → Z$ is a pointwise nonlinearity, which essentially means that the estimated latent variables $z$ are equal to a permutation and a pointwise nonlinearity of the original latents $z$.”

Our work provides identifiability up to permutation and scaling (linear) of the actual latent components. 

### Other key differences

1. All of these methods are optimised using variational inference. This makes our model much simpler to optimise. This is a crucial difference for practical use.
2. [2] exploits specific temporal or spatial structure in the encoded latents, while factorising the joint into a switching linear dynamical system, which is a different setting to ours.

Regarding the theoretical novelty of our identifiability results against iVAE, we note that our Theorem 3.2 for MTRN in Section 3.2 is the linear identifiability result for the multi-task regression network, which is completely different to iVAE. We believe that the reviewer meant the similarity between our Theorem 3.5 for MTLCM in Section 3.3.4 and iVAE. Below, we reproduce specific key equations (Eqs 60 and 61 in our paper) that expand upon the identifiability results of iVAE:

$$l_i^{-1}=\frac{\partial T_i (l^{-1}(z))}{\partial z_i}=\sum_{j=1}^k M_{i,j} \frac{\partial \bar{T_j} (z)}{z_i}=M_{i,i}+2M_{i,2i}z_i$$

$$2l_i^{-1}z_i=\frac{\partial T_{2i}(l^{-1}(z))}{\partial z_i}=\sum_{j=1}^k M_{2i,j} \frac{\partial \bar{T_j} (z)}{z_i}=M_{2i,i}+2M_{2i,2i}
z_i$$

both of which hold for $i=1,\cdots, d$ since $l^{-1}(z)=[l_1^{-1}(z_1),\cdots,l_d^{-1}(z_d)]$ is a linear pointwise function (rather than nonlinear point-wise as in iVAE) thanks to our linear likelihood (this does not hold in iVAEs because their mixing function f is nonlinear). By matching the coefficients, this implies that $M_{i,2i}=0$ and $M_{2i,i}=0$, which reduces the identifiability class from block-identifiable to permutation-identifiable.

We have clarified these key differences in the manuscript, and sincerely hope we have addressed the concerns in this regard

Reviewer comment (Theory):
Would be nice to have some references/citations if you make strong statements like this "This contrasts with current state-of-the-art approaches, 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." Which works are these? Why are they difficult to train, and where is that shown? And again last part about identifiability is not even true as shown above.

Author response:
1. On the statement contrast with current SOTA approaches, this statement referred to iVAE and iCaRL, the supporting evidence for which is in our experiments.  We have clarified this in the paper. Furthermore, the Lachapelle et al (2023) approach is also difficult to train due to its bilevel optimization objective.

Reviewer comment (Novelty):
"... Willetts & Paige (2021); Kivva et al. (2022) recently extend the results in unsupervised generative models to the case of models with mixture model priors." Not all these papers assume mixture model priors specifically.

Author response:
2. On the citation for identifiable unsupervised learning with mixture model priors, we have removed the two references which were not applicable for this statement.

Reviewer comment (Experiments):
Given the following statement, you should have included them as baselines: "concurrent work (Lachapelle et al., 2023; Fumero et al., 2023) has expanded this area of research to consider the multi-task and meta-learning settings."

Author response:
3. We are working to include the Lachapelle et al (2023) approach in our baselines. The Fumero et al 2023 paper does not propose identifiability results, and is less relevant since it focuses on meta-learning rather than multitask learning.

Reviewer comment (Writing):
"...existence of d independent ground-truth..." please be careful of the word independent and how it's used.

Author response:
4. “..existence of d independent…” has been reworded to “existence of d linearly independent ground-truth task-specific vectors..”

Reviewer comment (Writing):
"The marginal likelihood of $\psi$ under MTLCM is given by" imprecise language -- it's not the likelihood of $\psi$ as they are not random variables.

Author response:
The statement on marginal likelihood has been corrected.