Reconstruct and Match: Out-of-Distribution Robustness via Topological Homogeneity

Authors: chaoqi  chen
luyao  tang
hui  huang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 16:03:57

Abstract: Since deep learning models are usually deployed in non-stationary environments, it is imperative to improve their robustness to out-of-distribution (OOD) data. A common approach to mitigate distribution shift is to regularize internal representations or predictors learned from in-distribution (ID) data to be domain invariant. Past studies have primarily learned pairwise invariances, ignoring the intrinsic structure and high-order dependencies of the data. Unlike machines, human recognizes objects by first dividing them into major components and then identifying the topological relation of these components. Motivated by this, we propose Reconstruct and Match (REMA), a general learning framework for object recognition tasks to endow deep models with the capability of capturing the topological homogeneity of objects without human prior knowledge or fine-grained annotations. To identify major components from objects, REMA introduces a selective slot-based reconstruction module to dynamically map dense pixels into a sparse and discrete set of slot vectors in an unsupervised manner. Then, to model high-order dependencies among these components, we propose a hypergraph-based relational reasoning module that models the intricate relations of nodes (slots) with structural constraints. Experiments on standard benchmarks show that REMA outperforms state-of-the-art methods in OOD generalization and test-time adaptation settings.

Main paper: 
 introduction: 
 Distribution shift has emerged as a major challenge for the success of machine learning systems [34, 25] . Although deep learning models are believed to generalize well to in-distribution (ID) data, a well-trained model deployed in the open world often encounters out-of-distribution (OOD) data whose contexts may differ from the training distribution, resulting in dramatic performance degeneration and raising concerns about model safety in many high-staking applications, such as autonomous driving and medical diagnosis. This gives rise to the importance of OOD generalization [92, 76] , which aims to build a robust learning machine that can perform well in unseen test environments.Regarding OOD generalization, a central theme is how to learn general features from training data that can be extrapolated to test distributions. Following this idea, a plethora of OOD generalization methods have been proposed over the past few years, including domain alignment [42] , latent feature disentanglement [61, 48, 85] , meta-learning [40, 41] , invariant risk minimization [2, 1, 96] , and augmentation-guided invariant predictor [74, 83, 95] , to name a few. On the other hand, given the natural adaptivity gap between training and test distributions [15] , recent works [29, 30, 81, 11, 8] attempt to further enhance the source-trained model through test-time adaptation [46] , which leverages unlabeled target samples to update the model in an online manner. These two series of studies address distribution shifts during training and inference and have been proven to work synergistically [8] . (c) Our REMA introduces a sparse and discrete element -slot [51] -to serve as a bridge for both reconstruction and relational modeling, avoiding explicit disentanglement in the latent space.Albeit their general efficacy for various tasks, these prior efforts largely overlook the topological structure of the image data. As a consequence, existing OOD generalization methods highly rely on specialized regularization objectives and are significantly less interpretable than human vision systems. Unlike machines, human object recognition [4] typically involves initially decomposing objects into several major components (e.g., keypoints), followed by identifying the structural relationships among these components, and finally making predictions by comprehensively considering the main components and their inherent relations. To this end, a critical question remains open in the field:How to devise a unified framework that imitates the human vision process for OOD generalization?In this work, we believe that OOD generalization requires consideration of two key aspects: 1) how to represent the structure of data, and 2) how to model the relationships between different entities.Grounded on these insights, we propose Reconstruct and Match (REMA), a general framework to endow deep models with the capability of capturing the topological homogeneity of objects without using human prior knowledge or fine-grained attribute annotations. Figure 1 illustrates the motivation.To identify major components from objects, REMA introduces a slot-based reconstruction module to map dense pixels into a sparse set of slot vectors in an unsupervised manner. This module encourages the deep model to reduce unwanted redundancy and preserve those predominant parts (the words "component" and "part" are used interchangeably). Then, to model and reason the highorder dependencies among these components, we propose a hypergraph-based matching module that discovers the relations of slots with structural constraints, i.e., topological homogeneity of the object.The main contributions of this paper are summarized as follows:• We propose REMA, a novel OOD generalization framework to mitigate distribution shifts in deep learning models by imitating the human visual recognition process.• We introduce a self-supervised reconstruction process to identify informative parts from objects without additional supervision, and a hypergraph matching module to reason the high-order part-based object relationships and interactions across domains.• Experiments on six widely used benchmarks demonstrate that REMA outperforms state-ofthe-art methods in OOD generalization and test-time adaptation settings. 
preliminaries: 
 Problem Setup. Assume that X is the input space, Z is the latent space, and Y is the output space.The predictor f = h • g is comprised of a featurizer g : X → Z that learns to extract embedding features, and a classifier h : Z → Y that makes predictions based on the extracted features. The goal of OOD generalization is to find a predictor f : X → Y that generalizes well to all unseen target domains. In deep neural networks, empirical risk minimization (ERM) [72] is capable of learning highly predictive features, making it the simplest baseline for this problem.ERM Objective. We train the model by minimizing the empirical cross-entropy loss function:EQUATIONwhere P tr is the training distribution, Θ denotes the parameter space, and ℓ is the loss function. The trained model will be evaluated on a test set from a test distribution P ts , where P tr ̸ = P ts . 
proposed method: 
 In this section, we provide a detailed description of our Reconstruct and Match (REMA) framework. As shown in Figure 2 , REMA consists of two key modules tackling the main component discovery (Section 3.1) and high-order relational modeling and reasoning (Section 3.2).For OOD generalization/adaptation, a long-standing issue is how to distinguish between transferable (domain-agnostic) and non-transferable (domain-related) information in the data. An ideal solution is to disentangle latent representations in an unsupervised manner (see Fig. 1(a) ), which has been proven to be prohibitively difficult and even infeasible [50] . Therefore, most prior works (see Fig. 1(b) ) aim to learn domain-invariant representations through elaborate feature alignment modules, such as adversarial training [21] and pseudo-labeling [37] , without exploring the intrinsic composition of the features themselves. Despite their general efficacy, these methods may introduce a significant amount of redundant and noisy information during the alignment process. For humans, when identifying an object, we tend to first look for its main components, which can be regarded as a process of information compression from a learning perspective. Inspired by this, we aim to mimic the very first step of the human recognition process by identifying the main components from object images without human prior knowledge or fine-grained annotations.To be specific, we introduce a self-supervised reconstruction module based on Slot Attention [51] , termed Selective Slot-based Reconstruction (SSR), to assist deep models in learning a set of part-based representations for sparsely characterizing a target object. Formally, given an input x, g extracts a set of feature embeddings z ∈ R N ×dz , where N is the number of embedding and d z denotes embedding dimension. The Slot Attention module will first take a set of slots s ∈ R K×ds (K is the number of slots) and the feature map z, then project them to dimension d s by a linear transformation f k for slots and f q , f v for z. The Slot Attention will be trained as follows,EQUATIONwhere q = f q (z) ∈ R K×ds , k = f k (z) ∈ R N ×ds , and v = f v (z) ∈ R N ×ds denote the query, key and value vectors respectively. A ∈ R N ×K stands for the attention matrix. Unlike self-attention, the queries in slot attention are a function of the slots s ∼ N (s; µ, σ) ∈ R K×ds , and will be updated iteratively over T iterations. The slots are initialized by random sampling from a standard Gaussian distribution. Queries at iteration t are represented by qt = f q (s t ), and the slot updating process is: s t+1 := update(attn(q t , k), v). After each iteration, a Gated Recurrent Unit (GRU) is employed on the slot representations s t+1 to update their states, integrating new information while retaining relevant context from previous iterations.Since the number of slots K is manually predefined and fixed during training, attention networks may learn redundant or even incorrect associations, thus affecting the understanding of objects and scenes. Therefore, our SSR introduces a selective slot attention mechanism that dynamically adjusts the importance of slots for reconstruction. To quantitatively represent which slots are more important for reconstruction, we first introduce an importance score ρ for each slot. For each slot s i , we use a lightweight neural network h θ to predict an initial importance score:ρ init i = σ(h θ (s i )), where σ(•) is the sigmoid function, ensuring ρ init i ∈ [0, 1].To capture interactions among slots, we further introduce an interaction matrix W ∈ R K×K , where each element W ij is defined asW ij = softmax j (U s i ) ⊤ (V s j )/ √ d .Here, U and V are learnable projection matrices, and d is a scaling factor. The final importance score ρ i for each slot s i is then computed asρ i = K j=1 W ij ρ init j .We then scale each slot representation by its respective importance score to create a weighted slot representation si = ρ i • s i . To this end, the overall training objective of SSR can be formulated as:EQUATIONwhere λ is the balancing parameter. The regularization term encourages sparsity in the initial importance scores, guiding the model to utilize only the most relevant slots for reconstruction. On the other hand, the less relevant slots are down-weighted rather than discarded, which may help retain subtle information. Compared to continuous semantic representations, this training process is more manageable due to the discretization of the slots. Since we lack truly fine-grained annotations, the learned slots or "concepts" may not be as readily interpretable as individual words. It is worth noting that incorporating a strong visual feature extractor could significantly boost baseline performance.However, for the sake of fairness, we have opted not to use approaches, such as DINOSAUR [64] .Albeit we have obtained part-based representations in the form of slots, their inner-and interobject relations are still under-explored. Prior graph-based [9, 7] and non-graph [33] works only accommodate pairwise relationships (e.g., connection between two nodes in a simple graph), but are inadequate to model the high-order relations inherent in images. Addressing this issue, we introduce a simple yet effective technique-HORR-to model high-order topological relations with a collection of graph nodes and hyperedges, divided into three stages: construction, learning, and matching.Definition 1 (Topological Homogeneity) We conceptualize topological homogeneity between the same object across images as a (hyper)graph matching problem, which is mathematically formulated as a relaxed quadratic assignment problem [52] ,EQUATIONwhere A ∈ R n×n and B ∈ R m×m are the adjacent matrix encoding structure information of the graph G A and G B respectively, n and m are the number of graph nodes, || • || F is the Frobenius norm, X u ∈ R m×n is the unary affinity matrix and generally specified as the node affinity M aff , and X is the relaxed permutation matrix encoding node-to-node assignment.Hypergraph Construction. Given visual slot representations from the previous step, we carry out a linear transformation to obtain graph nodes V A and V B . This transformation is intended to map the visual features into the graph domain, ensuring effective graph matching. Then, we build hypergraph graphs G A and G B in both domains, modeling the topological structures inherent in the images, x A and x B , respectively. The hypergraph [17] is defined as:G A/B = (V A/B , E A/B , W), where V A/B is the node sets, E A/B is the hyperedge sets, and W ∈ R |E|×|E| is a diagonal matrix of edge weights. The hypergraph is denoted by an incidence matrix H ∈ R |V|×|E| , where H(v, e) = 1 indicates the node v ∈ e and H(v, e) = 0 indicates v / ∈ e. We adopt non-parametric density estimation, i.e., K-Nearest Neighbors (KNN), to establish hyperedge connections. The Euclidean distance is used to calculate the distance between node embeddings.Hypergraph Learning. After constructing the hypergraphs, we update the graph node features according to the connectivity defined by the hyperedges. Technically, we introduce hypergraph convolution [17] to perform message-passing and feature aggregation,EQUATIONwhere Θ represents the parameter to be learned during training. By doing so, we can explicitly model intricate correlations among different parts of an image (i.e., slots). On the other hand, in contrast to grids, sequences, and even graphs, this high-order modeling can more effectively capture relations among nodes, avoiding excessive or erroneous pairwise connections.Hypergraph Matching. Since the graph node features have been updated by Eq. ( 5), we introduce an affinity matrix W aff to measure the node correspondence between G A and G B . Following [19, 44] , we use the differentiable Sinkhorn layer [67] to calculate the affinity matrix Ŵaff , where each element indicates the degree of matching between pairs of nodes across graphs. To leverage the topological structures, we need a training objective to minimize the pairwise structural discrepancy between the hypergraphs [78] . In particular, we follow the Definition 1 to specify unary affinity matrix X u as the obtained node affinity matrix Ŵaff . Since the constructed hypergraphs naturally encode rich high-order relationships, they enhance the cross-domain topological matching process. Conversely, the matching process introduces additional structural knowledge to the current graph through message propagation. Formally, the hypergraph-based matching objective is formulated as follows:EQUATIONwhere the(i, j) element in Y Π ∈ R n×m is 1 if v A i ∈ G A and v B j ∈ G B belongto the same object class, otherwise it will be 0. Here, Y Π can be regarded as pseudo labels for matching due to the absence of ground-truth correspondence. Moreover, as the matching primarily targets high-level semantic relationships, we avoid imposing additional structural constraints [23, 43, 44] on this training objective, thereby significantly simplifying the training process.Putting everything together, the full training objective is formulated as follows:EQUATIONwhere α and β are hyper-parameters for balancing different loss terms. Our algorithm first trains deep models using the reconstruction objective to obtain discrete part representations with sufficient sparsity. Then it turns to model the cross-image components correlations as the graph matching problem where we further introduce the structural regularization term into the matching loss.Note that our approach can be directly applied to test-time adaptation, similar to those self-supervised methods [69, 49, 3, 20] . The difference lies in the fact that their proxy tasks are somewhat heuristic, such as predicting rotations, jigsaw puzzles, and random masking. In contrast, our method directly exploits the intrinsic structured properties of the data, making it more robust and versatile. 3) is set to 0.01 in all experiments. α and β in Eq. ( 7) is set to 10 and 0.1, respectively.Baselines. We compare REMA against two types of baseline methods.(1) OOD Generalization:We adopt the leave-one-domain-out evaluation protocol and follow the model selection strategy used in [25] . The main baselines have optimization-based CORAL [68] , MLDG [40] , CDANN [45] and MMD [42] , augmentation-based Mixstyle [95] , SagNet [55] and I-Mixup [80, 82, 84] , and so on. (2) Test-time adaptation: All methods are based on online batch-level test data adaptation setting and the following baselines are included: entropy-minimization: Tent [75] and SHOT [47] , pseudo-labeling: T3A [29] and TAST [31] , consistency-alignment: TSD [79] and TIPI [56] . Test-Time Adaptation. We compared REMA with existing advanced TTA methods, and the results are summarized in Table 2 and Table 3 . REMA consistently provides improvements on multiple datasets. Compared to TSD [79] , it achieves a 1.3% improvement on three pixel-level corruption datasets and a 3.1% improvement on three challenging cross-domain datasets. Compared to the SOTA, REMA has three major advantages: (1) Tent [75] or SHOT [47] based on entropy minimization has limitations, as they can overcome distribution shifts caused by pixel corruption but fail to handle cross-domain test data, even leading to performance degradation. However, REMA can handle both types of OOD data. (2) T3A [29] and TAST [31] based on pseudo-labels have performance limitations, and the pros and cons of pseudo-labels significantly impact the adaptation, which can be observed in their ceilings on two types of datasets. REMA enhances resistance to noisy labels by capturing higher-order semantic dependencies. ( 3) REMA converts dense pixels into sparse slots, accurately capturing domain-invariant features, which will not be affected by pixel corruption and domain shift. Implementing consistent optimization during test time makes REMA more robust, resulting in a 5.1% accuracy improvement across all benchmarks, i.e., ERM vs. REMA. Ablations of key modules in REMA. In this part, we provide the ablation results in Table 4 , investigating the independent and combined effects of SSR and HORR proposed in REMA. 
ablation studies: 
 As can be seen, incorporating SSR and HORR separately leads to improved generalization performance, demonstrating their contributions to abstraction and relational modeling. In addition, integrating SSR and HORR in our method yields the best performance, highlighting that the two modules systematically work together and reciprocate each other. Moreover, as shown in Figure 4 (a), the full REMA can provide more complete and accurate representations of objects.Predicted ResultsFigure 5 : Learned affinity vs. GT Analysis on continuous test-time adaptation. We empirically evaluate the continuous learning ability of REMA by comparing it with state-of-the-art test-time adaptation methods, namely Tent [75] , EATA [57] , and DeYO [38] . The results are presented in Fig. 4(b) , where the adaptation order is from left to right. We can see that the proposed REMA substantially and consistently outperforms all baseline methods as adaptation proceeds, revealing the robustness and the capability to handle open-world changes.Analysis on hypergraph matching. In Figure 5 , we visualize the learned doubly stochastic node affinity matrix and the ground-truth (GT) matrix. We observe that the proposed graph matching method effectively identifies the correct node affinity, revealing its effectiveness in handling cross-domain higher-order relationships. We extract feature embeddings using CIFAR-10C data, where the corruption type is snow with a severity level of 5. The different color stands for different classes. We can observe that our REMA exhibits better clustering patterns with both SSR and HORR modules. When SSR is removed, some closely related categories (e.g., animals) become highly mixed, indicating the advantage of SSR in extracting robust semantic representations. On the other hand, without HORR, intra-class variations increase, indicating some degree of ambiguity between different categories. Our HORR addresses semantic ambiguity by modeling topological homogeneity.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We have fully disclosed all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
related work: 
 OOD Generalization. The aim of OOD Generalization is to train a model using data from the source distribution so that it can perform well on an unknown target distribution. A series of works can be divided into three categories: (1) Minimizing distribution differences [42, 45, 68, 90] : using meta-learning to simulate distribution shift [40, 18] , or learning an invariant transformation based on adversarial learning [42, 45, 89, 94, 93] . (2) Domain-invariant representation learning: by analyzing the feature learning process of deep neural networks [16, 66, 27, 63, 12, 13] , it is inferred how ID data can be generalized to OOD data [6, 65] . By understanding the interactions between ERM and generalization, domain-invariant features are learned. (3) More generic OOD generalization: many studies focus on generalization under a relaxed assumption, such as open environments with unknown classes [8] , semi-supervised settings in the testing environment [86] , adapting foundation models [91] , and OOD data synthesis [60, 59, 22] , aiming to identify generalizable features from sophisticated test distributions. However, previous studies have only focused on the instances themselves and have not explored the inherent high-order semantic dependencies in ID and OOD data, overlooking the topological homogeneity between distributions.Test-time Adaptation. The goal of TTA [46] is to adjust the source trained model using test data without ground truth during the testing phase [47, 58, 10, 87, 79, 29] . Most of the existing methods directly perform gradient optimization on the test samples, with optimization objectives including prediction entropy [75] , self-training and stochastic restoring [77] , or based on normalization [54] , etc., making the model adapt to the dynamic target environment at each step. Recently works further consider the scenario of noisy outlier samples [24, 70] appearing in the test stream. However, previous studies have simply optimized the test instances based on self-training, ignoring the topological relationship between samples and the high-order semantic dependencies within the test data steam, leading to suboptimal adaptation.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [No] Justification: After internal review and patent approval, we will release the code.Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • 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. 
conclusion: 
 In this paper, we propose to achieve out-of-distribution robustness from the perspective of ensuring topological homogeneity between the same object class regardless of the surrounding environments.To achieve this goal, we propose a new framework REMA to first obtain a set of sparse and discrete representations from dense pixels and then model the high-order topological relations and dependencies via hypergraph (a generalized form of graph). Our experimental results reveal that REMA achieves superior performance on standard OOD generalization and test-time adaptation benchmarks.In the future, we aim to extend this framework to more complex and dynamically changing scenes.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
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 have provided these details in the supplementary.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.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: This paper received an Accept, two Weak Accepts, and a Borderline Accept. All reviewers agree that the work presents an interesting and innovative approach. Notably, they commend the novelty and well-motivated design of the selective slot-based reconstruction module and the hypergraph-based relational reasoning framework. After the rebuttal, the authors successfully addressed all the concerns and questions raised by the reviewers, resulting in all reviewers being satisfied post-rebuttal. Consequently, the AC has decided to Accept this paper, recognizing its potential impact on both the theoretical and practical aspects of out-of-distribution generalization tasks.

Reviews: review: paper_summary: paper_summary: A new methodology, REconstruct and MAtch (REMA) is introduced to learn a more robust and generalizable feature set in computer vision models.  REMA relies on a slot attention module to learn sparse embeddings of features which characterize a target object, and this module is then coupled with a high-order relational reasoning (HORR) module which creates a graph representation of the object.  This graph representation encodes how the sparse features relate to each other.  Because the feature encodings are both simplified through the sparse attention module and contain a learned topological relation from the HORR module, they are expected to be robust to domain drift. The robustness to domain drift is demonstrated using several benchmark datasets with consistent networks and training hyperparameters shared between them.  REMA consistently and clearly outperforms SOTA methods for OOD generalization; the results are convincing.
strength_weakness: strength_weakness: The report of the proposed method is exceedingly technical.  Discussion of the intuition motivating the method is restricted to a brief statement concerning the human visual system, and no connection back to the human visual system is made throughout the remaining discussion of the methodology.  No further motivation or context is provided for algorithmic design choices, which prevents the reader from understanding *why* the method works well. The emphasis is entirely on the *how*. I suspect that a significant literature review was performed as part of the algorithm design; sharing the literature review context for these design choices would have helped motivate and clarify the methodology.  It could easily have been included in the appendix without detracting attention from the main contribution of the work. Alternately, the context could have been included in the main result, and the implementation details left to the appendix.

The benefit of the HORR module is open to some question.  The impact of the HORR module vs the SSR module is discussed in Section 4.3, Table 4, and shown in Figure 6 with the tSNE embeddings.  The results of Table 4 are reported without uncertainty, and the values are close enough (within 2%-8.5% of the baseline without REMA) to make uncertainty a valuable indicator of the role of each module. Given the stochastic variation inherent to the tSNE algorithm, Figure 6c (SSR w/o HORR) and 6d (SSR w/ HORR) could be considered identical since there are no repeated embeddings reported, and no figure of merit describing the variance for repeated embeddings and clustering is included.  Visually, the embeddings are very close, so although the value of the SSR module is clearly demonstrated in disentangling latent features in Fig. 6c, the value of the HORR module has some doubt.  (An example of suitable metrics for quantifying the author's claim of better clustering in Figure 6d vs. Figure 6c might be DBSCAN or OPTICS applied to multiple embeddings, and quantified with the mean and std of the v-measure score for the clusters.)  The supplemental Figures 6-8 do little to clarify the situation; the authors state only that SSR and HORR behave differently depending on the type of corruption. An additional place to clarify the behavior of the individual SSR and HORR modules would have been in Figure 4a, where the grad cam results with and without the SSR and HORR modules could have been reported.  Finally, the comment in the supplement regarding training the SSR module first suggests that although including the HORR module may be beneficial, it is not necessarily key to the success of the work as no mention of training HORR first is made.

Overall, the results are extensive, and support the conclusion that REMA is an improvement on the SOTA for OOD robustness.  But the lack of contextualization and motivation for the algorithm in general, and the behavior of the different modules specifically, is a weakness of the paper.
questions: questions: 1. The stated goal of the Selective Slot-based Reconstruction (SSR) module is to create a sparse embedding of target features.  In section 3.3, it would seem that a variational auto encoder type approach is used to train the encoder.  Sparsity can already be considered forced due to the size of the feature vector/latent space of the VAE; a more sparse feature vector is just one that is smaller in this context.  How does the additional SSR module change the latent representations to make them more sparse?  Is it just removing specific frequency components?  Why is it not possible to do this directly on the latent space of the encoder through a specialized loss function without the use of an additional module?  

2. In section 3.1 "(A) standard MLP skip connection is applied" after the GRU; why? What is the intuition here?

3. What is the minimum and maximum connectivity of the hypergraphs constructed in Section 3.2?  How would this dimensionality relate to the dimensionality of the sparse embeddings?  Of the original feature vector extracted from the encoding module?

4. Please provide citation(s) for the statement in Supplemental section C.1 "sparse modeling based on slots may struggle to accurately separate the scene into several main components".  

5. In the supplemental section B.2, why is it reasonable to expect that the SSR and HORR modules have different effects based on the type of data corruption?  Shouldn't the keypoints of a target object (and therefore the sparse embedding of those key points as well as the topology of how those keypoints relate to each other) be unaffected by whether it is fog or frost corrupting the image?
limitations: limitations: The authors report limitations in the supplemental material.
paper_summary: paper_summary: The authors propose an approach to tackle OOD generalization via a method that tightly combines learning-based feature extraction with graph-based relationship modelling to explicitly learn and represent the topological structure of image data, achieving competitive results across a range of OOD and test-time adaptation benchmarks.
strength_weakness: strength_weakness: _TLDR; I do not see any severe ‘prohibitive’ weaknesses in this work, but have a few questions and requests that I’d like the authors to clarify & address._

- Some insight into what the slots actually represent would be insightful and elevate the paper, see below.
- Missing detail on complexity, please see question section below.
- Missing details regarding inference and training procedure, see below.
- Some inconsistencies in notation as well as some typos that should be corrected
questions: questions: **Main concerns, questions & potential improvements:**

**[Q1]**: Consistency in notation could be improved.   
 The authors initially use lower-case characters to indicate functions, e.g. in equation (1), and mathcal font to indicate sets (e.g. input, output, feature space/domain);   
However, this suddenly changes during the introduction of the method (l.94),  and mathcal is used to refer to transformations (e.g. K_beta, which is nothing else than a function);   
Then slightly later (l.135), mathcal refers to the ‘result’ of the linear transformation that essentially forms the graph nodes;   
$\textrightarrow$ I’d highly recommend to keep a consistent notation to avoid confusion (as it did confuse me); 

**[Q2]**: Some missing details regarding training & inference procedure. 
- How exactly is training performed? The authors mention that the “algorithm first trains deep models using the reconstruction objective” (l.171) – is the model (or a part of it) then frozen when the graph is created/HORR trained, or is HORR simply added and everything then trained/finetuned? The appendix section doesn’t make this clear to me either, and some more details here would be helpful.
- The authors mention in l.99 that the queries “will be refined during the T attention iterations”. How many iterations are employed in practice, and is this performed for each new image pair (i.e. each step) during training? 
- How exactly is the inference actually performed at test time? Does inference require pairs as employed during training, or what is the exact setup there? 

**[Q3]**: Details around complexity;   
Following up from Q2, what is the ‘complexity’ of the method in contrast to other related methods, e.g. in terms of inference time?  
Note that given the impressive performance and graph-based nature, this could be an interesting insight to the reader (to gauge potential trade-offs between multi- and 1-step reasoning methods)  

**[Q4]**: Details regarding slots / information bottleneck:   
- The authors mention in the appendix that ‘5’ slots have been typically used. Is this the same for all datasets? And how many were selected as relevant?
- Do these numbers change across datasets or even classes? And if so, could you provide possible insights into why, e.g. multiple objects / increased complexity / more components, or similar.

**[Q5]**: Insights into the actual 'slot correspondence'/representation:  
- 5 slots seems quite few to represent the content of an image. Is this due to the simplicity of the images, and do the authors have some intuition how this would change for `real-world' natural settings? 
- How do the slots actually 'align' with 'components' of objects: Taking the motivational picture of the horse, would your algorithm actually identify 'parts' of an object as a component, or rather different objects in an image, or entirely different? Some insight into these aspects would be highly interesting to actually see to which extent the inner workings align with what we would expect from humans (which is, after all, your underlying motivation) 

---

**Additional comments:**   
- typo l.92: embedding -> embeddings 
- capitalization l.109: we -> We
- typo l.216: methods […] is -> are
limitations: limitations: Limitations have been adequately addressed by the authors in the appendix;   
$\textrightarrow$ I highly appreciate the authors being honest and providing ‘proper’ limitations to their method in terms of potential applicability!
paper_summary: paper_summary: This paper presents REMA which designed to improve the robustness of deep learning models against out-of-distribution (OOD) data. REMA employs a selective slot-based reconstruction module to dynamically map dense pixels into a sparse set of slot vectors, enabling the identification of major components from objects in an unsupervised manner. Additionally, a hypergraph-based relational reasoning module is introduced to model high-order dependencies among these components, ensuring topological homogeneity. Experiments conducted on standard benchmarks demonstrate that REMA outperforms state-of-the-art methods in OOD generalization and test-time adaptation settings, highlighting its effectiveness in handling distribution shifts and enhancing the adaptability of deep models in real-world, non-stationary environments.
strength_weakness: strength_weakness: 1. The paper does not provide a detailed analysis of the computational overhead introduced by the new modules, which could impact the scalability of the approach for large-scale or real-time applications.

2. While the experiments cover several benchmark datasets, the generalizability of REMA to other types of datasets or more diverse real-world scenarios is not fully explored.

3. The effectiveness of REMA relies on several hyperparameters (e.g., the number of slots, attention iterations), and the paper does not thoroughly investigate the sensitivity of the model's performance to these parameters.
questions: questions: 1. Parameter Sensitivity: How sensitive is REMA's performance to the choice of hyperparameters such as the number of slots and attention iterations? Could an automated hyperparameter tuning method improve the robustness and generalization of the model further?

2. Computational Efficiency: What is the computational overhead associated with the selective slot-based reconstruction and hypergraph-based relational reasoning modules? How does this overhead impact the scalability and real-time applicability of REMA in large-scale or resource-constrained environments?
limitations: limitations: Yes
paper_summary: paper_summary: The paper proposes a method for improving domain generalization and test time adaptation by introducing a selective variant of slot attention, formulating the relationship between slots across images as topological homogeneity between hypergraphs constructed based on the slots, and thereby matching occurances of the same object between different images.
strength_weakness: strength_weakness: The proposed method consists of fairly many newly introduced steps.
While ablations show their necessity of the major modular building blocks of REMA for achieving the eventual results, these building blocks are not analyzed in detail.
For understanding them (and therefore REMA), it would be necessary to expound each individual step and confirm that it has the intended effect and that the explicit or implicit mathematical assumptions hold.
For example, the results of standard slot attention can be of varying quality, and it would be important to verify whether and how such imprecisions introduced there propagate through the next steps.
The analysis in Figure 5 goes in this direction, but it is not clear what exactly is probed and measured here. A similar analysis as in for the other components of the method would be helpful.

**Main issue of the paper**: There exist different communities in ML which study the generalization performance of DNNs to distribution shifts in very different settings. Unsupervised Domain Adaptation (UDA) is separate from Domain Generalization (DG) and both are separate from OOD generalization. Putting papers which study either of those in the same Table is wrong, misleading and highly confusing. All of these settings have their benefits and challenges and the numbers are simply incomparable if the initial conditions, such as which data is available during training, differ this strongly. The authors conflate and confuse all three settings on multiple occasions which makes the paper confusing and the results are impossible to understand. It is further completely unclear in which of the three their proposed method falls into. Below, I explain the issue in much more detail and provide specific problematic text instances. The paper needs a major revision which includes reworking all sections and most Tables.

line 183: "For OOD generalization, we leverage the three most widely used benchmark datasets." The datasets PACS and Office-Home are usually used in Unsupervised Domain Adaptation (UDA), not OOD generalization. Datasets commonly used to benchmark OOD generalization are datasets like ImageNet-R, ImageNet-Sketch, ObjectNet or ImageNet-V2. The choice of the baselines the authors compare to also gives the impression that they confuse domain adaptation with OOD generalization: CORAL and DANN are UDA methods.


line 66: "The goal of OOD generalization is to find a predictor f : X 7→ Y that generalizes well to all unseen target domains." This is correct, but then the used baselines are wrong because they assume access to the target domains and are trained using this information. It is not clear to me how the UDA methods (shown in Table 1) can be implemented without this access.


The setting of UDA differs drastically from OOD generalization: In UDA, we attempt to learn a model which performs well on the source domain and a set of target domains, using unlabeled data from the target domains. That is, in UDA, we assume access to the unlabeled test set at training time. This differs from the OOD generalization setting: Here, we want to train a model on a source dataset and then test its generalization to an unseen test set. That is, we assume no knowledge about the test time distribution shift. This is a crucial difference and it is wrong to confuse the two terms. The implemented baselines, e.g. CORAL and DANN, are used in a UDA setup, not in an OOD generalization setup. In this light, the wording "Following common practice, the model selection is based on a training domain validation set" is meaningless because the common practice for model selection differs depending on which setting we consider.


The related work section is similarly confusing. On the one hand, the authors wish to review works on OOD generalization where the goal is "to train a model using data from the source distribution so that it can perform well on an unknown target distribution". [41] and [42] use the DG setting, i.e. assumes access to n-1 target domains. [63] assumes access to all target domains. It is unclear why [61] is a good citation for domain-invariant learning. I haven't checked all citations, but the ones I checked **all** assumed access to the target domains.


The setting the baselines use in Table 1 is inconsistent. CORAL and DANN have been trained in an UDA setting. That is, they assume access to unlabeled target data at test time from a target domain. VNE and IRM are used in a domain generalization setting where they assume access to n-1 domains and aim to perform well on an unseen domain. These are two very different settings which are incomparable to each other and it is wrong and confusing to put them in the same table without any discussion. Further, the authors aim to test OOD generalization where no access to unseen target domain data is available (line 66) which is incompatible with either UDA or domain generalization.


The baseline numbers for the cited papers are wrong. I checked the numbers of CORAL and VNE from Table 1 with the numbers in the original papers and they do not match. If the authors reimplemented all the baselines, they need to state this and discuss where the performance differences come from.


Checking papers-with-code for the PACS dataset (https://paperswithcode.com/sota/domain-generalization-on-pacs-2), the best numbers are close to 100% for the domain generalization setting. Since the authors use DG benchmarks in their table 1, I assume that using other benchmarks in the DG domain would be valid as well. If we filter for RN50 architectures only, the best number is 90.5% which is higher compared to the best number reported by the authors.


Looking at the Office-Home dataset (https://paperswithcode.com/sota/unsupervised-domain-adaptation-on-office-home), the best number here is 90%. Here, the setting is UDA. But as mentioned before, it is not clear which setting should be used because the authors include both UDA and DG benchmarks in their Table 1. The best number in the UDA setting is higher than what the authors report in Table 1. There are actually five separate benchmarks for OfficeHome on papers with code (Domain Generalization, Domain Adaptation, Universal Domain Adaptation, UDA, Partial DA) and it is not clear which one should be used.


It is not clear how the models are trained. The test sets are comprised of different domains. Do the authors train their models on all of them or on e.g. two of the domains and test against the rest or on one of the domains and then test against the rest?


For test-time adaptation, the authors missed important works which perform better than their method. SLR achieves an error of 48.7% on the highest severity of ImageNet-C [A]. ETA and EATA [B] achieve an error of around 52%. 
[A] Mummadi et al. "Test-time adaptation to distribution shift by confidence maximization and input transformation"
[B] Niu et al. "Efficient test-time model adaptation without forgetting."


Press et al. [C] showed that most test-time adaptation methods collapse when adapting for long periods of time. The authors should test their model on the proposed CCC benchmark to analyze whether their method also suffers from the collapse. The authors are writing that they test their method in  the continuous test-time adaptation scenario in Fig. 4b, but it is entirely unclear on which dataset they are testing it. Is it the CCC benchmark? Notably, if the authors are testing their approach on the Continual Test-time adaptation benchmark from [D] (although I am just guessing here), Press et al. showed in [C] that one needs to adapt for longer time periods to show-case collapse and the adaptation periods proposed in [D] are insufficient. The authors must state clearly which continual learning benchmark they are using and I would suggest to report numbers on both [C] and [D].
[C] Press et al. "RDumb: A simple approach that questions our progress in continual test-time adaptation"
[D] Wang et al. "Continual test-time domain adaptation"
questions: questions: The paper requires a major revision to make it much more clear which setting the authors are targetting: Is it:
- OOD generalization as suggested by the title? Then, the authors **cannot** assume any knowledge about the distribution shift at test time. Then the baselines in Table 1 cannot be used as they do assume access to the target domains.
- Unsupervised Domain Adaptation as suggested by the baselines in Table 1, i.e. DANN and CORAL? Then, the authors need to reformulate their title, abstract, motivation etc. and remove all instances of "OOD generalization" because this is very confusing otherwise. The authors also need to remove the Domain Generalization baselines from Table 1.
- Domain Generalization as suggested by the IRM and VNE baselines in Table 1? Then, the authors need to remove the UDA baselines from their Table 1 and remove instances of mentioning OOD generalization from the paper.


The authors need to add Test-Time adaptation baselines to their Table and improve their continual Test-Time evaluation.

The authors need to make more comprehensive evaluations on the influence of the different components of their loss.
limitations: limitations: The limitations section in the appendix is very short. A clear limitation is that the method introduces several hyperparameters which need tuning. This imposes an additional computational overhead which has not been discussed.

Another limitation which can be fixed is that the authors did not study the continual learning setting to analyze whether their method collapses when adapting for long time periods.

In 2024, most people work with large-scale pretrained models such as e.g. CLIP. It is not clear how this method could be used on those models and whether improving ResNet50's performance is relevant nowadays.

Comments: title: Thank you for the responses.
Thanks for the rebuttal
Appreciation for Your Revised Review and Understanding
I apologize for the misunderstandings in my review
comment: Thank you for your insightful feedback and constructive comments, which have been invaluable in enhancing our manuscript. We will incorporate the additional results and discussions in the final version.
We are glad to hear that our response helped resolve your questions. Thank you again for your time and constructive feedback!
I'd like to thank the authors for their responses and the additional provided information & insights -- especially the visualisations regarding slots & image regions; 

Having read the other reviews and rebuttal, I will stick with my original rating and recommend weak acceptance
After reviewing the rebuttal and considering the comments from other reviewers, I will raise my score. My questions have been satisfactorily addressed. Thank you to the authors.
Dear Reviewer EaaQ,

Thank you for your thoughtful reconsideration of our paper. We appreciate your willingness to update your evaluation and withdraw your previous concerns. Your revised understanding and increased score significantly aid in the review process, and we are grateful for your efforts to resolve these issues.

Best regards,

Submission4013 Authors
Dear authors,

I would like to acknowledge major misunderstandings of the paper and the relevant literature on my part. I was not aware that DANN and CORAL are now also DG benchmarks. Seeing those in the table contributed to me thinking the paper is based in UDA. I apologize for missing this and will raise my score to 6 and withdraw my concerns.

Best, reviewer Eaaq