One Pool Is Not Enough: Multi-Cluster Memory for Practical Test-Time Adaptation

Test-time adaptation (TTA) adapts pre-trained models to distribution shifts at inference using only unlabeled test data. Under the Practical TTA (PTTA) setting, where test streams are temporally correlated and non-i.i.d., memory has become an indispensable component for stable adaptation, yet existing methods universally store amples in a single unstructured pool. We show that this single-cluster design is fundamentally mismatched to PTTA: a stream clusterability analysis reveals that test streams are inherently multi-modal, with the optimal number of mixture components consistently far exceeding one. To close this structural gap, we propose Multi-Cluster Memory (MCM), a plug-and-play framework that organizes stored samples into multiple clusters using lightweight pixel-level statistical descriptors. MCM introduces three complementary mechanisms: descriptor-based cluster assignment to capture distinct distributional modes, Adjacent Cluster Consolidation (ACC) to bound memory usage by merging the most similar temporally adjacent clusters, and Uniform Cluster Retrieval (UCR) to ensure balanced supervision across all modes during adaptation. Integrated with three contemporary TTA methods on CIFAR-10-C, CIFAR-100-C, ImageNet-C, and DomainNet, MCM achieves consistent improvements across all 12 configurations, with gains up to 5.00% on ImageNet-C and 12.13% on DomainNet. Notably, these gains scale with distributional complexity: larger label spaces with greater multi-modality benefit most from multi-cluster organization. GMM-based memory diagnostics further confirm that MCM maintains near-optimal distributional balance, entropy, and mode coverage, whereas single-cluster memory exhibits persistent imbalance and progressive mode loss. These results establish memory organization as a key design axis for practical test-time adaptation.

💡 Research Summary

Test‑time adaptation (TTA) seeks to adjust a pretrained model at inference time using only unlabeled test data, thereby avoiding costly retraining for every new domain shift. While early TTA work assumed a single, static target domain, more recent efforts have introduced continual TTA (CTTA) and, most recently, practical TTA (PTTA), which models the realistic scenario where test samples arrive as a temporally correlated, non‑i.i.d. stream. In PTTA, each mini‑batch covers only a narrow slice of the evolving target distribution, making a memory buffer essential for stable adaptation. Existing memory‑based TTA methods, however, store samples in a single unstructured pool (single‑cluster memory, SCM), implicitly assuming that the target distribution is unimodal.

The authors first demonstrate that this assumption is fundamentally flawed. By sliding a window over PTTA streams from CIFAR‑100‑C, ImageNet‑C, and other benchmarks, they extract three lightweight pixel‑level descriptors (channel‑wise mean/variance, spatial mean, and color histograms) and fit Gaussian mixture models (GMMs) with varying numbers of components K. Bayesian Information Criterion (BIC) consistently selects K* values between 5 and 10, indicating that even within a single corruption type the data are intrinsically multimodal. This empirical evidence motivates a redesign of the memory architecture.

Multi‑Cluster Memory (MCM) is introduced as a plug‑and‑play framework that organizes stored samples into up to Kmax clusters, each representing a distinct mode of the target distribution. The core of MCM consists of three complementary mechanisms:

1. Descriptor‑based Cluster Assignment – Each incoming sample is summarized by a channel‑wise mean‑variance vector dₓ. The Euclidean distance between dₓ and all existing cluster centroids Dₖ is computed; the sample is assigned to the nearest cluster if the distance is below a threshold τ, otherwise a new cluster is spawned. τ controls granularity: smaller values yield finer mode separation, while larger values allow broader clusters. Because the descriptor scale is bounded across datasets, τ can be set once per descriptor type without per‑dataset tuning.

2. Adjacent Cluster Consolidation (ACC) – When the number of clusters reaches Kmax, MCM merges the most similar pair of adjacent clusters (i.e., clusters created consecutively in time). The adjacency constraint reduces the search from O(K²) to O(K) and leverages the temporal continuity of PTTA streams: neighboring clusters are likely to belong to the same or smoothly transitioning domain, so merging them minimally disturbs overall mode coverage. The merged cluster retains the N samples with the lowest prediction uncertainty, preserving quality while respecting the per‑cluster capacity N.

3. Uniform Cluster Retrieval (UCR) – During adaptation, MCM draws an equal number of samples from each cluster to form the adaptation mini‑batch. This balanced sampling prevents the model from over‑fitting to any single dominant mode and ensures that the loss is computed on a representation of the full target manifold.

When a cluster reaches its capacity N, MCM replaces the least valuable sample using a scoring function that extends the RoTTA heuristic (age and entropy) with an additional term proportional to the descriptor distance from the cluster centroid. This encourages intra‑cluster compactness while still favoring recent, uncertain samples.

The authors integrate MCM with three state‑of‑the‑art memory‑based TTA methods—RoTTA, PeTTA, and ResiTTA—and evaluate on four benchmark suites (CIFAR‑10‑C, CIFAR‑100‑C, ImageNet‑C, DomainNet) under the PTTA protocol. Across all 12 baseline‑dataset configurations, MCM yields consistent improvements, achieving an average error reduction of 2.96 %. The most pronounced gains appear on datasets with larger label spaces and higher intrinsic multimodality: up to 5.00 % absolute reduction on ImageNet‑C and 12.13 % on DomainNet.

To quantitatively assess memory quality, the authors propose a GMM‑based diagnostic framework that measures three statistics on the stored samples: (i) imbalance ratio (how evenly samples are distributed across clusters), (ii) entropy of the cluster‑level descriptor distribution, and (iii) mode coverage (the proportion of GMM components adequately represented). MCM consistently attains near‑optimal values on all three metrics, whereas SCM exhibits persistent skew, lower entropy, and periodic loss of modes as adaptation proceeds. Ablation studies confirm that each component of MCM contributes meaningfully: removing ACC leads to memory overflow and degraded mode coverage; disabling UCR causes performance collapse on highly multimodal streams; and setting τ too low or too high harms either memory efficiency or mode discrimination.

The paper’s contributions are threefold: (1) providing the first systematic evidence that PTTA streams are multimodal, thereby justifying a shift away from single‑pool memory; (2) proposing a lightweight, descriptor‑driven multi‑cluster memory architecture with explicit mechanisms for assignment, consolidation, and balanced retrieval; (3) introducing a principled diagnostic toolkit for evaluating memory representativeness, and demonstrating that performance gains stem from improved distributional coverage rather than mere increases in capacity.

Limitations and future directions are acknowledged. The current centroid is a simple arithmetic mean of channel statistics, which may be insufficient for highly non‑Gaussian or non‑linear mode structures. Incorporating richer feature embeddings (e.g., intermediate CNN activations) could capture subtler domain shifts at the cost of higher computation. Moreover, the threshold τ is fixed a priori; adaptive τ scheduling based on streaming statistics could further improve flexibility. Finally, extending MCM to jointly optimize memory organization and model parameters (e.g., via bi‑level optimization) is an exciting avenue for research.

In summary, Multi‑Cluster Memory reframes memory design from a flat buffer into a structured, mode‑aware repository, enabling test‑time adaptation systems to faithfully track and exploit the multimodal nature of real‑world data streams. The framework is simple, computationally efficient, and compatible with existing TTA pipelines, making it a practical step toward robust deployment of deep models under continual distribution shift.