Towards Understanding Extrapolation: a Causal Lens

Authors: lingjing  kong
guangyi  chen
petar  stojanov
haoxuan  li
eric p xing
kun  zhang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 22:37:54

Abstract: Canonical work handling distribution shifts typically necessitates an entire target distribution that lands inside the training distribution.
However, practical scenarios often involve only a handful target samples, potentially lying outside the training support, which requires the capability of extrapolation.
In this work, we aim to provide a theoretical understanding of when extrapolation is possible and offer principled methods to achieve it without requiring an on-support target distribution.
To this end, we formulate the extrapolation problem with a latent-variable model that embodies the minimal change principle in causal mechanisms.
Under this formulation, we cast the extrapolation problem into a latent-variable identification problem.
We provide realistic conditions on shift properties and the estimation objectives that lead to identification even when only one off-support target sample is available, tackling the most challenging scenarios.
Our theory reveals the intricate interplay between the underlying manifold's smoothness and the shift properties.
We showcase how our theoretical results inform the design of practical adaptation algorithms. Through experiments on both synthetic and real-world data, we validate our theoretical findings and their practical implications.

Main paper: 
 introduction: 
 Extrapolation necessitates the capability of generalizing beyond the training distribution support, which is essential for the robust deployment of machine learning models in real-world scenarios. Specifically, given access to a source distribution D src := p(x src , y src ) with support X src := Supp(p src (x)) and one or a few out-of-support samples x tgt / ∈ X src , the goal of extrapolation is to predict the target label y tgt . For example, if the training distribution includes dog images, we aim to accurately classify dogs under unseen camera angles, lighting conditions, and backgrounds. While intuitive for humans, machine learning models can be brittle to minor distribution shifts [1] [2] [3] [4] .Addressing distribution shifts has garnered significant attention from the community. Unsupervised domain adaptation under covariate shifts addresses the shift of the marginal distribution p(x) across domains. However, canonical techniques such as importance sampling and re-weighting [5] [6] [7] [8] [9] are predicated on the assumptions of overlapping supports Supp(p tgt (x)) ⊂ Supp(p src (x)) and the availability of the entire target marginal distribution p tgt (x). Similarly, domain generalization [10] [11] [12] assumes access to multiple source distributions p src (x, y) whose supports jointly cover the target distribution. In addition to these methods, test-time adaptation (TTA) [13] [14] [15] [16] is particularly relevant to our discussion of extrapolation. TTA addresses out-of-distribution test samples at the individual sample level. Canonical methods include updating the source model with entropy-based or selfsupervised losses on target samples. However, most TTA research focuses on empirical aspects, with limited theoretical formalization [17] . Most related to our work, Kong et al. [18] and Li et al. 2) where s potentially changes all pixels in the images, such as the camera angle in the example. In this case, we can identify the invariant variable c under a moderate amount of shift until the shift becomes excessive. For instance, the back view of the cat in the figure could be confused with other animals. Figure (c ) illustrates the sparse shift condition (Theorem 4.4) where s influences a limited number of pixels, such as the background in the example. In contrast to the dense shift, we can identify c under the sparse shift regardless of its severity. In the figure, there is no ambiguity of the class "cow" even though the background has changed to the moon. [19] propose theoretical frameworks to characterize distribution shifts and explore conditions for identifying latent changing factors. However, these frameworks assume access to multiple source distributions with overlapping supports, which are not directly applicable to the extrapolation problem considered in this work, where we have potentially only one out-of-support target sample x tgt .Since the target y tgt can be arbitrarily out of the source support (Figure 1a ), extrapolation is illposed without proper assumptions on the relationship between the source and the target. In this work, we formulate a latent-variable model that encodes a minimal change principle to address this ill-posedness. Specifically, we assume that a latent variable z determines x such that x := g(z). The minimal change principle entails the following two assumptions on the generating process. 1) The out-of-support nature of x tgt stems from only a subspace of z, denoted as s, while the complement partition c for the target sample x tgt is within the training support, i.e., z := [c, s], s tgt / ∈ S src , and c tgt ∈ C src . 2) The changing variable s only controls non-semantic attributes in x and thus doesn't influence the label y, i.e., y := g y (c). This formulation attributes the seemingly complex shifts in the pixel space of x to simple intrinsic changes (in the sense that this change involves only s) in the latent space, allowing us to reason about the transfer via the invariant variable c. Under our formulation, extrapolation amounts to identifying invariant latent variables c, with which a model f : c → y trained on the labeled source dataset can be directly applied to the target sample.In light of this formulation, we investigate the identifiability conditions of the invariant variable c. We propose two sets of conditions addressing two regimes of the influence from s on x. We refer to the case where all dimensions x i (e.g., all pixels in an image) could be influenced by the changing variable s as the dense influence and the case where only a limited subset of dimensions x i is affected as the sparse influence. Specifically, our first condition (Dense Influence, Theorem 4.2) states that if s's influence is dense, then assuming that c takes on only finite values, extrapolation requires the manifold associated with each value of c (i.e., g(c, •) over s) to be adequately separable, and the target changing variable s tgt to be close to the source support S src . Intuitively, if images from two classes (two values of c) are sufficiently distinguishable, such as cats and dogs, we can still recognize the class of the target sample x tgt , even if it has undergone moderate unseen shifts on all pixels like camera angles and positions controlled by s (Figure 1b ). Our second condition (Sparse Influence, Theorem 4.4) states that if s's influence is sufficiently sparse, then extrapolation can occur regardless of the distance of s tgt to the source support S src . Intuitively, we can recognize the class "cow" even if only the background is changed, regardless of the severity (Figure 1c ).Our theory provides insights into the interaction between the underlying manifold's smoothness, the out-of-support distance, and the nature of the shift. We conduct synthetic data experiments to validate our theoretical results. Moreover, we discuss the relationship between our results and TTA approaches. In particular, we apply our theoretical insights to improve autoencoder-based MAE-TTT [20] and observe noticeable improvements. Additionally, we demonstrate that basic principles (sparsity constraints) from our framework can benefit the state-of-the-art TTA approach TeSLA [21] . Our empirical results not only show the practical viability of our theory but also pave the way for further advancements in the field.In summary, our contributions are threefold:• We formulate the extrapolation task as a latent-variable identification problem. Our latent-variable model encodes complex changes in observed variables to a partition of latent variables, allowing us to reason about transferability from the source to the target through latent variable identification.• We provide identification guarantees for the proposed latent-variable model, including shifts of distinct properties (dense vs. sparse) and corresponding conditions on the generating process. Our theory provides an essential understanding of when latent-variable identification is possible without accessing an entire target distribution and assuming overlapping supports as in prior work [18, 19] .• Inspired by our theory, we propose to add a likelihood maximization term to autoencoder-based MAE-TTT [20] to facilitate the alignment between the target sample and the source distribution. In addition, we propose sparsity constraints to enhance the state-of-the-art TTA approach TeSLA [21] . We validate our proposals with empirical evidence. 
related work: 
 Extrapolation. Out-of-distribution generalization has attracted significant attention in recent years. Unlike our work, the bulk of the work is devoted to generalizing to target distributions on the same support as the source distribution [22, 23, 8] . Recent work [24] [25] [26] [27] investigates extrapolation in the form of compositional generalization by resorting to structured generating functions (e.g., additive, slot-wise). Another line of work [28] [29] [30] studies extrapolation in regression problems and does not consider the latent representation. Saengkyongam et al. [31] leverage a latent variable model and linear relations between the interventional variable and the latent variable to handle extrapolation.In this work, we formulate extrapolation as a latent variable identification problem. Unlike the semi-parametric conditions in prior work, our conditions do not constrain the form of the generating function and are more compatible with deep learning models and tasks.Latent-variable identification for transfer learning. In the latent-variable model literature, one often assumes latent variables z generate the observed data x (e.g., images, text) through a generating function. However, the nonlinearity of deep learning models requires the generating function to be nonlinear, which has posed major technical difficulty in recovering the original latent variable [32] .To overcome this setback, a line of work [33] [34] [35] [36] assumes the availability of an auxiliary label u for each sample x and under different u values, each component z i of z experiences sufficiently large shift in its distribution. Since this framework assumes all latent components' distributions vary over distributions indexed by u, it does not assume the existence of some shared, invariant information across distributions, which is often the case for transfer learning tasks. To address this issue, recent work [18, 19] introduce a partition of z into an invariable variable c and an changing variable s (i.e., z := [c, s]) such that c's distribution remains constant over distributions. They show both c and s can be identified and one can directly utilize the invariant variable c for domain adaptation. However, their techniques crucially rely on the variability of the changing variable s, mandating the availability of multiple sufficiently disparate distributions (including the target) and their overlapping supports. These constraints make them unsuitable for the extrapolation problem. In comparison, our theoretical results give identification of the invariant variable c (the on-support variable in the extrapolation context) with only one source distribution p src (x) and as few as one out-off support target sample x tgt through mild assumptions on the generating function, directly tackling the extrapolation problem.Please refer to Section A1 for more related work. 
extrapolation and latent-variable identification: 
 Given the labeled source distribution p(x src , y src ), our goal is to make predictions on a target sample x tgt outside the source support (x tgt ̸ ∈ X src ). While more target samples would provide better information about the distribution shift, in practice, we often have only a handful of samples to work with. Therefore, we focus on the challenging scenario where only one target sample x tgt is available.Making reliable predictions on out-of-support samples x tgt is infeasible without additional structure.Real-world problems where humans successfully extrapolate often follow a minimal change principle: they involve sparse, non-semantic intrinsic shifts despite complex raw data changes. For example, a person who has only seen a cow on a pasture can recognize the same cow on a beach, even if the background pixels change significantly. Here, the cow corresponds to the part of the latent variable that remains within the support of the source data, which we call the invariant variable c (c tgt ∈ C src ), while the background change corresponds to the complement that drifts off the source support, which we call the changing variable s (s tgt / ∈ S src ). Clearly, extrapolation is impossible if the intrinsic shift is dense (i.e., all dimensions change, z = s) or semantic (i.e., y is a function of s). For instance, if the variable s also alters the cow's appearance drastically, making it unrecognizable, extrapolation fails. We define the data-generating process to encapsulate this minimal change principle, as follows:EQUATIONFigure 2 : The data-generating process. The invariant latent variable c and the changing latent variable s jointly generate the observed variable x. The dashed line indicates potential statistical dependence.In this process, the latent space z ∈ Z ⊂ R dz comprises two subspaces: the invariant variable c ∈ C ⊂ R dc and the changing variable s ∈ S ⊂ R ds . We define Z := Z src ∪ {z tgt } as the source support augmented with the target sample z tgt and similarly X and S. The invariant variable c encodes shared information between the source distribution p(x src ) and the out-of-support target sample x tgt , while the changing variable s describes the shift from the source support X src . Hence, c tgt ∈ C src and s tgt / ∈ S src . The variables z := [c, s] jointly generate the observed variable x ∈ X ⊂ R dx through an invertible generating function g : R dz → R dx . Furthermore, we assume that the label y originates from the invariant variable c. This assumption reflects the reality that factors such as camera angles and lighting do not affect the object's class in an image.Our latent-variable model adheres to the minimal change principle in two key ways: (1) the target sample's out-of-support nature arises from only a subset of latent variables s, and (2) these changing variables s are non-semantic, thus not influencing the label y.Extrapolation and identifiability. Under this framework, extrapolation is possible if we can identify the true invariant variable c in both the source distribution p src (x) and the target data x tgt . This allows us to learn a classifier f cls : c → y on the labeled source distribution p src (x, y). Since the target sample's invariant variable falls within the source support (c tgt ∈ C src ), this classifier f cls can be directly applied to the target sample c tgt . Thus, the task of extrapolation reduces to identifying the invariant variable c in both the source distribution p(x src ) and the target sample x tgt . In Section 4, we explore the conditions for identifying the invariant variable c.Given the above reasoning, we define identifiability in Definition 3.1 (i.e., block-wise identifiability [37, 24] ) which suffices for extrapolation. Definition 3.1 (Identifiability of the Invariant Variable c). For any x 1 and x 2 , their true invariant variables c 1 , c 2 are equal if and only if the estimates ĉ1 , ĉ2 are equal: c 1 = c 2 ⇐⇒ ĉ1 = ĉ2 . 
identification guarantees for extrapolation: 
 In this section, we provide two sets of conditions on which one can identify the invariant variable c and discuss the intuition and implications.As discussed in Section 3, we need to identify the target sample x tgt with source samples x src that share the same invariant variable values with the target sample, i.e., c src = c tgt . This enables us to obtain the label of x tgt by assigning the label of such x src . The shift between the source distribution p(x src ) and the target sample x tgt originates from the out-of-support nature of the changing variable s tgt , i.e., s tgt / ∈ S src , it is crucial to impose proper assumptions on the influence of s on x so that x retains sufficient footprints of c for identification beyond the source support.We denote the Jacobian matrix of the generating function g as J g (z) and x's dimensions under the influence of s as I s (z) := {i ∈ [d x ] : ∃j ∈ {d c + 1, . . . , d z }, s.t., [J g (z)] i,j ̸ = 0}. We note that the set I s (z) is a function of z, since the influenced dimensions may vary over z. Intuitively, if s influences x in a dense manner, i.e., large |I s (z)| for potentially all dimensions x i , there may not be dimensions of x serving as clear signatures of c, thereby increasing the difficulty of identify c. Additionally, the degree to which the changing variable s is out-of-support plays a critical role -the further the target changing variable-the further the target changing variable s tgt deviates from the source distribution support S src , the more severe and unpredictable the shift becomes, making it harder to retrieve c. In the following, we formalize conditions on the influence of s from these two perspectives, revealing an interesting trade-off and interaction between these factors.Notations. The true generating process involves c, s, distributions p, and g (Equation 1), we define their statistical estimates with ĉ, ŝ, p, and ĝ through the objectives we will introduce. 2 We assume that the estimation process respects the conditions of the corresponding true-generating process.We begin by investigating scenarios where there are no constraints on the number of dimensions of x (i.e., the number of pixels) influenced by the changing variable s, i.e., potentially large |I s (z)|, which we term as dense shifts. For images, these shifts encompass global transformations such as changes in camera angles and lighting conditions that could potentially affect all pixel values (Figure 1b ).Understanding the problem. As dense shifts could influence all the dimensions of x, every dimension could be out of the source support and there might not be dimensions of x that solely contain the information of c. Consequently, relying on any subset of x dimensions to infer the original c becomes untenable. For instance, consider a scenario where the source distribution contains frontalview images of a cat, while the target sample portrays the same cat from a side view (Figure 1b ). The model cannot recognize these two images as the same cat (the same c) by matching a specific part of the side view, say the cat's nose, to samples in the source distribution because this cat's nose only shows up as a front view and can be vastly different in terms of the pixel region and values. The model cannot match specific features such as the cat's nose, between the side-view target and the source distribution, as the pixel region and values for the nose drastically differ.Our approach. For the reasons above, we need to constrain such dense changes so that even when all dimensions are affected, the target sample adheres to some intrinsic structure determined by the underlying c tgt and remains distinguishable from samples of c ̸ = c tgt In many real-world distributions, we can interpret c as the embedding vector of classes or other categories, with each c value indexing a manifold g(c, •) over s. If manifolds are smooth and sufficiently separable from each other, they should exhibit limited variations in the adjacent region to the training support, avoiding confusion between distinct categories. For example, there exists a noticeable distinction between cats and lions, such that moderate illumination changes would not cause confusion until illumination significantly obscures distinguishing features. In the following, we formalize these structures by assuming a finite cardinality of c and constraining the distance of s tgt to the support S src .Additional notations. We denote with J u an upper bound of the Jacobian spectrum norm: ∥J g (z)∥ ≤ J u on the support. In Appendix A2, we show J u < ∞ due to Assumption 4.1-i and Assumption 4.1ii. We denote with D(c 1 , c 2 ) the ℓ 2 distance between two manifolds on the support boundary:D(c 1 , c 2 ) := inf s1,s2∈Bd(Ssrc) ∥g(c 1 , s 1 ) -g(c 2 , s 2 )∥, where we denote the boundary of source support with Bd(S src ). We denote with D(s, S src ) the minimal ℓ 2 distance between s and the source support S src , i.e., D(s, S src ) := inf ssrc∈Ssrc ∥s -s src ∥. .Discussion on the conditions. As discussed above, the main conditions revolve two key factors: the discrete structure of the invariant distribution of p(c) in Assumption 4.1-iii and the off-support distance of the changing variable s in Assumption 4.1-v. The discrete structure of p(c) is applicable to many real-world scenarios, especially classification tasks where the semantic invariant information often manifests as discrete class labels or other categorical distinctions. While this assumption is typically valid, it can be extended to encompass continuous dimensions in the invariant variablec := [c c , c d ]where c c and c d stand for the continuous and discrete dimensions, respectively. In such cases, we can group the continuous dimensions c c with the changing variable s and the same proof would give rise to the identification of the discrete part c d , which suffices for classification tasks. The off-support distance condition involves the smoothness of the generating function g, where a smoother generating function allows more leeway for the target changing variable s tgt to deviate. When s controls the camera angle, one may be able to recognize a slightly sided view of cats after seeing front views in the source until the s deviates too far and all images become back views, potentially leading to confusion with other animals (Figure 1b ). Assumption 4.1-i ensures that the generating process preserves the latent information, which is widely adopted in the literature [18, 19, 35, 36, 33, 38] . Specifically, this guarantees that manifolds indexed by distinct values of c are separate from each other, maintaining strictly positive distances between them. Assumption 4.1-ii,iv are technical conditions mirroring realities that pixels values are bounded and the changing variable s often represent attributes that vary gradually across its support (e.g., lighting and angles). Theorem 4.2 (Extrapolation under Dense Shifts). Assuming a generating process in Equation 1, we estimate the distribution with model (ĝ, p(ĉ), p(ŝ)) with the objective:EQUATIONUnder Assumption 4.1, the estimated model can attain the identifiability in Definition 3.1.Proof sketch. We estimate the generative process through maximum likelihood estimation on the source distribution p(x) = p(x). Under Assumption 4.1-i,ii,iii,iv, we can establish the identification of c on the support [38] , i.e., ĉ1 = ĉ2 ⇐⇒ c 1 = c 2 , for x 1 , x 2 ∈ X src . This implies that all samples x src on a given source manifold g(c, •) share identical values of ĉ. In the objective, we maximize the likelihood p(ĉ tgt ) to drive ĉtgt to match one of the discrete values of ĉ with nontrivial probability mass. Given the identification of c on the support, this equates to assigning x tgt to a manifold g(c * , •) with c * ∈ C. Our task now switches to ensuring that this is the correct manifold for x tgt , i.e., c * = c tgt . To accomplish this, we select the estimated model that uses the minimal offsupport distance on ŝ (i.e., D(ŝ, Ŝsrc )) to explain the off-support nature of x tgt . This also embodies the minimal change principle. This and Assumption 4.1-v guarantee that only the correct manifold (g(c tgt , •)) can effective capture x tgt , thereby facilitating the desired identification. In practice, these constraints can be implemented through Lagrange multipliers. Full proof is given in Appendix A2.We now examine cases where the changing variable s influences only a subset of dimensions of x, i.e., a limited |I s (z)|, which we refer to as sparse shifts. For image distributions, these shifts include local corruptions or background changes that do not alter foreground objects (Figure 1c ).Additional notations. We define the index set I c (z) under the influence of c and the indices under the the exclusive influence of c as I c\s (z) := I c (z) \ I s (z).Understanding the problem. In contrast to the dense-shift scenario, here we have a non-trivial subset of dimensions [x] I c\s (z) that are unaffected by the changing variable s. Consequently, if these dimensions carry sufficient information about c, we can exploit them to directly recover the true c, regardless of the distance [x] Is(z) deviates from the support. In contrast, in the dense-shift scenario, we need to constrain the out-of-support distance of s and assume the discreteness of c. Consider a scenario where a fixed c represents a specific cow and s controls only the background. Despite the variation in the target background (e.g., desert or space), we can effectively match the cow in the target image to the correct source images (see Figure 1c ). While this may seem intuitive for humans, it is nontrivial for machine learning models to automatically recognize the region [x] I c\s (z) , especially given its potential variation across z.Our approach. For image classification, [x] I c\s (z) corresponds to foreground objects (or a portion) unaffected under sparse changes induced by s (e.g., background changes). Humans can recognize this region because the pixels within it are strongly correlated (e.g., cow features). This observation motivates us to formalize such dependence structures in natural data to enable automatic identification. iv [Mechanistic Dependence]: For all z, any nontrivial partition P 1 , P 2 of the dimensions I c\s (z) yields dependence between the sub-matrices of the Jacobian J g (z): rank([J g (z)] I c\s (z)) < rank([J g (z)] P1 (z)) + rank([J g (z)] P2 (z)).Discussion on the conditions. Assumption 4.3-iii stipulates that the influence of s is sparse, either in terms of dimension counts or in its intersection with the influence from c. It is noteworthy that while the influence is sparse, its location can vary over images, as indicated by the dependence of I s on z.Consequently, it can capture diverse image corruptions and background changes. Assumption 4.3ii ensures that [x] I c\s is sufficiently informative about c. For instance, it precludes scenarios where a sparse corruption alters the top stroke of "7" to resemble "1", rendering the uncorrupted region fundamentally unidentifiable. Assumption 4.3-iv enforces the dependence alluded to in our previous discussion: the unaffected dimensions [x] I c\s exhibit mechanistic dependence across them, characterized by the Jacobian rank [39] . Thus, generating separate parts of an object necessitates more capacity than generating the entire object, as the dependence across the two parts can inform each other's generation. This inherent dependence enables the identification of the unaffected region. Theorem 4.4 (Extrapolation under Sparse Shifts). Assuming a generating process in Equation 1, we estimate the distribution with model (ĝ, p(ĉ), p(ŝ)) with the objective: It's worth noting that unlike the global shift case (Theorem 4.2), here we do not place a constraint on the out-of-support-ness of s tgt , a point we empirically verify in Section 6.3.EQUATIONGenerative adaptation. Our theoretical framework, inherently a generative model, can be implemented through auto-encoding over the source distribution and the target. Akin to our estimation framework, MAE-TTT [20] trains a masked auto-encoding model (f enc and f dec ) on the source distribution and adapts to target samples through the auto-encoding objective. Consequently, we have f dec (f enc (x)) ≈ x for x ∈ X src ∪ {x tgt }, which approximates the distribution-matching aspect of our estimation objectives Equation 2 and Equation 3.Despite the resemblance on the reconstruction objective, MAE-TTT does not explicitly perform the representation alignment as our objectives -a classifier f cls is only trained on the labeled source distribution, which takes in f enc 's output ẑ and produces logit values. In addition, our objectives entail maximizing the target likelihood p(ĉ tgt ) to align ĉtgt to the source support Ĉsrc . As large logit values indicate the sample is close to distribution modes [40, 41] and f enc is enforced invertible through auto-encoding, we can interpret minimizing the entropy of f cls (ẑ tgt ) as filtering ẑtgt to obtain ĉtgt and driving it towards the modes of p(ĉ). Therefore, we implement the entropy minimization loss inf -log y f cls (ẑ tgt ) y log(f cls (ẑ tgt ) y ) as a surrogate for maximizing p(ĉ). We show that this significantly boosts the performance of MAE-TTT in Section 6.1.Regularization. While our objectives simultaneously involve the source distribution and the target sample, the source distribution may not be accessible during adaptation. Aggressive updates on the target sample may distort the source information stored in the model and ultimately impair the performance. To address this, we propose to impose regularization on the source-pretrained backbone during adaptation to enforce minimal changes and preserve the source information. In Section 6.2, we instantiate this with low-rank updates and sparsity constraints, showcasing the resultant benefits. 
synthetic data experiments: 
 In this section, we conduct synthetic data experiments on classification to directly validate the theoretical results in Section 4. We present additional experiments on regression in Section A4.2.Experimental setup. We generated the synthetic data following the generative process in Equation 1, with d c = 4 and d s = 2. We focus on binary classification and sample class embeddings c 1 and c 2 from N (0, I c ) and N (2, I c ) respectively. We sample s src from a truncated Gaussian centered at the origin and sample s tgt at multiple distances from the origin. For the dense-shift case, we concatenate c and s and feed them to a well-conditioned 4-layer multi-layer perceptron (MLP) with ReLU activation to obtain x. For the sparse-shift case, we pass c to a 4-layer MLP to obtain a 4-d vector. We duplicate 2 dimensions of this vector and add s to it. The final x is the concatenation of the 4-d vector and the 2-d vector. We sample 10k points for the source distribution and 1 target sample for each run. We perform 50 runs for each configuration and compute the accuracy on the target samples. More details can be found in Appendix A4.Results and discussions. We compared our method with iMSDA [18] and a model trained only on source data. The results in both dense and sparse shift settings are summarized in Table 1 . Our method consistently outperforms both baseline methods (nearly random guesses) by a large margin on all sub-settings, validating our theoretical results. The results on iMSDA suggest that directly applying domain-adaptation methods to the extrapolation task may result in negative effects for lack of the target distribution in their training. 
real-world data experiments: 
 We provide real-world experiments to validate our theoretical insights for practical algorithms (Section 4.3) and theoretical results (Section 4.2). More results can be found in Appendix A5. As discussed in the first implication in Section 4.3, we incorporate an entropy-minimization loss to MAE-TTT and compare it with the original MAE-TTT.Experimental setup. We conduct experiments on ImageNet-C [45] and ImageNet100-C [46] with 15 different types of corruption. For the baseline, we utilize the publicly available code of MAE-TTT.In our approach, we do not directly integrate the entropy-minimization loss into the MAE-TTT framework. This is because the training process of self-supervised MAE relies on masked images, whereas entropy-minimization requires the classification of the entire image. To address this, we introduce additional training steps with unmasked images and apply the entropy-minimization loss during these steps. Specifically, the training process for each test-time iteration is split into two stages.We first follow the MAE-TTT approach by inputting masked images and training the model using reconstruction loss. In this stage, only the encoder is updated. Then, we input full images (32 in a batch) and optimize the model with the entropy minimization loss following SHOT [43] . In this stage, both the encoder and classifier are optimized. The learning rates for both stages are set the same. Comparison with baselines. In Table 3 , we compare our method with the baseline MAE-TTT [20] and other baselines therein. We can observe that our algorithm largely boosts the performance of the MAE-TTT baseline over most corruption types. This corroborates our theoretical insights and showcases its practical value. Understanding entropy-minimization steps. Table 4 presents the results of entropy-minimization with different training steps. The results indicate that the additional entropy-minimization steps significantly enhance the performance of the MAE-TTT framework, demonstrating the synergy between auto-encoding and entropy-minimization as indicated in our theoretical framework.As suggested by the second implication in Section 4.3, we integrate sparsity constraints into the stateof-the-art TTA method, TeSLA/TeSLA-s [21] . Although our theoretical results rely on a generative model, we demonstrate that our implications are also applicable to discriminative models.Experimental setup. We conduct experiments on the CIFAR10-C, CIFAR100-C, and ImageNet-C datasets [45] , following the protocols outlined for TeSLA and TeSLA-s [21] , with and without training data information. In the pre-train stage, we apply the ResNet50 [47] as the backbone network and follow prior work [14, 44] to pre-train it on the clean CIFAR10, CIFAR100, and ImageNet training sets, with joint contrastive and classification losses. In the test-time adaptation process, we adopt the sequential TTA protocol as outlined in TTAC [44] and TeSLA [21] . This protocol prohibits the change of training objectives throughout the test phase. To encourage sparsity, we add low-rank adaptation (LoRA) modules [48] to the backbone network, which limits the adaptation to low intrinsic dimensions. Beyond LoRA, we further implement a masking layer with corresponding sparsity constraint (ℓ 1 loss) to filter out redundant changes. More details can be found in Appendix A5.Results analysis. The average error rates under 15 corruption types for all CIFAR10-C, CIFAR100-C, and ImageNet-C datasets are summarized in Table 2 . We can observe that sparsity constraints consistently improve performance over the current SOTA method, TeSLA/TeSLA-s, across all three datasets. The lightweight nature of the sparsity constraint and its consistent performance enhancements make it a valuable addition. This demonstrates the potential of sparsity constraints as a versatile, plug-and-play module for enhancing existing TTA methods. To investigate the trade-off between the shift scope (dense vs. sparse) and severity, we simulate different levels of corruption severity and corrupted region sizes and evaluate a classical TTA method TENT [15] on these configurations. Following [45] , we inject impulse noise to the CIFAR10 dataset, with noise levels ranging from 1 to 10 to simulate various severity levels. To control the shift's scope, we crop regions of various sizes and introduce corruption only to this region. Figure 3 displays classification error curves under various shift severity levels and region sizes. We can observe that classification errors rise with increasing noise levels and region sizes. Notably, for large block sizes (dense shifts), the performance dramatically declines and even collapses as the severity level rises, whereas the performance remains almost constant over all severity levels in the sparse shift regime, verifying the theoretical conditions for Theorem 4.2 and Theorem 4.4. 
conclusion and limitations: 
 In this work, we characterize extrapolation with a latent-variable model that encodes a minimal change principle. Within this framework, we establish clear conditions under which extrapolation becomes not only feasible but also guaranteed, even for complex nonlinear models in deep learning.Our conditions reveal the intricate interplay among the generating function's smoothness, the out-ofsupport degree, and the influence of the shift. These theoretical results provide valuable implications for the design of practical test time adaptation methods, which we validate empirically.Limitations: On the theory aspect, the Jacobian norm utilized in Theorem 4.2 only considers the global smoothness of the generating function and thus may be too stringent if the function is much more well-behaved/smooth over the extrapolation region of concern. Therefore, one may consider a refined local condition to relax this condition. On the empirical side, our theoretical framework entails learning an explicit representation space. Existing methods without such a structure may still benefit from our framework but to a lesser extent. Also, our framework involves several loss terms including reconstruction, classification, and the likelihood of the target invariant variable. A careful re-weighting of these terms may be needed during training. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We disclose our compute resources in Appendix A4 and Appendix A5.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work is on the understanding the fundamental aspect of machine learning, posing no direct societal impacts.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification:Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite all the employed codebase in Appendix A4 and Appendix A5.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: "Towards Understanding Extrapolation: a Causal Lens " In this section, we discuss some related topics including extrapolation, latent-variable identification, and test-time adaptation.Extrapolation. Out-of-distribution generalization has attracted significant attention in recent years. Unlike our work, the bulk of the work is devoted to generalizing to target distributions on the same support as the source distribution [22, 23, 8] . Recent work [24] [25] [26] [27] investigates extrapolation in the form of compositional generalization by resorting to structured generating functions (e.g., additive, slot-wise). Another line of work [28] [29] [30] studies extrapolation in regression problems and does not consider the latent representation. Saengkyongam et al. [31] leverage a latent variable model and assumes a linear relation between the intervention variable and the latent variable to handle extrapolation. In this work, we formulate extrapolation as a latent variable identification problem. Unlike the semi-parametric conditions in prior work, our conditions do not constrain the form of the generating function and are more compatible with deep learning models and tasks. We demonstrate that our conditions naturally lead to implications benefiting practical deep-learning algorithms.Latent-variable identification for transfer learning. Identifying latent variables in a causal model has become one canonical paradigm to formalize and understand representation learning in the deep learning regime. Typically, one would assume some latent variables z generate the observed data x (e.g., images, text) through a generating function. However, the nonlinearity of deep learning models requires the generating function to be nonlinear, which has posed major technical difficulty in recovering the original latent variable [32] . To overcome this setback, a line of work [33] [34] [35] [36] assumes the availability of an auxiliary label u for each sample x and under different u values, each component z i of z experiences sufficiently large shift in its distribution. This condition leads to componentwise identification of z, i.e., each estimate ẑi is equivalent to z π(i) up to an invertible mapping for a permutation function π :[d z ] → [d z ].Since this framework assumes all latent components' distributions vary over distributions indexed by u, it doesn't assume the existence of some shared, invariant information across distributions, which is often the case for transfer learning tasks. To address this issue, recent work [18, 19] introduce a partition of z into an invariable variable c and an changing variable s (i.e., z := [c, s]) such that c's distribution remains constant over distributions. They show both c and s can be identified and one can directly utilize the invariant variable c for domain adaptation. However, their techniques crucially rely on the variability of the changing variable s, mandating the availability of multiple sufficiently disparate distributions (including the target) and their overlapping supports. These constraints make them unsuitable for the extrapolation problem.In comparison, our theoretical results give identification of the invariant variable c (the on-support variable in the extrapolation context) with only one source distribution p src (x) and as few as one out-off support target sample x tgt through mild assumptions on the generating function, which directly tackles the extrapolation problem.Test-time adaptation. Test-time Adaptation (TTA) aims at adapting models trained on a source domain to align with the target domain during testing [49] [50] [51] [52] [53] [54] [55] . It is broadly classified based on whether the training objective is modified. Test-time Training (TTT) methods [13, 14, 44, 56, 57] , including TTT [13] and TTT++ [14] , proficiently adjust models to target domains by implementing similar self-supervised learning strategies on both training and testing data. In contrast, Sequential Test-Time Adaptation [15, 54, 55, [58] [59] [60] [61] [62] [63] (sTTA) garners significant interest due to its practicality, notably its one-pass sequential inference and no training objective access. Research in sTTA primarily concentrates on two facets: the selection of model parameters for adaptation and the refinement of pseudo-labeling techniques for enhanced efficiency. For instance, TENT [15] fine-tunes the Batch Normalization (BN) layers by minimizing entropy, SHOT [16] adjusts the backbone network while maintaining a static classifier, and T3A [64] updates the classifier prototype. Moreover, a burgeoning line of research [65, 21, 44, 50, 15, 16] focuses on deriving more robust self-training signals through improved pseudo labeling strategies. For example, TTAC [44] employs clustering techniques to extract more accurate pseudo labels. Despite the prominent recent development, these algorithms tend to be brittle and sensitive to hyper-parameter tuning [66] and limited in theoretical understanding [17] .Our work offers formalization and understanding to fill in this gap. We show that insights inferred from our theory can indeed benefit existing TTA algorithms, which hopefully will serve as the first step to bridge the theory and practice for TTA algorithms.A2 Proof for Theorem 4.2 .We first present Lemma A1 from Kong et al. [38] which establishes the discrete information on the source support and serves as the starting point in the proof of Theorem 4.2.Lemma A1 (Source discrete subspace identification [38] ). Assuming a generating process in Equation 1, we estimate the distribution with model (ĝ, p(ĉ), p(ŝ)). Under Assumption 4.1 i,ii,iii,iv, it follows that the estimated variable ĉ takes on values from {ĉ k } K k=1 where each value corresponds uniquely to one value of the true variable c, i.e., c = c k ⇐⇒ ĉ = ĉk . (2)Under Assumption 4.1, the estimated model can attain the identifiability in Definition 3.1.Proof for Theorem 4.2. Lemma A1 shows that the discrete invariant variable c is identifiable on the source distribution.In the following, we show that the target's invariant variable c tgt is identifiable if s tgt does not drift too far away from the source support S src . Suppose that x tgt resides on both manifolds g(c k , •)and g ′ (c k ′ , •)where k ̸ = k ′ . The generating function g ′ ∈ G belongs to the generating function class and behaves exactly the same as g on the source support, i.e., g ′ = g over C × S src . We define the minimal distance D(c k , c k ′ ) between the two manifolds on support boundaries, i.e.,D(c k , c k ′ ) := min s1,s2∈Bd(Ssrc) ∥g(c k , s 1 ) -g(c k ′ , s 2 )∥ > 0. Since x tgt lives on both manifolds g(c k , •) and g ′ (c k ′ , •), we can express it as x tgt = g(c k , s tgt ) = g ′ (c k ′ , s ′ tgt ).We define s src ∈ arg min s∈Ssrc ∥s -s tgt ∥ and s ′ src ∈ arg min s∈Ssrc ss ′ tgt as two closest points on the source support to s tgt and s ′ tgt respectively. It follows thatEQUATIONwhere h := s tgt -s src and h ′ := s ′ tgt -s ′ src .It follows from Equation 4g(c k ′ , s ′ src ) -g(c k , s src ) = 1 0 J g(c k ,•) (s src + t • h)dt h - 1 0 J g(c k ′ ,•) (s ′ src + t • h ′ )dt h ′ ; =⇒ 1 0 J g(c k ,•) (s src + t • h)dt h - 1 0 J g(c k ′ ,•) (s ′ src + t • h ′ )dt h ′ ≥ D(c k , c k ′ ); =⇒ 1 0 J g(c k ,•) (s src + t • h)dt h + 1 0 J g(c k ′ ,•) (s ′ src + t • h ′ )dt h ′ ≥ D(c k , c k ′ ); =⇒ J u (∥h∥ + ∥h ′ ∥) ≥ D(c k , c k ′ ); =⇒ max{∥h∥ , ∥h ′ ∥} ≥ D(c k , c k ′ ) 2J u .(EQUATIONfor the true generating function g. Therefore, x tgt can only be explained by one manifold, which we denote as g(c tgt , •).Finally, we show that the objective Equation 2 guarantees that the solution ĉtgt corresponds to the true c tgt . We suppose that c tgt = c k which corresponds to ĉk for a specific k ∈ [K]. First, we note that ĉtgt could only take on values from {ĉ k } k∈[K] due to the naturalness constraint sup p(ĉ d ). Also, the correct solution ĉk is always a feasible solution to the objective Equation 2, since ĝ can take on the true generating function g. We can apply the contradiction argument in the second step. That is, if there exists another plausible solution ĉk ′ in addition to the true solution ĉk , we would haveEQUATIONwhere the definitions are analogous to those in step two and decorated with • to indicate the difference. As noted above, the correct solution ĉk is upper-bounded by ĥ< D(ĉ k ,ĉ k ′ ) 2 Ĵu, since this is attainable when the estimated generating function is the true function, i.e., g = ĝ. Thus, the false solution ĉk ′ would always yield ĥ′ ≥ D(ĉ k ,ĉ k ′ ) 2 Ĵu > ĥ , which would be excluded by the distance-minimizing term min ĝ,ŝsrc∈ Ŝsrc ∥ŝ tgt -ŝsrc ∥. Therefore, we have shown that the estimated ĉtgt would correspond to the correct c k .A3 Proof for Theorem 4.4 iv [Mechanistic Dependence]: For all z, any nontrivial partition P 1 , P 2 of the dimensions I c\s (z) yields dependence between the sub-matrices of the Jacobian J g (z): rank([J g (z)]I c\s (z)) < rank([J g (z)] P1 (z)) + rank([J g (z)] P2 (z)).Theorem 4.4 (Extrapolation under Sparse Shifts). Assuming a generating process in Equation 1, we estimate the distribution with model (ĝ, p(ĉ), p(ŝ)) with the objective:EQUATIONUnder Assumption 4.3, the estimated model can attain the identifiability in Definition 3.1. Lemma A1 (Brady et al. [39] ). Let g, ĝ : R dz → R dx be smooth and invertible. Then, for anyz ∈ R z , S ⊂ R ds , rank([J g (z)] S ) = rank([J ĝ (ẑ)] S ), where ẑ := ĝ-1 • g(z).Proof. The proof consists of two steps. In the first step, we show the identification of the index set I c\s (z) ⊂ [d x ] over which x receives only c' influence. That is, ĝ maps the estimated invariant variable ĉ to x dimensions generated by the true invariant variable c alone. In the second step, we show that the objective sup p(ĉ tgt ) assigns to the estimated invariant variable for the target sample ĉtgt the source-distribution estimated value ĉsrc that is also generated by c tgt . Recall the notation z := [c, s]. We denote the latent source support as Z src and the set augmented with the target sample as Z := Z src ∪ {z tgt }Step 1. We first show that ĝ cannot map the estimated invariant variable ĉ to x dimensions generated by both the invariant variable c and the changing variable s.We show this by contradiction. Suppose that ∃z * ∈ Z, such thatEQUATIONWe partition I c\s (z * ) into two disjoint sets I c\s,1 := {i ∈ I c\s (z * )|i ∈ Îc\s (ẑ * )} and I c\s,2 := {i ∈ I c\s (z * )|i / ∈ Îc\s (ẑ * )} base on the overlap with Îc\s (ẑ * ). By the supposed condition Equation 7, it follows that I c\s,1 ̸ = ∅ and I s,1 ̸ = ∅.We now show that I c\s,2 is nonempty by contradiction. Suppose that I c\s,2 is empty. It follows that I c\s = I c\s,2 ⊂ Îc\s . By definition, we know that I s,1 ⊂ Îc\s . Thus, A := I c\s ∪ I s,1 ⊂ Îc\s . This inclusion implies that rank([J g (z * )] A,:EQUATIONSince I c\s ∩ I s,1 = ∅, [J g (z * )] I c\s ,dc+1:dz = 0, and each row of [J g (z * )] Is,dc+1:dz is nonzero by definition, we have that rank([J g (z * )] A,:EQUATIONwhich contradicts Equation 8. Therefore, I c\s,2 is nonempty. Since we have rank([J g (z * )] I c\s ) = d c and (I c\s,1 , I c\s,2 ) forms a pair of nonempty partition, Assumption 4.3-iv implies that rank([J g (z * )] I c\s,1 ,:EQUATIONAs the g's and ĝ's Jacobian matrix ranks are related (Lemma A1), Equation 10 implies that rank([J ĝ (ẑ * )] I c\s,1 ,:EQUATIONSince I c\s,1 ⊂ Îc\s,1 and I c\s,2 ∩ Îc\s,1 = ∅ by definition and [J ĝ (ẑ * )] I c\s,2 ,dc+1:dz has a full-row rank (Assumption 4.3-iii), it follows that rank([J ĝ (ẑ * )] I c\s ,:EQUATIONHowever, Lemma A1 implies that rank([J ĝ (ẑ * )] I c\s ,:EQUATIONThus, we have arrived at a contradiction. We have shown that ĝ maps the estimated invariant variable ĉ to x dimensions generated by the true invariant variable c alone, i.e., Îc\s (ẑ) ⊂ I c\s (z). Moreover, if an index i ′ from the ŝ's region Îs (ẑ) also belonged to I c\s (z) , i.e., i ′ ∈ Îs (ẑ) ∩ I c\s (z), then we would have rank([J g (z)] I c\s (z),: ) = rank([J ĝ (ẑ)] I c\s (z),: ) > rank([J ĝ (ẑ)] Îc\s (ẑ),: ) = d c . Thus, it follows that we can identify the index set exclusively influenced by c over z ∈ Z:EQUATIONStep 2. By definition, each value of ĉ determines the region [x] Îc\s (ĉ,ŝ) . Due to Equation 14, we have [x] Îc\s (ĉ,ŝ) = [x] I c\s (c,s) , where c can only take on a unique value according to Assumption 4.3-ii. Therefore, there exists an one-to-one mapping h ĉ from ĉ to c over Z. Further, since the maximum likelihood estimation is performed over the entire source distribution p(x) (Equation 3), the image of h ĉ equal to the entire C, thus is onto. Thus, we have shown that there exists an invertible mapping h ĉ : ĉ → c valid over Z (including the target sample), resulting from our objective (Equation 3).We employ a variational auto-encoder [67] whose encoder and decoder are both 4-layer MLP with 32 dimensions and leaky ReLu (α = 0.2). Following Equation 2and Equation 3, we implement reconstruction loss, KL loss on the source distribution, the likelihood loss on the target sample, and a classification loss on the source data. For the dense case, we implement an additional distance loss to minimize the ℓ 2 distance of ŝtgt to the center of the source support (which is the origin in our case). The source-only baseline is trained only with classification loss. The iMSDA implementation is adopted directly from the source code of Kong et al. [18] . We train all methods with Adam [68] and learning rate 2e -3 for 25 epochs. We fix the loss weights λ cls = 1, λ recons = 0.1, λ tgt_likelihood = 0.1, and λ s_distance = 0.01 (for dense shifts) overall distance configurations. We only tune λ KL from the interval {1e -1, 1e -2, 1e -3}. We run synthetic data experiments on one Nvidia L40 GPU and each run consumes less than 2 minutes.In addition to the classification experiments, we evaluate our model on regression in this section.Data generation. The regression target y is generated from a uniform distribution U (0, 4). We sample 4 latent invariant variables c from a normal distribution N (y, I c ). Two changing variables in the source domain s src are sampled from a truncated Gaussian centered at the origin. In the target domain, changing variables s tgt are sampled at multiple distances (e.g., {18, 24, 36}) from the origin. For dense shifts, observations x are generated by concatenating c and s and feeding them to a 4-layer MLP with ReLU activation. For sparse shifts, only two out of six dimensions of x are influenced by the changing variable s. We generate 10k samples for training and 50 target samples for testing (one target sample accessed per run).Model. We make two modifications on the classification model in Section 5. First, we substitute the classification head with a regression head (the last linear layer). Second, we replace the cross-entropy loss with MSE loss. We fix the loss weights of MSE loss and KL loss at 0.1 and 0.01 for all settings, respectively, and keep all other hyper-parameters the same as in the classification task. We use MSE as the evaluation metric.Table A1 displays the evaluation results. We can observe that our method consistently outperforms the baseline and maintains its performance over a wide range of shift distances. In contrast, the baseline that directly uses all the feature dimensions degrades drastically when the shift becomes severe. This indicates that our approach can indeed identify the invariant part of the latent representation, validating our theoretical results. When applying entropy minimization in the MAE-TTT framework [20] , we did not directly integrate the entropy-minimization loss. The self-supervised MAE training process relies on masked images, whereas entropy minimization requires classifying the entire image. To address this, we introduced additional training steps using unmasked images and applied the entropy-minimization loss during these steps. Specifically, the training process for each test-time iteration is split into two stages: 1) Stage One: We follow the MAE-TTT approach by inputting masked images and training the model using reconstruction loss. In this stage, only the encoder is updated. 2) Stage Two: We input full images (32 in a batch) and optimize the model with the entropy minimization loss following SHOT [43] . In this stage, both the encoder and classifier are optimized. The learning rates for both stages are set to be the same. The experiments are conducted with the PyTorch 1.11.0 framework, CUDA 12.0 with 4 NVIDIA A100 GPUs.Here, we provide the implementation details of our modification to add sparsity regularization. In the pre-train stage, we apply the ResNet50 [47] as the backbone network and follow [14, 44] to pre-train it on the clean CIFAR10, CIFAR100, and ImageNet training sets, with joint contrastive and classification losses. In the test-time adaptation process, we adopt the sequential TTA protocol as outlined in TTAC [44] and TeSLA [21] . This protocol prohibits the change of training objectives throughout the test phase. Moreover, all testing data be processed in a sequential manner (one-pass), ensuring each data point is passed through the adaptation process exactly once.Our method is built upon TeSLA [21] , and follows most of its hyperparameters. Thus, we only discuss the extra hyperparameters we involved, including the low-rank dimension r, the ratio of learning rate factor for soft frozen ratio l = lr lora lr backbone , and the ratio of sparsity loss ratio s . The details are shown in Table A2 . It was observed that the constraints on minimal change need to be more stringent as the complexity of the data increases. The experiments are conducted with the PyTorch 1.13.0 framework, CUDA 11.7 with an NVIDIA A100 GPU.We assess the impact of the recoverability of the invariant variable on Test-Time Adaptation (TTA) methods. To do this, we compare the performance of TTA methods using supervised and unsupervised pre-trained models that have similar ImageNet classification accuracy. Our goal is to validate whether For this detailed analysis, we employ ResNet50 models pre-trained in both supervised and selfsupervised manners, such as MoCo [70] . For the MoCo model, we apply the linear probe to fit the classifier. For a fair comparison, we select checkpoints from both pre-training methods that have similar ImageNet accuracies, approximately 69.7%. We follow the open-source TTA Benchmark [71] to evaluate both models using different downstream TTA methods. ImageNet-C is used as the evaluation dataset, and all hyperparameters are set to their default values.The performance of the different pre-trained models is summarized in Table A3 . Methods using supervised pre-training outperform those using unsupervised pre-training by a significant margin, indicating that the invariant variables learned from annotated labels play a crucial role in enhancing test-time adaptation.In Table A4 , we provide the detailed performance with corruption-wise classification error rates on all CIFAR10-C, CIFAR100-C, and ImageNet-C datasets. Specifically, we report results under seed 0 on all 15 testing corruptions including Gaussian Noise, Shot Noise, Impulse Noise, Defocus Blur, Glass Blur, Motion Blur, Zoom Blur, Snow, Frost, Fog, Brightness, Contrast, Elastic Transformation, Pixelate, and JPEG Compression. • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification:Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).We slightly abuse the notation p to denote density functions for continuous variables or delta functions for discrete variables.The code is provided here. 


Decision: After reviewing the paper, the feedback from the reviewers, and the authors' thorough rebuttals, I recommend accepting the paper. The authors have effectively addressed the concerns raised, including those from Reviewer 5kN3 (I read both the review and rebuttal). The work tackles an important problem and presents a solid theoretical foundation with empirical validation. Given these strengths, I believe the paper merits acceptance.

Reviews: review: paper_summary: paper_summary: * With regard to extrapolability in classification problems, they introduce a reasonable assumption that the latent factor generating the distribution shift only affects the input x but does not affect the label y is introduced and the identifiability of the latent factor is proved under additional reasonable assumptions.
* The interaction between the smoothness of the generating function, the distance out of support, and the nature of the shift (is the shift restricted to some pixels of the image?) is clarified.
* They discuss the relationship with test-time adaptation, extend using sparsity constraint, and verify its empirical performance.
strength_weakness: strength_weakness: * Methodological improvement and its empirical superiority to the existing test-time adaptation itself is marginal.
questions: questions: Q1. Is the reported error bar in Table 2 standard error or standard deviation?
Q2. Does it matter if the marginal distribution of the label $p(y)$ changes in the test phase?
limitations: limitations: Most of the important results are for the image classification problem, where the class-conditioned distribution p(x|y) is separated for each y and x has much information, and thus only unchanged pixels have enough information of $y$ under the sparse change assumption. It seems not to be simply extended to numerical predictions.
paper_summary: paper_summary: The authors discuss the problem of domain adaptation or distribution shift in the case when only a single point in the shifted domain is available rather than the full distribution. The authors approach such a problem from the perspective of a latent variable model, which assumes that the observations are generated from two latent variables: one is invariant under domain shifts (and carries informations about the classification label) and the other one changes under domain shifts but is irrelevant for classification. However, the generator function mixes these two latent variables making the classification under domain shifts difficult. The authors theoretically discuss the identifiability of the invariant latent variables in two scenarios, when the irrelevant for classification latent variables affects all dimensions in the input or only a subset of them. These theoretical results are validated in experiments on synthetic and real-world data.
strength_weakness: strength_weakness: - I am not sure I completely understood the connection between the theory and the actual learning objective. Specifically, I would appreciate if the authors could elaborate a bit more on how the Eq. (2) and (3) are connected to the learning objective in Sec. 3.3 (also see the Questions part of the review)
questions: questions: - It is very interesting that you derive a bound of how far the target point can be from the source manifold in Assumption 3.1-v. I wonder how tight do you think this bound is?
- Could you provide a bit more intuition why Assumption 3.3.-iv allows us to "identify the unaffected dimension indices [...] with our estimated model", I am not sure I understood this point?
- Did I understand correctly that in practice the estimation algorithms (as described in Sec. 3.3.) reduce to training an auto-encoder to match the data distribution, and a classifier in the latent space? So the key idea is that "raising" the classification to the latent space allows us to separate invariant latents from the nuisance variables?
limitations: limitations: The limitations are adequately addressed
paper_summary: paper_summary: The paper addresses the problem of out-of-support extrapolation and presents identification results within a latent variable framework. This framework resembles invariant learning, where one subset of latent variables directly causes the labels, while another subset, termed style latents, undergoes a distribution shift at test time. The authors provide identification results for scenarios where this distribution shift affects the entire image (global shift case) and where it affects only specific regions of the image (local shift case). Additionally, they connect their identification results to the entropy maximization objective used in test time adaptation (TTA) algorithms. These theoretical aspects are verified through experiments on popular TTA benchmarks, such as CIFAR-C and ImageNet-C.
strength_weakness: strength_weakness: * My main concern is with the technical soundness of the paper, particularly regarding certain claims and connections. I believe the connection between the identification results and the test-time adaptation (TTA) objective is problematic.  Specifically, Theorems 3.2 and 3.4 require that the learned distribution of the data ($\hat{p}(x)$) matches the true distribution ($p(x)$). However, the TTA methods are not generative models, and therefore they do not estimate the distribution of the data. However, TTA methods are not generative models and do not estimate data distributions, so they cannot enforce the constraint $p(x) = \hat{p}(x)$.  Therefore, it is unclear how the identification results provide insights into TTA methods. Identification concerns the existence of a unique solution for perfect estimation under the learning objective, but the learning objectives for the identification analysis and TTA methods are different. While the authors discuss this in Section 3.3, however, the issue is deeper,  they do not address the density estimation constraint and its connection to test-time adaptation comprehensively.

* The definition of identification considered by the authors (Definition 2.1) is problematic because it does not relate the estimated latents ($\hat{c}$) to the true latents ($c$). This implies that the estimated latents could be an arbitrary function of the true latents, which is concerning as it makes it unclear why such identifiability would be desirable. The identification criteria state that the estimated content latents are the same whenever the true content latent variables are the same. However, this does not necessarily imply a meaningful relationship between the estimated and true latents. For instance, block identification [1] ensures that the estimated content latents are a function of only the true content latents, establishing a clear relationship. If we cannot relate the estimated latents to the true latents, it is unclear whether the inferred content latents ($\hat{c}$) are genuinely informative about the true content latents ($c$). This lack of a meaningful connection undermines the practical value of the identifiability results.

References:

[1] Lachapelle, Sébastien, Divyat Mahajan, Ioannis Mitliagkas, and Simon Lacoste-Julien. "Additive decoders for latent variables identification and cartesian-product extrapolation." Advances in Neural Information Processing Systems 36 (2024).

[2] Khemakhem, Ilyes, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. "Variational autoencoders and nonlinear ica: A unifying framework." In International conference on artificial intelligence and statistics, pp. 2207-2217. PMLR, 2020.

[3] Zimmermann, Roland S., Yash Sharma, Steffen Schneider, Matthias Bethge, and Wieland Brendel. "Contrastive learning inverts the data generating process." In International Conference on Machine Learning, pp. 12979-12990. PMLR, 2021.
questions: questions: Please refer to the weaknesses section in my review above for more details regarding my concerns and queries.

* Why does the identification analysis apply for TTA when its learning objective is different from the learning objective considered in the identification analysis?

* How do the identification criteria considered in this study relate to standard identification criteria used in causal representation learning/ non-linear ICA [1, 2, 3]? Does satisfying the proposed criteria imply we learn meaningful content variables such that they are related to the true content variables?

* Following the results in Table 2 regarding the effect of adding sparsity regularization, it is unclear how TeSLA + SC provides an improvement over TeSLA, as the difference in mean performance between the two approaches is within the standard error. The authors should consider revising line 357 to accurately reflect this observation.
limitations: limitations: Yes, the authors have properly addressed the limitations and impact of their work.
paper_summary: paper_summary: This work addresses the challenge of extrapolation in scenarios where only a few target samples are available. It aims to provide a theoretical understanding and methods for effective extrapolation without needing a target distribution within the training support. The approach involves a latent-variable model based on the minimal change principle in causal mechanisms. The study identifies conditions under which identification is possible even with a single off-support target sample. Experiments with synthetic and real-world data validate the theoretical and practical aspects of the findings.
strength_weakness: strength_weakness: 1. Table 2 appears to omit results from TeSLA-s, which outperforms their TeSLA+SC. Thus their claims of SOTA are in question. Table 1 of:
     * Devavrat Tomar, Guillaume Vray, Behzad Bozorgtabar, and Jean-Philippe Thiran. Tesla: Test-time self-learning with automatic adversarial augmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20341–20350, 2023.

2. The paper lacks expected sections. There is no related work, making it difficult to contextualize their work within the field. Additionally, some of the introduction could be better placed in a background section, specifically lines 41-67.

3. The authors do not provide any source code.

4. More elaboration on the limitations would enhance the paper. The authors seem to have overlooked reflecting sufficiently on potential limitations of their method. For instance, since everything operates in latent space, the method would only be applicable to models that possess a latent representation.

5. Minor typos:
     * What do the numbers in Table 1 represent? Accuracy? AUROC?
     * Line 127: Should it not be, "... values c as x_{src}"?
     * Line 127: Missing period.
     * Line 321: It would be helpful to provide a link to the code for MAE-TTT.
     * Line 356: Missing period.
questions: questions: How would you apply your method for regression problems? Do you think it would perform well?
limitations: limitations: See weakness 4.

Comments: title: Thanks a lot for your recognition and valuable suggestions
More results on the regression setting
More results on the regression setting
comment: Dear Reviewer zpbG,

We really appreciate your recognition of our work and your kind words, and we are happy to hear that your concerns have been addressed. Again, thank you for your valuable suggestions which have undoubtedly contributed to improving the quality of our paper.

Many thanks,

The authors of #7667
Thank you for your response. My concerns have been addressed. I have already scored high enough and will maintain my score.
Dear reviewer 5kN3,

Since the discussion period will end in a few hours, we will be online waiting for your feedback on our rebuttal, which we believe has fully addressed your concerns. 

We would highly appreciate it if you could take into account our response when updating the rating and having discussions with AC and other reviewers.

Thank you so much for your time and efforts. Sorry for our repetitive messages, but we're eager to ensure everything is addressed.

Authors of # 7667
Dear Reviewer zpbG,

We were wondering whether your technical had been properly addressed by our responses so far. Please let us know if you have further questions or concerns that we can address. Thank you for your engagement with our work.

Many thanks,

Authors of 7667
​​Once again, we are grateful for your time and effort for reviewing our paper. Since the discussion period will end in around a day, we are very eager to get your feedback on our response. We understand that you are very busy, but we would highly appreciate it if you could take into account our response when updating the rating and having a discussion with AC and other reviewers.

Thanks for your time,

Authors of # 7667
Thank you so much for your encouraging words and we are truly grateful for your dedicated time and constructive comments!
Thank you very much for a thorough rebuttal! I confirm my positive view on this paper and happy to increase my score.
Dear Reviewer 5kN3,

As the discussion deadline approaches, we are wondering whether our responses have properly addressed your concerns? Your feedback would be extremely helpful to us. If you have further comments or questions, we hope for the opportunity to respond to them.

Many thanks,

7667 Authors
> Proof on Definition 2.

Thank you for your positive feedback and we are happy to see this major concern resolved!

> Standard deviations for TeSLA.

Thanks for the question! We’ve included in our revision standard deviations for TeSLA performances as follows. 
| Method | CIFAR10-C   | CIFAR100-C  | ImageNet-C |
|--------|-------------|-------------|------------|
| TeSLA  | 12.5 ± 0.04 | 38.2 ± 0.03 | 55.0 ± 0.17 |
Just as you can see, the standard deviations are all relatively small, so they should not affect our conclusion.

---

Thank you for your thoughtful comments and remarks and we’d highly appreciate your further feedback!
Thank you for your detailed feedback and we are delighted that we have cleared your previous concerns!

> Writing suggestions.

We highly appreciate your suggestion on the writing. Thanks to your suggestion, we have made the following modifications:

1. We have replaced current lines 70-72 with: “In particular, we apply our theoretical insights to improve autoencoder-based MAE-TTT [24] and observe noticeable improvements on TTA tasks. We also demonstrate that basic principles (sparsity constraints) from our framework can benefit state-of-the-art TTA approach TeSLA [27].”

2. We have replaced the contribution lines 82-83 with: “Inspired by our theory, we propose to add a likelihood maximization term to autoencoder-based MAE-TTT [24] to facilitate the alignment between the target sample and the source distribution. In addition, we propose sparsity constraints to enhance state-of-the-art TTA approach TeSLA [27]. We validate our proposals with empirical evidence.” 

3. We have re-organized lines 275-287 in Section 3.3 to make the relationship with MAE-TTT clearer. In particular, we make lines 276-280 as a separate paragraph to highlight the connection (the reconstruction objective) and lines 280-287 another paragraph to specify the distinction and our proposal, starting with “Despite the resemblance on the reconstruction objective, MAE-TTT doesn’t explicitly perform the representation alignment as our objectives (2)(3).”

We hope these modifications would make the message clearer and we would appreciate your further feedback!


> MAE-TTT +entropy for other benchmark datasets.

Thanks for the great question! 

In light of your question, we have started running MAE-TTT + entropy minimization on ImageNet-C. Given our computing resource, MAE-TTT on ImageNet will require 20 days to complete and we may not be able to provide full results by the end of the discussion period. Nevertheless, we will include this result in our revision. Thank you for your understanding!

For CIFAR-10/100-C: MAE-TTT requires pre-trained and fine-tuned MAE checkpoints, which are only available for ImageNet (please see the original paper [a]). We are concerned that directly using the ImageNet pre-trained checkpoints on CIFAR-10/100-C would lead to unfair comparisons and also require a significant amount of fine-tuning for the transfer between datasets. Thus, we decide to focus on ImageNet-C dataset for now, and this is the most representative and challenging benchmark. 

We thank you for your understanding and patience!

[a] Masked Autoencoders Are Scalable Vision Learners. He et al. CVPR 2022.
Also, thanks for the proof on Definition 2 and its connection with block identifiability! The argument is correct and addresses my concerns. 

Regarding my final question about improvement with sparsity regularization over TeSLA, thanks for the clarification, my earlier statement was incorrect. However, I am still not convinced since it also depends on the standard deviation in the performance of the method TesLA as well. For example, if the performance of TeSLA on CIFAR10-C is $12.5 \pm 0.3$, then the confidence intervals would overlap. But I do not see this as a major concern though I encourage authors to report the deviation in the error rate of all baselines as well.
Thanks for your detailed response! I have a better understanding of the implications of the theoretical results and how authors intended to connect them with prior learning methods for TTA (hence adjusted my rating as well). A suggestion would be to change the introduction and especially the contribution to highlight this more, it is heavily centered around the connection between theoretical results and the existing TTA approaches. The authors could mention more explicitly that they propose an improvement over MAE-TTT inspired by their theory and evaluate it on TTA benchmarks. Furthermore, section 3.3 can be improved and perhaps made into a separate section, with more details MAE-TTT being related to the theoretical analysis and proposed changes on top of it.

I would be most interested in understanding the performance of MAE-TTT +  since it directly conforms to the theory. This prompts another question: why did the authors not experiment with the proposed MAE-TTT + entropy minimization for benchmarks in Table 2? MAE-TTT is only explored in section 5.1 on the ImageNet100-C dataset and is not included for the benchmarks in Section 5.2 Wouldn't that be an important question to understand how the approach following the theoretical analysis compares with other state-of-the-art approaches? I am happy to increase my score further if the authors can perform experiments for the same.
We now provide more results on the regression task we’ve included in the rebuttal. This setup corresponds to the sparse shift setting in our theory where only two out of six dimensions of $\mathbf{x}$ are influenced by the changing variable $\mathbf{s}$. In comparison, the setting we’ve included in the rebuttal corresponds to the dense shift case. The distinction is described in lines (lines 300-304). The other setups are identical. We directly adopt the hyper-parameters from the dense shift setting without modifications. The results are also measured in MSE (the lower the better).

| Sparse Shift Distance | 18     | 24     | 30     |
|----------------------|--------|--------|--------|
| Baseline             | 1.835 | 3.323 | 5.841 |
| Our Method           | 1.145 | 1.476 | 1.596 |

We can observe that like the dense shift case, our method outperforms the baseline consistently and maintains the performance over a wide range of shift severities. In contrast, the baseline that directly uses all the feature dimensions degrades drastically when the shift becomes severe. This indicates that our approach can indeed identify the invariant part of the latent representation, validating our theoretical results.

---

Please let us know if we have resolved your concerns – thank you!
We now provide more results on the regression task we’ve included in the rebuttal. This setup corresponds to the sparse shift setting in our theory where only two out of six dimensions of $\mathbf{x}$ are influenced by the changing variable $\mathbf{s}$. In comparison, the setting we’ve included in the rebuttal corresponds to the dense shift case. The distinction is described in lines (lines 300-304). The other setups are identical. We directly adopt the hyper-parameters from the dense shift setting without modifications. The results are also measured in MSE (the lower the better).

| Sparse Shift Distance | 18     | 24     | 30     |
|----------------------|--------|--------|--------|
| Baseline             | 1.835 | 3.323 | 5.841 |
| Our Method           | 1.145 | 1.476 | 1.596 |

We can observe that like the dense shift case, our method outperforms the baseline consistently and maintains the performance over a wide range of shift severities. In contrast, the baseline that directly uses all the feature dimensions degrades drastically when the shift becomes severe. This indicates that our approach can indeed identify the invariant part of the latent representation, validating our theoretical results.

---

Please let us know if we have resolved your concerns – thank you!