Context is All You Need

Abstract: 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.

• The paper introduces CONTXT, a method that adapts neural network feature representations using cached context vectors for efficient test-time adjustment.
• It demonstrates robust improvements in classification accuracy, achieving gains up to 25% under severe domain shifts through minimal arithmetic adjustments.
• The method applies to both image classification and large language models by steering hidden activations without retraining or adding extra parameters.

Contextual Augmentation via Feature Transforms: The CONTXT Method

Introduction and Motivation

Artificial Neural Networks (ANNs) deployed in real-world systems frequently encounter domain shift, leading to a degradation in generalization and reliable performance under out-of-distribution (OOD) settings. The practical requirements of domain generalization (DG) and test-time adaptation (TTA) are fundamentally challenged by expensive, unstable, and resource-intensive solutions that depend on retraining or complex adaptation mechanisms. The paper "Context is All You Need" (2604.04364) introduces CONTXT (Contextual augmentatiOn for Neural feaTure X Transforms), a lightweight and broadly applicable feature-level adaptation procedure designed to steer ANN representations using contextual priors while imposing negligible computational or engineering burden.

CONTXT is motivated by analogs with neurobiological systems, particularly the role of hippocampal-prefrontal interactions in context-appropriate behavior and memory recall. The method operationalizes these insights through direct arithmetic manipulation of hidden features, enabling robust inferences under context shift with a minimal protocol. The mechanism eschews retraining, sparse autoencoders, or paired prompts, setting it apart from prominent activation-steering and test-time adaptation strategies.

Methodology

CONTXT centers on localizing and modulating the context in feature space at chosen network layers. For a context $\kappa$ (characterized, for example, by domain, style, sentiment, or persona), a representative context vector $c_{\ell, \kappa}$ is obtained by averaging the features $h_\ell(x)$ over all samples $x$ drawn from $\kappa$ at a specific layer $\ell$. The current input's feature $h_\ell(x)$ is then compared to $c_{\ell, \kappa}$ to compute an index $d_{\ell, \kappa}(x) = c_{\ell, \kappa} - h_\ell(x)$.

Feature adaptation is performed by adding a scaled version of this index:

$$\tilde{h}_\ell(x) = h_\ell(x) + \alpha \, d_{\ell,\kappa}(x),$$

where the scalar $c_{\ell, \kappa}$0 is optimized on a small validation set to inject ($c_{\ell, \kappa}$1) or suppress ($c_{\ell, \kappa}$2) contextual priors. The mechanism extends naturally to mixtures of multiple contexts, supporting flexible and fine-grained adaptation.

Figure 1: CONTXT index computation and application across network layers and contexts.

The CONTXT layer introduces negligible overhead: the context vectors can be cached, only light vector arithmetic is needed at inference, and no additional parameters are introduced. Critically, this enables plug-in deployment in frozen or pretrained models without risk of catastrophic forgetting.

Main Empirical Results

Image Classification

The empirical evaluation involves both canonical OOD classification benchmarks and representative failure cases. For the motivating "cow on a beach" scenario (an ImageNet-trained VGG19 misclassifies an OOD cow image), the method demonstrates that injecting farm context and removing spurious beach context realigns the hidden representation, recovering the correct label (ox) and increasing model confidence within a broad range of index magnitudes. Injecting irrelevant context (city) provides no correction, underscoring the selectivity of the method.

PACS and CCT Benchmarks

Systematic domain adaptation experiments on PACS and CCT reveal that CONTXT yields monotonic and stable improvements under distribution shift, particularly when in-domain injection is coupled with OOD context removal. Sweeps over parameter strengths produce accuracy gains up to 10% and stabilize tuning, with the most pronounced improvements in the most challenging domains.

Figure 2: OOD correction and the necessity of semantically-relevant context indices for recovering classification performance.

Domain-wise Effects

Detailed analyses confirm that the method preserves in-domain accuracy and preferentially boosts accuracy in severe domain shift situations (e.g., +20% on PACS/Cartoon, +25% on CCT/Location 108). These effects directly address a principal limitation of existing approaches: robust, interpretable test-time steering without retraining.

Figure 3: CONTXT parameter sweep visualizing improvements as a function of in-domain injection and OOD removal strengths.

LLMs

The generality of CONTXT is evidenced in LLM applications. Here, persona and sentiment shifts are enforced by constructing context vectors from short, interpretable prompts (e.g., "be extremely positive"). Application of these indices to the activations at early and mid layers enables control over generated sentiment and persona, reliably flipping sentiment predictions in up to 80% of Yelp review test cases without loss in form (self-BLEU preservation).

Figure 4: CONTXT-based sentiment steering enables reliable transfer of sentiment in text outputs with minimal degradation.

Compared to contemporary activation steering and sparse autoencoder (SAE) approaches, CONTXT is notable for (1) its single-token formulation, (2) independence from token-level alignment, and (3) persistence of the steer over long completions. The absence of need for auxiliary model training, backpropagation, or prompt engineering simplifies integration and lowers cost.

Figure 5: Qualitative shift in outputs achieved by CONTXT manipulation of internal representations in Llama-3.

Figure 6: Tabulated example completions and persona shifts achieved by CONTXT application.

Significance and Theoretical Implications

The work establishes that feature space adaptations, governed by compact context vectors, can linearly modulate both discriminative and generative neural processing. Unlike methods requiring data augmentation, domain adversarial training, or test-time optimization, CONTXT delivers robust adaptation with minimal resource consumption, rigorous interpretability, and fine control over context. By drawing from neurobiological models of hippocampal-prefrontal context modulation, the approach links biological principles with tractable interventions in artificial networks.

Empirical results contradict the common expectation that downstream reweighting or linear transforms in feature space are inherently brittle or limited to toy tasks; substantial and interpretable performance gains are demonstrated across heterogeneous architectures and tasks.

Limitations and Future Directions

The method currently presumes knowledge of the test-time context, though results suggest that small, context-representative validation sets suffice for effective adaptation. The authors outline directions for dynamically learning context encodings in an online fashion, leveraging novelty or change-detection modules, echoing hippocampal mechanisms of experience-driven structure learning. Extensions could incorporate context-memory modules with plasticity for on-the-fly context inference, further narrowing the gap between flexible biological and artificial generalization.

Conclusion

The CONTXT method offers a principled, low-overhead approach to context-driven adaptation in neural networks. By steering internal activation patterns with simple, cached context indices, it achieves robust performance under distribution shift without retraining, fine-tuning, or complex engineering. The results underscore the untapped capacity of feature space interventions for both discriminative and generative tasks. The approach suggests a path toward architectures where stable knowledge is decoupled from deployment-specific context, advancing both our theoretical understanding and practical capabilities in robust AI.