Authors: Jean Erik Delanois, Shruti Joshi, Ryan Golden, Teresa Nick, Maxim Bazhenov
Organizations: Department of Computer Science & Engineering, University of California, San Diego, La Jolla, California, USA · Department of Medicine, University of California, San Diego, La Jolla, California, USA · Microsoft Corporation, Redmond, WA, USA
Artificial Neural Networks (ANNs) are increasingly deployed across diverse real-world settings, where they must operate under data distributions that differ from those seen during training. This challenge is central to Domain Generalization (DG), which trains models to generalize to unseen domains without target data, and Test-Time Adaptation (TTA), which improves robustness by adapting to unlabeled test data at deployment. Existing approaches to address these challenges are often complex, resource-intensive, and difficult to scale. We introduce CONTXT (Contextual augmentatiOn for Neural feaTure X Transforms), a simple and intuitive method for contextual adaptation. CONTXT modulates internal representations using simple additive and multiplicative feature transforms. Within a TTA setting, it yields consistent gains across discriminative tasks (e.g., ANN/CNN classification) and generative models (e.g., LLMs). The method is lightweight, easy to integrate, and incurs minimal overhead, enabling robust performance under domain shift without added complexity. More broadly, CONTXT provides a compact way to steer information flow and neural processing without retraining.
Figures & tables
Figure 1: CONTXT : Contextual augmentation via feature transforms. (a) At a chosen layer, compare the current feature vector h to a precomputed contextual feature representation c to form an ”index” (their difference) d=c−h . (b) Add a scaled version of this index, αd , to the features; α>0 injects the context while α<0 removes it. (c) Mix multiple contexts by linearly combining indices with separate scalars, e.g. αidi+αjdj .
Figure 2: ”Cow on a beach” example. (a) A representative input image shown alongside contextual examples. (b/c) The vertical axis reports the model’s maximum softmax confidence, while the horizontal axis sweeps the strength of the farm / city index injection; each subplot corresponds to a different fixed level of beach context removal (strength annotated above each panel). For both injection and removal, α=0 indicates that no context is applied. Text above each curve denotes the top-1 predicted class at that setting. Correct application of CONTXT yields proper classification.
Figure 3: Baseline accuracy for the CCT (a) and PACS (b) models. Models were trained on a single domain (Location 38 / Photo), performance on the training domain is highest while accuracy quickly degrades when tested in other domains.
Figure 4: Accuracy heatmaps for CCT (a) and PACS (b). Vertical axis: out-of-domain removal strength; horizontal axis: in-domain injection strength. Color encodes change in mean test accuracy averaged across all domains (trained and untrained) compared to baseline. CONTXT can improve performance by about 10%.
Figure 5: Domain-wise change in accuracy on CCT (a) and PACS (b). Source domains show zero shift - Photo in PACS and Location 38 in CCT - while most unseen target domains exhibit substantial improvements.
Figure 6: Flip rate (percentage of reviews whose predicted sentiment changes after rewriting) versus Self-BLEU between rewritten and original reviews for Llama 8B (a) and 70B (b). When asked to rephrase a review, the baseline (no CONTXT ) preserves sentiment, whereas models provided with opposing sentiment CONTXT flip the classification while maintaining fluency. Points along each trajectory correspond to increasing context-shift magnitude (lowest at the bottom of the plot).
Neural network robustness studies how models withstand adversarial inputs, data shifts, noise, and hardware faults while maintaining reliable predictions.
This paper studies continual test-time adaptation (CTTA), the task of adapting a model to constantly changing unseen domains in testing while preserving previously learned knowledge. Existing CTTA methods mostly focus on adaptation to the current test domain only, overlooking generalization to arbitrary test domains a model may face in the future. To tackle this limitation, we present a novel online test-time domain generalization framework for CTTA, dubbed TestDG. TestDG aims to learn features invariant to both current and previous test domains on the fly during testing, improving the potential for effective generalization to future domains. To this end, we propose a new model architecture and a test-time adaptation strategy dedicated to learning domain-invariant features, along with prototype selection and update mechanisms for effectively managing information from previous test domains. TestDG achieved state of the art on three of four public CTTA benchmarks and performed on par with the best prior method on ImageNet-C. Moreover, it showed superior generalization to unseen test domains.
Test-time adaptation (TTA) aims to align a model to shifting test domains using only unlabeled streaming data. Most existing methods implicitly infer a single global domain distribution, ignoring the multidimensional and sample-specific nature of real-world domain shifts, leading to fragile adaptation. We propose DOME, an effective domain encoder that explicitly models each sample's domain in a zero-shot manner. DOME leverages vision-language pretraining to extract dense, continuous representations, parameterizes domains as distributional variables, and introduces a momentum-updated sparse domain bank for disentangled supervision. By injecting these explicit domain cues into downstream models, even a basic entropy-minimization TTA strategy achieves state-of-the-art performance across ImageNet-C, ImageNet-R, and ImageNet-Sketch, outperforming complex TTA approaches. Our results demonstrate that robust adaptation stems not from intricate adaptation algorithms, but from explicit, structured domain representation.
Domain Generalization (DG) aims to learn representations that remain robust under out-of-distribution (OOD) shifts and generalize effectively to unseen target domains. While recent invariant learning strategies and architectural advances have achieved strong performance, explicitly discovering a structured domain-invariant subspace through second-order statistics remains underexplored. In this work, we propose CPCANet, a novel framework grounded in Common Principal Component Analysis (CPCA), which unrolls the iterative Flury-Gautschi (FG) algorithm into fully differentiable neural layers. This approach integrates the statistical properties of CPCA into an end-to-end trainable framework, enforcing the discovery of a shared subspace across diverse domains while preserving interpretability. Experiments on four standard DG benchmarks demonstrate that CPCANet achieves state-of-the-art (SOTA) performance in zero-shot transfer. Moreover, CPCANet is architecture-agnostic and requires no dataset-specific tuning, providing a simple and efficient approach to learning robust representations under distribution shift. Code is available at https://github.com/wish44165/CPCANet.