1 Introduction↩︎

Predictive models deployed in feedback loops with the systems that produce their training data face a basic problem: the distribution they are trained on diverges from the one they are asked to score at serving time. In this work we focus on recommender systems, where the (user, item, context) tuples appearing in training logs are those a deployed policy already chose to surface, and are further filtered by user engagement. This phenomenon, broadly referred to as exposure bias in the recommendation literature [1], means that at serving time the same model is asked to score a much broader and differently distributed candidate set. The gap between training and serving-time distributions is widely understood to be a primary reason that offline ranking metrics are often unreliable predictors of online performance. Causal representation learning (CRL) [2] offers a principled solution. If a model’s predictions rely only on causal mechanisms that generate the outcome, they should remain valid under shifts that leave those mechanisms intact, including the policy-induced shift between training and serving distributions. Classical CRL aims to recover latent causal factors with strong identifiability guarantees [3], [4], but the assumptions it relies on—interventional or randomised data, anchor variables with known causal semantics, or multiple environments—are difficult to satisfy on observational recommendation logs. A second strand applies CRL to recommendation directly, including in production. [5] treat user feature drift as an intervention via a counterfactual VAE that infers latent user state and performs post-intervention inference at serving time. [6] block-identify causal treatment representations via labelled contrastive pairs that stratify on confounders and outcomes. DCCL [7] and Rec4Ad [8] disentangle specific bias mechanisms (popularity, sample-selection) with components tied to those mechanisms, and have been deployed online at Kuaishou and Taobao respectively.

Similar to the recsys-CRL works above, we assume that the confounders are observed and present in the training logs, encoded in the user, context, and item features that the deployed policy conditions on. In this case, full identification is more than what is needed to address the specific failure mode introduced by exposure bias. The spurious pathway can be suppressed directly by enforcing that the learned representation carries no information about the confounder beyond what the supervised target already requires (Figure 1). This narrower goal does not recover latent factors or enable counterfactual queries, but it does break the dependence that drives the offline-to-online gap, and it does so using only the existing observational training data. We instantiate this idea by learning a treatment representation that remains predictive of the outcome conditional on the observed confounder, while penalising mutual information between the learned representation and the confounder. The resulting criterion keeps the outcome-relevant causal content of the treatment and removes the confounder-dependent non-causal leakage that standard rankers may otherwise absorb. We make the following contributions:

• We propose an information-theoretic disentanglement criterion and prove that, at the population level, it favors representations that retain the causal content of the treatment and discard confounder-dependent non-causal information. We then derive a tractable variational lower bound that makes the criterion optimisable from finite observational data alone.

• We introduce a deployment-ready CRL recipe for robustness under distribution shift: a straightforward modification of any standard supervised model, trained on the existing confounded log, with no inference-time cost and no requirement for randomised or interventional data.

• We validate the same recipe at three scales: a synthetic SCM, the KuaiRand benchmark [9], and a live A/B test on a production personalised playlist generation surface with millions of users on the music streaming platform Spotify. Across all three settings, a capacity-matched CRL variant matches the baseline on in-distribution metrics and outperforms it whenever evaluation moves off the training distribution: on a held-out random-exposure split, under intervention, and in a live serving production environment.

2 Method↩︎

2.1 Setting↩︎

We adopt the Structural Causal Model framework of [10]. A causal model is a tuple\(\mathcal{M}=(\mathbf{V},\mathbf{U},\mathcal{F},P(\mathbf{U})),\) where \(\mathbf{V}\) is the set of endogenous variables, \(\mathbf{U}\) is the set of exogenous noise variables with joint distribution \(P(\mathbf{U})\), and \(\mathcal{F}=\{f_V\}_{V\in\mathbf{V}}\) is a collection of structural assignments \(V := f_V(\mathrm{Pa}(V),U_V).\) In our setting, \(\mathbf{V}=\{T,Y,Z\}\), where \(Z\) is an observed confounder (in recommendation, a user representation), \(T\) is the treatment to be ranked (an item), and \(Y\) is the outcome of interest. As illustrated in Figure 1, the observed treatment is an entangled mixing \(T=m(T_C,T_{nC})\) of latent causal components \(T_C\) and non-causal components \(T_{nC}\). We make two structural assumptions: (A1) \(Y \perp T_{nC} \mid T_C, Z\) meaning that the outcome depends on the treatment only through its causal latents, given the confounder; (A2) \(I(T_{nC}; Z \mid T_C) > 0\): which we call unique spurious leakage: the non-causal latents retain information about the confounder that is not already contained in the causal latents. This is the setting that arises in recommender systems whenever the deployed policy has correlated non-causal item-side factors with user-side signal, creating a spurious backdoor path \(T \leftarrow T_{nC} \leftarrow Z \to Y\) that a standard predictor learning \(\mathbb{E}[Y \mid T, Z]\) will absorb. Such structures arise also in many causal representation learning settings where spurious correlations between treatment and confounders lead to biased models if not explicitly addressed.

While maximising \(I(T;Y \mid Z)\) ensures the learned representation is predictive of the outcome, it may still preserve spurious signals as it does not penalise representations for carrying extra non-causal information. To address this, we introduce a more targeted objective that trades predictive utility against confounder dependence.

2.2 A New Objective for Causal Disentanglement↩︎

We propose to learn a representation \(g(T)\) that maximizes the following: \[\label{eq:Jg} \tcboxmath[ colback=gray!5, colframe=gray!55, boxrule=0.6pt, arc=1.5mm, left=2mm, right=2mm, top=1.5mm, bottom=1.5mm ]{ J(g) = \underbrace{I(g(T);Y\mid Z)}_{\text{Utility}} - \lambda\, \underbrace{I(g(T);Z)}_{\text{Penalty}} }\tag{1}\] called the Disentanglement Criterion, where \(\lambda \ge 0\) is a hyperparameter that controls the trade-off between predictive utility and confounder invariance. This objective balances two forces: the utility term rewards representations that are predictive of the outcome \(Y\) beyond \(Z\), thereby encouraging the preservation of \(T_C\); the penalty term discourages information shared with \(Z\), thereby suppressing \(Z\)-correlated non-causal factors. Under our structural assumptions, the penalty term is exactly the channel through which \(T_{nC}\) leaks into the prediction via the spurious path \(T_{nC}\leftarrow Z\to Y\). An equivalent bounded reparameterization is given in Appendix 7. We first verify that the criterion ranks the ideal causal representation above the naive entangled one.

Proposition 1 compares two fixed representations. We next formalise the more general mechanism through which the disentanglement criterion removes non-causal information. The observed treatment \(T=m(T_C,T_{nC})\) may already lose information about \(T_C\), as does a learned representation \(g(T)\). We therefore define \(\phi(T_C)\) as the information about \(T_C\) recoverable from \(g(T)\) via a deterministic map \(r\), i.e.\(r(g(T))=\phi(T_C)\). A causally purified representation is then \(\bar g:=r(g)=\phi(T_C)\). Because \(\bar g\) is a deterministic function of \(g\), the information in \(g\) decomposes exactly into the recovered causal content \(\phi(T_C)\) and a residual beyond it. We assume that this residual carries no additional predictive information about \(Y\) given \(Z\), but remains informative about the confounder, which are the representation-level counterparts of A1 and A2: \(I(g(T);\,Y\mid\phi(T_C),Z)=0,~I(g(T);\,Z\mid\phi(T_C))>0.\) The following proposition shows that stripping the residual strictly improves \(J\).

However, the form of \(\phi\) remains unspecified. Should it preserve all information in \(T_C\) (i.e.be invertible), or could a lossy compression be preferable? Since \(T_C\) is correlated with both \(Y\) (good for utility) and \(Z\) (costly for the penalty), any lossy \(\phi\) reduces the penalty at the cost of some utility, with the trade-off governed by \(\lambda\). To formalise this, let \(\phi_\text{id}\) denote the identity (uncompressed) and let \(\phi\in\mathcal{H}\) denote any non-invertible alternative transformation, where \(\mathcal{H}\) is the class of all non-invertible lossy functions applicable to \(T_C\). For each \(\phi\in\mathcal{H}\) we define \[\begin{align} \Delta U(\phi) := I(T_C;Y\mid Z)-I(\phi(T_C);Y\mid Z), \qquad \Delta P(\phi) := I(T_C;Z)-I(\phi(T_C);Z). \end{align}\] where \(\Delta U(\phi)\) is the utility lost by compressing the causal component, while \(\Delta P(\phi)\) is the reduction in confounding penalty gained by that compression. We assume a no-free-lunch condition according to which every nontrivial lossy transformation \(\phi\in\mathcal{H}\) under consideration incurs a genuine utility loss and yields a genuine penalty reduction: \[\begin{align} \label{eq:nofree} \Delta U(\phi)>0~\text{and}~\Delta P(\phi)>0. \end{align}\tag{2}\] We further define the critical penalty weight as the ratio of the two: \[\begin{align} \Lambda(\phi) := \frac{\Delta U(\phi)}{\Delta P(\phi)}. \end{align}\]

Proposition 3 shows that after non-causal leakage has been removed, sufficiently small penalty weights \(\lambda\) preserve the full causal component, and \(\phi_\text{id}\) is optimal; compression of \(T_C\) becomes preferable only once the penalty reduction outweighs the corresponding utility loss. Finally, we provide a concrete sanity check for the preceding claims in Appendix 8 for a linear Gaussian SCM where all mutual-information terms can be computed in closed form: the appendix verifies that \(J\) strictly prefers the causal representation \(T_C\) over the naive entangled representation \((T_C,T_{nC})\), and illustrates the lossless regime by varying the amount of stochastic compression applied to \(T_C\).

2.3 A Tractable Objective For Causal Disentanglement↩︎

The criterion in 1 cannot be optimised directly since both \(I(g(T);Y\mid Z)\) and \(I(g(T);Z)\) are unknown and intractable to compute from finite data. We therefore derive a tractable surrogate by replacing each term with a computable bound: a lower bound on the utility and an upper bound on the penalty so that the surrogate is a valid lower bound on \(J(g)\).

2.3.0.1 Lower bound on the utility.

For any conditional distribution \(q(y\mid g(T),z)\), the standard variational bound on conditional mutual information (CMI) [11] gives \[\begin{align} I(g(T);\,Y\mid Z) \;\ge\; \mathbb{E}_{p(g(T),y,z)}\!\bigl[\log q(y\mid g(T),z)\bigr] + H(Y\mid Z), \end{align}\] which motivates learning \(q\) to approximate the true posterior \(p(y\mid g(T),z)\). We instantiate this via a contrastive InfoNCE estimator [12]: for each anchor \((g(T),y^+,z)\) we draw \(K\) negative outcomes \(y^-_1,\ldots,y^-_K\) from \(p(y\mid z)\) and define \[\begin{align} \mathcal{L}_{\mathrm{utility}}(\theta) := -\mathbb{E}\!\left[ \log\frac{e^{f_\theta(g(T),\,y^+)/\tau}}{e^{f_\theta(g(T),\,y^+)/\tau} +\sum_{k=1}^K e^{f_\theta(g(T),\,y^-_k)/\tau}} \right]. \end{align}\] The InfoNCE loss directly lower-bounds the CMI [12]: \[I(g(T);\,Y\mid Z) \;\ge\; \log(K{+}1) - \mathcal{L}_{\mathrm{utility}}(\theta) =: -\mathcal{L}_{\mathrm{utility}}(\theta) + C_Y, \quad C_Y := \log(K{+}1). \label{eq:infonce}\tag{3}\] In the recommender setting this is exactly the in-batch cross-entropy ranking loss with negatives drawn from the same user context, so \(\mathcal{L}_{\mathrm{utility}}\) is the loss the production ranker is already trained on.

2.3.0.2 Upper bound on the penalty.

For any representation \(g(T)\) and critic \(d_\phi\), the NCE-CLUB estimator [13] upper-bounds the mutual information between \(g(T)\) and \(Z\): \[\label{eq:nceclub} I(g(T);\,Z) \;\le\; \underbrace{ \mathbb{E}_{p(g(T),z)}\bigl[d_\phi(g(T),z)\bigr] -\mathbb{E}_{p(g(T))p(z)}\bigl[d_\phi(g(T),z)\bigr] }_{=:\,I_{\mathrm{NCE\text{-}CLUB}}(g(T);\,Z;\,d_\phi)}.\tag{4}\] The first term scores positive pairs \((g(T),z)\sim p(g,z)\), while the second scores negative pairs drawn independently from \(p(g)p(z)\). The bound is tight when \(d_\phi\) recovers the log-density ratio \(d^\star(g,z)=\log p(z\mid g)-\log p(z)\). In practice, we approximate this via a learnable critic \(d_\phi(g(T),z)\approx\log q_\phi(z\mid g(T))-\log p(z)\), learned by discriminating joint samples \((g,z)\sim p(g,z)\) from independent samples \((g,z)\sim p(g)p(z)\). Further details are provided in Appendix 9.

Substituting 3 and 4 into 1 , yields a tractable variational lower bound for \(J(g)\).

The resulting objective augments the standard ranking loss with an explicit mutual-information penalty that discourages dependence between the learned representation and the confounder, unlike IPM-based methods [14] which require explicit population matching.

3 Related work↩︎

Our work draws on causal representation learning, debiasing for recommender systems, and adversarial mutual-information estimation. As discussed in §1, our work is closely related to [5], [6], DCCL [7], and Rec4Ad [8]. In particular, [6] requires labelled contrastive pairs built by stratifying on confounders and outcomes, whereas the present work uses a single information-theoretic penalty with no labelled-pair, propensity, or component-specific machinery.

3.0.0.1 Treatment-effect estimation with neural networks.

CFR/TARNet [14] balances treatment-conditional representations via integral-probability metrics. CEVAE [15] addresses unobserved confounders with a deep generative model. Dragonnet [16] adds a propensity-prediction head to encourage representations sufficient for treatment-effect identification. Structured-treatment work targets graph- or text-valued \(T\) [17]–[19].

3.0.0.2 Debiasing causal recommendation.

IPS reweighting [20], [21] and its variants—asymmetric tri-training [22], instrumental-variable methods such as IViDR [23] and iDCF [24], and explicit causal-graph formulations [25], [26], correct exposure bias by reweighting samples or estimating an inverse mechanism. Earlier representation-side precursors disentangle bias-specific signals: causal embeddings [27] for the treated/control split and DICE [28] for interest vs.conformity. The industrial deployments DCCL and Rec4Ad extend this to multi-component pipelines tied to specific bias mechanisms (popularity, sample-selection).

3.0.0.3 Adversarial MI estimation and invariance.

The practical estimator combines gradient reversal [29] with the InfoNCE lower bound [12]. Alternative MI estimators include MINE [30] and the CLUB upper bound [31]; NCE-CLUB [13] provides the upper-bound machinery used in our diagnostics (Appendix 9). MI-based representation objectives are also at the core of Deep InfoMax [32]. The disentanglement criterion \(J(g)\) is complementary to invariant risk minimisation [33], [34]: IRM identifies predictors that are invariant across multiple environments, whereas we suppress dependence on a single observed confounder when no environment partition is available.

3.0.0.4 Causal bottlenecks and abstractions.

Our objective is related to causal information-bottleneck methods, which retain task-relevant causal information while discarding irrelevant variation. The Causal Information Bottleneck [35] learns causal variable abstractions by compressing an input while preserving causal control over a target, and Structural Causal Bottleneck Models [36] assume that causal effects between high-dimensional SCM variables are mediated by lower-dimensional bottlenecks of their causes. These perspectives connect to causal abstraction [37], which seek coarser variables or models that are interventionally consistent with the original system, together with methods that learn such maps [38]–[42]. In the language of Figure 1, such methods would ask whether the entangled \(T=m(T_C,T_{nC})\) admits a coarser description, ideally \(T_C\), that preserves interventional consistency. Our setting is adjacent but deliberately weaker: with \(m\) fixed and unobserved, we learn a map \(T \mapsto g(T)\) sufficient for the recoverable causal content \(\phi(T_C)\), while removing the \(Z\)-correlated non-causal leakage carried by \(T_{nC}\). We do not identify \(T_C\), learn \(m\), or verify model-level consistency; we learn a representation that retains the causal information needed for prediction.

4 Experiments↩︎

We evaluate the proposed method in three settings of increasing realism. The synthetic SCM (§4.1) shows that the method successfully isolates the causal components of the treatment representation on a problem where the causal/non-causal split is known. The KuaiRand experiment (§4.2) evaluates the method on real recommendation data, training on a confounded algorithmic-exposure log and measuring offline performance under a held-out random-exposure distribution. The online A/B test (§4.3) verifies that the same offline robustness translates into measurable online gains on a deployed production model.

The objective in ?? bounds the penalty \(I(g(T);Z)\) from above via NCE-CLUB [13]. In the experiments that follow we instead reduce an InfoNCE lower bound [12] on the same MI via gradient reversal [29], and monitor the NCE-CLUB upper bound as a diagnostic. Figure 2 shows both estimators on a production training run where they decay together once gradient reversal becomes active, with the gap closing as residual MI shrinks. This is a convenient implementation detail that works well in our production setup.

4.1 Synthetic SCM↩︎

We instantiate a linear additive SCM with a confounder \(X \in \mathbb{R}^{10}\), a treatment \(T \in \mathbb{R}^{10}\) decomposable into \(d_c = 5\) causal and \(5\) non-causal dimensions, and an outcome \(Y\) that depends on \(T_C\) and \(X\) but not \(T_{nC}\). Full data-generating equations are given in Appendix 10. Since \(Y\) is invariant to interventions on \(T_{nC}\), a model that has correctly isolated the causal part of the treatment should also be invariant to those interventions. We measure the intervention sensitivity \(|\hat{Y}(T) - \hat{Y}(T')|\) where \(T'\) has its non-causal coordinates resampled, and the prediction MAE on \(Y\). In Figure 3 we sweep the outcome noise \(\sigma_Y\) from \(0\) to \(1\), and show how the baseline CATE model is sensitive to interventions on the \(T_{nC}\) non-causal signal, whereas the CRL variant’s intervention sensitivity sits essentially at zero across the entire range. The MAE is matched between the two methods.

4.2 KuaiRand↩︎

The KuaiRand-Pure [9] dataset contains \(\sim\)​1.4M algorithmic-exposure interactions and \(\sim\)​1.2M random-exposure interactions over the same \(27{,}285\) users and \(7{,}583\) items. The algorithmic split is heavily confounded; the random split is, by construction, an unbiased sample from the user \(\times\) item space. The pair gives us a direct measurement of distribution shift on a problem where the confounding is real, large, and known to live primarily in which (user, item) pairs are observed rather than in any single feature.

We use a standard DeepFM [43] as our click-through rate (CTR) backbone, with learned \(16\)-dimensional user-id and video-id embeddings, \(\sim\)​15 continuous video features passed through an MLP encoder, and a DNN tower with hidden sizes \([32, 16]\). DeepFM is a standard baseline in the recommendation literature [44], and has been applied to the KuaiRand dataset [45], [46]. Our CRL variant adds a bilinear MI critic between the item representation and the user representation, attached via gradient reversal. Architecture and capacity are otherwise identical between baseline and CRL. Training details are in Appendix 11.

In Table 1, Biased validation AUC measures ranking quality on the training distribution, and Random test AUC is evaluated on the held-out random-exposure split. Comparing the random test AUC across baseline and CRL variants indicates the robustness to distribution shift. We additionally report Recall@5 and NDCG@5 [23], [24]. For each user in the random-exposure test split we rank the items they were exposed to by predicted click probability; Recall@\(K\) is the fraction of that user’s positives (is_click=1) that appear in the top-\(K\), and NDCG@\(K\) uses binary relevance with discount \(1/\log_2(i+1)\) at rank \(i\) and IDCG computed over the user’s true positives. Both are averaged over users with at least one positive and at least five exposed items in the random split.

Over \(500\) seeds, the baseline and CRL variants are close on the training distribution: biased validation AUC is \(0.7574 \pm 0.0012\) for the baseline and \(0.7508 \pm 0.0109\) for the CRL variant (\(\Delta = -0.0066\), intervals here and below are per-seed standard deviations). On the unconfounded random-exposure test split, the CRL variant attains AUC \(0.6582 \pm 0.0056\) versus \(0.6453 \pm 0.0034\) for the baseline, an improvement of \(+0.0129\) AUC, roughly \(3.8\times\) the baseline’s per-seed standard deviation. The improved Random test AUC shows that the CRL model generalises better under distribution shift. This pattern—offline parity on the in-distribution split combined with gains on a held-out random-exposure split—is similar to what [47] call effective robustness: improvement under shift not predicted by in-distribution accuracy alone. In the online A/B test (Section 4.3) we show how this improved generalisation ability translates to real gains in listener engagement.

4.3 Online A/B test on personalised playlist generation↩︎

The KuaiRand experiment established that the CRL variant generalises better to a held-out distribution than its capacity-matched baseline. The strongest evidence that this offline robustness translates into real-world impact is an online A/B test on a production personalised playlist generation surface at Spotify, conducted under standard online controlled-experiment methodology with millions of users.

The test compares two versions of the production session-ranking model: the deployed baseline and a capacity-matched CRL variant trained with the disentanglement objective of §2. The CRL variant uses an adversarial loss to decorrelate the model’s internal representation of the candidate track from its representations of session- and user-level context, mirroring the recipe used in the KuaiRand experiment.

Table 2 reports the headline engagement metrics, measured over two weeks during the test. The CRL variant increases track streams by \(+0.75\%\) and minutes played by \(+0.50\%\) while decreasing track skips by \(-0.61\%\), with all three changes significant at the \(95\%\) level. The magnitude of these gains is of the same order as the improvement that the production ranker, itself a ground-up rebuild, delivered over its mature predecessor.

The result mirrors the KuaiRand finding. In both settings the CRL and baseline variants are architecturally and capacity-matched, and in both they perform similarly on the in-distribution offline metric. The CRL variant’s advantage materialises only when evaluation moves to a distribution that differs from the training one: the random-exposure split on KuaiRand, and the live serving distribution in production. This is precisely the regime in which the theory predicts a causal representation should help, and it is the regime that matters at deployment.

4.4 Discussion↩︎

The three experiments tell a consistent story. The synthetic SCM shows that the proposed adversarial penalty removes the non-causal components of the treatment representation while preserving the causal ones, robustly across outcome-noise levels. The KuaiRand experiment shows that this approach has the expected consequence on real data: indistinguishable in-distribution offline performance, but a meaningful robustness gain under a known distribution shift. The online A/B test shows that the same robustness gain translates to a production setting, where it shows up as improved listener engagement. Note that in all three settings, standard offline evaluation on the training distribution is essentially uninformative about whether the CRL variant is better. The synthetic MAE is unchanged, the KuaiRand biased-validation AUC is unchanged, and the offline metrics for the production model and its CRL variant are unchanged. The CRL impact is apparent only when evaluated against a distribution other than the training one. Offline-only comparisons of recommendation models can systematically under-detect the kind of generalisation improvements that matter at deployment.

5 Conclusion↩︎

We proposed a CRL-motivated method for making predictive models more robust under the distribution shift characteristic of on-policy recommendation. At its core is the Disentanglement Criterion, a new information-theoretic objective whose population-level optimum discards non-causal treatment latents; we characterise the regime under which the full causal component is preserved and derive a tractable variational lower bound that reduces to a straightforward modification of any standard supervised model. Evaluated on a synthetic SCM, the KuaiRand benchmark, and a live A/B test on a production music-streaming recommender, the method matches in-distribution offline performance in all three settings and improves robustness under distribution shift. The A/B test shows significant gains in online listener engagement. Limitations include the inability to empirically verify A1 and A2 on real data, a single product surface for the online evaluation, and the potential need for more expressive critics under richer confounding structures. Natural extensions include temporal confounder structures, characterising when in-distribution metrics cease to predict online gains, stress-testing across additional production recommenders, and connecting the criterion to causal abstraction frameworks by studying when \(g(T)\) can be interpreted as an interventionally consistent abstraction from \(T\) to \(\phi(T_C)\).

6 Proofs for §2↩︎

Throughout, we use the structural assumptions introduced in §2. The observed treatment decomposes as \[T=m(T_C,T_{nC}),\] the outcome is conditionally independent of the non-causal latents, \[\boldsymbol{(A1)}\qquad Y\perp T_{nC}\mid(T_C,Z),\] and the non-causal latents contain residual confounder information beyond the causal latents, \[\boldsymbol{(A2)}\qquad I(T_{nC};Z\mid T_C)>0.\]

Proof of Proposition 1↩︎

We compare the ideal representation \(g^\star(T)=T_C\) with the naive representation \(\tilde{g}(T)=T=(T_C,T_{nC})\).

6.0.0.1 Utility.

By the chain rule for conditional mutual information, \[\begin{align} I(\tilde{g}(T);Y\mid Z) &= I(T_C,T_{nC};Y\mid Z) \\ &= I(T_C;Y\mid Z) + I(T_{nC};Y\mid T_C,Z). \end{align}\] The second term is zero by A1. Hence \[I(\tilde{g}(T);Y\mid Z) = I(T_C;Y\mid Z) = I(g^\star(T);Y\mid Z).\]

6.0.0.2 Penalty.

Again, by the mutual-information chain rule, \[\begin{align} I(\tilde{g}(T);Z) &= I(T_C,T_{nC};Z) \\ &= I(T_C;Z)+I(T_{nC};Z\mid T_C). \end{align}\] By A2, \(I(T_{nC};Z\mid T_C)>0\). Therefore \[I(\tilde{g}(T);Z)>I(T_C;Z)=I(g^\star(T);Z).\]

Equivalently, using the entropy chain rule, \[\begin{align} I(T_C,T_{nC};Z) &= H(T_C,T_{nC})-H(T_C,T_{nC}\mid Z) \\ &= \bigl[H(T_C)+H(T_{nC}\mid T_C)\bigr] - \bigl[H(T_C\mid Z)+H(T_{nC}\mid T_C,Z)\bigr] \\ &= I(T_C;Z)+I(T_{nC};Z\mid T_C). \end{align}\] Thus, the excess penalty of the naive representation is exactly the unique spurious leakage term.

6.0.0.3 Conclusion.

The two representations have equal utility, but the naive representation has a strictly larger penalty. Hence \[\begin{align} J(g^\star)-J(\tilde{g}) &= \bigl[I(g^\star;Y\mid Z)-I(\tilde{g};Y\mid Z)\bigr] + \lambda\bigl[I(\tilde{g};Z)-I(g^\star;Z)\bigr] \\ &= \lambda\,I(T_{nC};Z\mid T_C)>0. \end{align}\] Therefore \(J(g^\star)>J(\tilde{g})\). 0◻

Proof of Proposition 2↩︎

Consider a representation satisfying the recoverability assumption: \[g(T)=h\bigl(\phi(T_C),\ell(T_{nC})\bigr),\] where the causal component \(\phi(T_C)\) is deterministically recoverable from \(g(T)\). That is, there exists a map \(r\) such that \[r(g(T))=\phi(T_C).\] Define the purified representation by \[\bar g(T):=r(g(T))=\phi(T_C).\] Thus \(\bar g\) is a deterministic function of \(g\). At the representation level, we assume that the residual information in \(g(T)\) beyond the recovered causal component contains no additional outcome information but does contain residual confounder information: \[\boldsymbol{(R1)}\qquad I(g(T);Y\mid \phi(T_C),Z)=0,\] and \[\boldsymbol{(R2)}\qquad I(g(T);Z\mid \phi(T_C))>0.\]

6.0.0.4 Utility preservation.

Since \(\bar g=\phi(T_C)\) is a deterministic function of \(g\), the pair \((g,\bar g)\) contains the same information as \(g\). Hence \[I(g;Y\mid Z)=I(g,\bar g;Y\mid Z).\] By the chain rule for conditional mutual information, \[\begin{align} I(g,\bar g;Y\mid Z) &= I(\bar g;Y\mid Z) + I(g;Y\mid \bar g,Z). \end{align}\] Since \(\bar g=\phi(T_C)\), the second term is \[I(g;Y\mid \bar g,Z) = I(g;Y\mid \phi(T_C),Z),\] which is zero by R1. Therefore \[I(g;Y\mid Z)=I(\bar g;Y\mid Z).\] Thus purification preserves the utility term.

6.0.0.5 Penalty reduction.

Again, since \(\bar g\) is a deterministic function of \(g\), \[I(g;Z)=I(g,\bar g;Z).\] Applying the mutual-information chain rule, \[\begin{align} I(g,\bar g;Z) &= I(\bar g;Z)+I(g;Z\mid \bar g). \end{align}\] Using \(\bar g=\phi(T_C)\), we obtain \[I(g;Z)-I(\bar g;Z) = I(g;Z\mid \bar g) = I(g;Z\mid \phi(T_C)).\] By R2, this quantity is strictly positive. Hence \[I(\bar g;Z)<I(g;Z).\]

6.0.0.6 Objective comparison.

Combining utility preservation and penalty reduction, \[\begin{align} J(\bar g)-J(g) &= \bigl[I(\bar g;Y\mid Z)-I(g;Y\mid Z)\bigr] + \lambda\bigl[I(g;Z)-I(\bar g;Z)\bigr] \\ &= \lambda\,I(g;Z\mid \phi(T_C))>0. \end{align}\] Therefore \(J(\bar g)>J(g)\) for every \(\lambda>0\). 0◻

General Predictive Purification Result↩︎

The purification result in the main text assumes that the causal component \(\phi(T_C)\) is recoverable from the representation. We now show that a more general purification principle holds without this assumption.

Let \[G := g(T)\] be an arbitrary representation. We call \(\bar G=s(G)\) a predictive purification of \(G\) if it is a deterministic function of \(G\) satisfying \[Y\perp G\mid(\bar G,Z).\] In words, once \(\bar G\) and \(Z\) are known, the remaining information in \(G\) does not provide additional predictive information about \(Y\).

6.0.0.7 Existence.

Under standard regularity conditions, such a statistic exists. One canonical construction is the conditional response profile \[s(G) := \bigl\{p(Y\mid G,Z=z):z\in\mathcal{Z}\bigr\}.\] This object records, for each value of \(G\), the conditional law of \(Y\) across values of the confounder. By construction, \[p(Y\mid G,Z)=p(Y\mid s(G),Z),\] which implies \[Y\perp G\mid(s(G),Z).\] This statistic may be infinite-dimensional and is used only as a population-level object.

6.0.0.8 Connection to the causal setting.

The result above is purely information-theoretic and applies to any representation. Its causal interpretation follows from assumption A1. For the full entangled latent representation \[G=(T_C,T_{nC}),\] we have \[p(Y\mid T_C,T_{nC},Z)=p(Y\mid T_C,Z),\] so the predictive purification depends only on the causal component \(T_C\), up to outcome-equivalence. Thus, in the SCM considered in this work, the population target of purification can be written as \[g(T)=\phi(T_C).\]

6.0.0.9 Relation to the main result.

Proposition 4 shows that the disentanglement criterion removes information that is irrelevant for predicting \(Y\mid Z\) but still informative about the confounder, without requiring any structural assumptions on \(g\). The main-text result strengthens this conclusion by imposing a recoverability assumption, which allows the purified representation to be identified explicitly as a function of the causal component \(T_C\).

7 Reparameterisation of the Disentanglement Criterion↩︎

The objective in 1 , \[J(g)=I(g(T);Y\mid Z)-\lambda I(g(T);Z),\] uses the unbounded penalty weight \(\lambda\in[0,\infty)\). It is often more convenient to work with an equivalent bounded parameter \(\gamma\in[0,1]\).

7.0.0.1 Reparameterisation.

For \(\lambda>0\), maximising \(J(g)\) is equivalent to minimising \[-J(g)=\lambda I(g(T);Z)-I(g(T);Y\mid Z).\] Dividing by the positive constant \(1+\lambda\) does not change the minimiser, so the same solutions are obtained by minimising \[\frac{\lambda}{1+\lambda}I(g(T);Z) - \frac{1}{1+\lambda}I(g(T);Y\mid Z).\] Define \[\gamma:=\frac{1}{1+\lambda}\in(0,1),\] so that \(1-\gamma=\lambda/(1+\lambda)\). The equivalent bounded objective is \[L_\gamma(g) := (1-\gamma)I(g(T);Z) - \gamma I(g(T);Y\mid Z), \label{eq:Jprime}\tag{5}\] which we minimise over \(g\). The limiting case \(\lambda=0\) corresponds to \(\gamma=1\), i.e. pure utility maximisation. Conversely, \(\lambda\to\infty\) corresponds to \(\gamma\to 0\), i.e. pure penalty minimisation.

7.0.0.2 Optimality condition under \(\gamma\).

Using the same notation as Proposition 3, compare the uncompressed representation \(T_C\) with a lossy alternative \(\phi(T_C)\). The uncompressed representation is preferred under \(L_\gamma\) when \[L_\gamma(T_C)<L_\gamma(\phi(T_C)).\] Expanding this inequality gives \[(1-\gamma)\Delta P(\phi) - \gamma\Delta U(\phi) <0,\] or equivalently \[\gamma > \frac{\Delta P(\phi)}{\Delta U(\phi)+\Delta P(\phi)} =: \Gamma(\phi).\] Thus \(\Gamma(\phi)\) is the critical utility weight above which the full causal component is preferred to the lossy compression \(\phi\).

7.0.0.3 Correspondence with Proposition 3.

The \(\lambda\)- and \(\gamma\)-parameterisations describe the same trade-off. Since \[\gamma=\frac{1}{1+\lambda},\] the condition \[\lambda<\Lambda(\phi) = \frac{\Delta U(\phi)}{\Delta P(\phi)}\] is equivalent to \[\gamma> \frac{\Delta P(\phi)}{\Delta U(\phi)+\Delta P(\phi)} = \Gamma(\phi).\] Thus the low-\(\lambda\) lossless regime in Proposition 3 is the same as the high-\(\gamma\) regime in Proposition 5.

8 Toy Example↩︎

We give a closed-form sanity check for the Disentanglement Criterion: \[J(g)=I(g(T);Y\mid Z)-\lambda I(g(T);Z)\] in a linear additive Gaussian model. The treatment consists of causal and non-causal components, but only the causal ones affects the outcome. Consider the following SCM: \[\begin{align} Z &= \epsilon_Z, \\ T_C &= \alpha Z + \epsilon_C, \\ T_{nC} &= \beta Z + \epsilon_N, \\ Y &= \rho T_C + \delta Z + \epsilon_Y, \end{align}\] where all noise variables are mutually independent and Gaussian: \[\epsilon_Z\sim\mathcal{N}(0,\sigma_Z^2),\quad \epsilon_C\sim\mathcal{N}(0,\sigma_C^2),\quad \epsilon_N\sim\mathcal{N}(0,\sigma_N^2),\quad \epsilon_Y\sim\mathcal{N}(0,\sigma_Y^2).\] The observed treatment is \[T=(T_C,T_{nC}).\] By construction, the outcome depends on the treatment only through the causal component: \[Y\perp T_{nC}\mid (T_C,Z).\] Moreover, under non-degenerate parameters, e.g. \(\beta\neq0\), \(\sigma_Z^2>0\), \(\sigma_C^2>0\), and \(\sigma_N^2>0\), the non-causal component contains residual information about the confounder beyond the causal component: \[I(T_{nC};Z\mid T_C)>0.\] Indeed, since the variables are jointly Gaussian, \[I(T_{nC};Z\mid T_C) = \frac{1}{2} \log \frac{\mathrm{Var}(T_{nC}\mid T_C)}{\mathrm{Var}(T_{nC}\mid Z,T_C)}.\] Now \[\mathrm{Var}(T_{nC}\mid Z,T_C)=\sigma_N^2,\] whereas \[\mathrm{Var}(T_{nC}\mid T_C) = \beta^2\mathrm{Var}(Z\mid T_C)+\sigma_N^2,\] with \[\mathrm{Var}(Z\mid T_C) = \frac{\sigma_Z^2\sigma_C^2}{\alpha^2\sigma_Z^2+\sigma_C^2}.\] Therefore, \[I(T_{nC};Z\mid T_C) = \frac{1}{2} \log\left( 1+ \frac{\beta^2\sigma_Z^2\sigma_C^2}{\sigma_N^2(\alpha^2\sigma_Z^2+\sigma_C^2)} \right)>0.\] Thus the Gaussian SCM satisfies assumptions A1 and A2 from the main text. Consequently, we now compare the ideal causal representation with the naive entangled representation: \[g_1(T)=T_C,~~~g_2(T)=T=(T_C,T_{nC}).\] For scalar jointly Gaussian variables, \[I(U;V\mid W) = \frac{1}{2}\log \frac{\mathrm{Var}(V\mid W)}{\mathrm{Var}(V\mid U,W)}.\] Conditioned on \(Z\), \[Y=(\rho\alpha+\delta)Z+\rho\epsilon_C+\epsilon_Y,\] so \[\mathrm{Var}(Y\mid Z)=\rho^2\sigma_C^2+\sigma_Y^2.\] Conditioned on \((T_C,Z)\), \[Y=\rho T_C+\delta Z+\epsilon_Y,\] and therefore \[\mathrm{Var}(Y\mid T_C,Z)=\sigma_Y^2.\] Hence \[\begin{align} I(T_C;Y\mid Z) &= \frac{1}{2} \log \left( \frac{\rho^2\sigma_C^2+\sigma_Y^2}{\sigma_Y^2} \right) \\ &= \frac{1}{2} \log \left( 1+\frac{\rho^2\sigma_C^2}{\sigma_Y^2} \right). \end{align}\] Since \(Y\perp T_{nC}\mid(T_C,Z) \implies I(T_C,T_{nC};Y\mid Z)=I(T_C;Y\mid Z)\). So the ideal and naive representations have identical utility.

For the penalty term, \[I(T_C;Z) = \frac{1}{2} \log \left( \frac{\alpha^2\sigma_Z^2+\sigma_C^2}{\sigma_C^2} \right) = \frac{1}{2} \log \left( 1+\frac{\alpha^2\sigma_Z^2}{\sigma_C^2} \right).\] By the mutual-information chain rule, \[I(T_C,T_{nC};Z) = I(T_C;Z)+I(T_{nC};Z\mid T_C).\] The second term is strictly positive by the calculation above, so \[I(g_2(T);Z)>I(g_1(T);Z).\] Therefore, for any \(\lambda>0\), \[\begin{align} J(g_1)-J(g_2) &= \lambda \bigl[ I(g_2(T);Z)-I(g_1(T);Z) \bigr] \\ &= \lambda I(T_{nC};Z\mid T_C)>0. \end{align}\] Thus, in this Gaussian SCM, the criterion strictly prefers the causal representation \(T_C\) over the naive entangled representation \((T_C,T_{nC})\). Figure 4 shows the same effect empirically in a multivariate simulation where the representation progressively removes coordinates of \(T_{nC}\). Retaining more non-causal information leaves the utility unchanged but increases the confounding penalty, so the corresponding objective values decrease more rapidly as \(\lambda\) grows.

After the non-causal component has been removed, the remaining question is whether the optimal representation should preserve all of \(T_C\) or compress it. To model lossy compression analytically, consider the stochastic representation channel \[G = T_C+\epsilon_G, \qquad \epsilon_G\sim\mathcal{N}(0,\sigma_G^2),\] with \(\epsilon_G\) independent of all other variables. This can be viewed as a Gaussian Markov kernel from \(T_C\) to \(G\), where \(\sigma_G^2\) controls the amount of compression. The penalty is \[P(\sigma_G^2) := I(G;Z) = \frac{1}{2} \log\left( 1+ \frac{\alpha^2\sigma_Z^2}{\sigma_C^2+\sigma_G^2} \right),\] and the utility is \[U(\sigma_G^2) := I(G;Y\mid Z) = \frac{1}{2} \log\left( \frac{\rho^2\sigma_C^2+\sigma_Y^2}{ (\rho^2\sigma_C^2+\sigma_Y^2) - \frac{\rho^2\sigma_C^4}{\sigma_C^2+\sigma_G^2} } \right).\] Both terms decrease as \(\sigma_G^2\) increases: compression removes information about the confounder, but it also removes predictive information about the outcome. The objective is \[J(\sigma_G^2) = U(\sigma_G^2)-\lambda P(\sigma_G^2).\] At the uncompressed boundary \(\sigma_G^2=0\), \[\frac{dP}{d\sigma_G^2}\bigg|_{\sigma_G^2=0} = -\frac{\alpha^2\sigma_Z^2}{2\sigma_C^2(\sigma_C^2+\alpha^2\sigma_Z^2)},\] and \[\frac{dU}{d\sigma_G^2}\bigg|_{\sigma_G^2=0} = -\frac{\rho^2}{2\sigma_Y^2}.\] Thus the uncompressed representation is locally optimal whenever \[\frac{dJ}{d\sigma_G^2}\bigg|_{\sigma_G^2=0} = \frac{dU}{d\sigma_G^2}\bigg|_{\sigma_G^2=0} - \lambda \frac{dP}{d\sigma_G^2}\bigg|_{\sigma_G^2=0} \le 0.\] Equivalently, \[\lambda \le \lambda_{\mathrm{crit}} := \frac{ \left| \frac{dU}{d\sigma_G^2} \right|_{\sigma_G^2=0} }{ \left| \frac{dP}{d\sigma_G^2} \right|_{\sigma_G^2=0} }.\] This threshold is the ratio between the marginal utility loss and the marginal penalty reduction at zero compression. Figure 5 visualises this trade-off: for small \(\lambda\), preserving the full causal component is optimal; for larger \(\lambda\), the penalty becomes sufficiently expensive that a noisy compressed representation can achieve a higher objective value. This is the Gaussian analogue of the lossless regime in Proposition 3.

9 NCE-CLUB Upper Bound and Critic Interpretation↩︎

To construct a tractable surrogate for the disentanglement objective \[J(g)=I(g(T);Y\mid Z)-\lambda I(g(T);Z),\] we require an upper bound on the penalty term \(I(g(T);Z)\).

9.0.0.1 General NCE-CLUB bound.

We use the NCE-CLUB estimator [13]. For a pair of random variables \((X_1,X_2)\), NCE-CLUB provides the upper bound \[I(X_1;X_2) \le I_{\mathrm{NCE\text{-}CLUB}}(X_1;X_2),\] where, for the optimal critic \(f^\star\), \[\begin{align} I_{\mathrm{NCE\text{-}CLUB}}(X_1;X_2) := \mathbb{E}_{(x_1,x_2^+)\sim p(x_1,x_2)} \bigl[f^\star(x_1,x_2^+)\bigr] - \mathbb{E}_{x_1\sim p(x_1),\,x_2^-\sim p(x_2)} \bigl[f^\star(x_1,x_2^-)\bigr]. \end{align}\] The first expectation is over joint, or positive, pairs; the second is over independently sampled, or negative, pairs.

9.0.0.2 Application to the penalty term.

Applying this construction to the penalty with \[X_1=g(T),\qquad X_2=Z,\] gives \[I(g(T);Z) \le I_{\mathrm{NCE\text{-}CLUB}}(g(T);Z).\] In practice, the optimal critic \(f^\star\) is unknown, so we use a learned critic \(d_\phi\) and obtain the plug-in estimator \[\begin{align} I_{\mathrm{NCE\text{-}CLUB}}(g(T);Z;d_\phi) &:= \mathbb{E}_{p(g(T),z)} \bigl[d_\phi(g(T),z)\bigr] \\ &\quad - \mathbb{E}_{p(g(T))p(z)} \bigl[d_\phi(g(T),z)\bigr]. \end{align}\] This is the quantity used as the penalty estimator in the main text.

9.0.0.3 Interpretation of the critic \(d_\phi\).

The optimal critic corresponds to the log-density ratio \[d^\star(g,z) = \log\frac{p(z\mid g)}{p(z)} = \log p(z\mid g)-\log p(z),\] which is intractable because the true conditional \(p(z\mid g)\) is unknown. We therefore approximate it with a learnable critic \(d_\phi(g,z)\), which can be interpreted as a variational estimate \[d_\phi(g_\theta(T),z) \approx \log q_\phi(z\mid g_\theta(T))-\log p(z),\] where \(q_\phi(z\mid g)\) is a learned conditional model.

The critic is trained to distinguish positive pairs \[(g,z)\sim p(g,z)\] from negative pairs \[(g,z)\sim p(g)p(z).\] This training aims to move \(d_\phi\) toward the optimal log-ratio and thereby tighten the plug-in estimate.

9.0.0.4 Adversarial training dynamics.

The resulting optimisation has a min–max form. The critic parameters \(\phi\) are trained to maximise the NCE-CLUB objective, improving the estimate of the treatment–confounder mutual information. The encoder parameters \(\theta\), through \(g_\theta\), are trained to minimise the same penalty term, thereby removing information about \(Z\) from the learned representation. Combined with the utility loss, this encourages representations that remain predictive of \(Y\) while suppressing confounder-dependent non-causal information.

10 Synthetic SCM: experimental details↩︎

We generate 1,000 training samples, and 1,000 validation samples from the following linear additive SCM with \(d = 10\), \(d_c = 5\): \[\begin{align} X &\sim \mathcal{N}(0, I_d), \\ T_C &= X_{0:d_c} + \epsilon_{T_C}, & \epsilon_{T_C} &\sim \mathcal{N}(0, I_{d_c}), \\ T_{nC} &= X_{d_c:d} + \epsilon_{T_{nC}}, & \epsilon_{T_{nC}} &\sim \mathcal{N}(0, I_{d-d_c}), \\ T &= [T_C, T_{nC}], \\ Y &= \mathbf{1}^\top T_C + \mathbf{1}^\top X + \epsilon_Y, & \epsilon_{Y,i} &\overset{\mathrm{iid}}{\sim} \mathcal{N}(0, \sigma_Y^2). \end{align}\]

The predictor is a linear model \(\hat{Y} = w_X^\top X + w_T^\top T + b\). We take the treatment representation to be the per-coordinate contribution of \(T\) to the prediction, \(h_T = T \odot w_T \in \mathbb{R}^{10}\), where \(w_T\) is the treatment slice of the predictor’s own weight vector. The MI critic therefore decorrelates \(X\) from the components of \(T\) that the predictor actually uses. Under gradient reversal this pushes \(w_T\) toward the causal coordinates of \(T\). The critic itself is bilinear: two bias-free linear projections \(W_g, W_x \in \mathbb{R}^{10 \times 10}\) to a shared embedding space, with InfoNCE scores given by the full \(B \times B\) matrix \(S = (W_g h_T)(W_x X)^\top\) over a minibatch of size \(B\). We train with lr \(3\times10^{-4}\), batch size \(64\), \(4{,}000\) epochs, and GRL schedule with \(\lambda = 0\) for epochs 0–200, cubic ramp \(0 \to \lambda_{\max} = 0.5\) over epochs 200–1000, held constant thereafter.

11 KuaiRand: experimental details↩︎

We use the KuaiRand-Pure [9] (Zenodo 10439422): \(27{,}285\) users, \(7{,}583\) items, \(\sim\)​1.4M algorithmic-exposure rows and \(\sim\)​1.2M random-exposure rows. We train on \(80\%\) of the algorithmic-exposure data, and use the remaining \(20\%\) as our validation set for early stopping. We calculate the AUC on the held-out random-exposure data. Our model is a DeepFM [43] backbone with \(16\)-dim user_id/video_id embeddings. A treatment encoder MLP \(d_t \to 32 \to 16\) ingests \(\sim\)​15 continuous video features (video_duration, one-hot video_type, \(11\) engagement-stat columns). Item representation \(g\): encoded video features (\(16\)) concatenated with video_id embedding (\(16\)), \(32\) dims total. User representation \(u\): \(18\) user-level features concatenated with user_id embedding (\(16\)), \(34\) dims total. DeepFM tower with DNN hidden sizes \([32, 16]\), dropout \(0.2\), FM second-order embed dim \(k = 8\). The CRL variant adds a bilinear critic between \(g\) and \(u\) via gradient reversal. The architecture and capacity are identical between the two configurations. With train with Adam, lr \(2 \times 10^{-3}\), weight decay \(10^{-6}\), and batch size \(4{,}096\). The GRL schedule is step-based: \(\lambda = 0\) for steps 0–400, linear ramp \(0 \to \lambda_{\max} = 0.3\) over steps 400–600, and held constant thereafter.

References↩︎

---

1. Equal contribution↩︎

2. Work done during internship at Spotify.↩︎