Prediction-Balanced Reservoir Sampling

• The paper introduces PBRS as a novel method for maintaining a prediction-balanced buffer to enable robust continual test-time adaptation on non-i.i.d. streams.
• PBRS employs prediction-balanced insertion and class-conditioned reservoir sampling to effectively mitigate overfitting and class imbalance in dynamic environments.
• Empirical evaluations demonstrate that PBRS outperforms prior TTA methods by significantly lowering error rates across benchmarks such as CIFAR10-C and ImageNet-C.

Prediction-Balanced Reservoir Sampling (PBRS) is an algorithmic method designed to mitigate overfitting and improve generalization in continual test-time adaptation (TTA) under temporally correlated (non-i.i.d.) data streams. Emerging from the NEED-robust Online TESt-time-adaptation (NOTE) framework, PBRS systematically maintains a small, temporal memory buffer of test samples that reflect a class-balanced, nearly i.i.d.-like subsample of the non-i.i.d. test stream, as inferred from the model’s own predictions. This enables robust adaptation of normalization statistics in the presence of severe class-imbalance and temporal correlation (Gong et al., 2022).

1. Motivation and Problem Setting

Continual TTA presupposes a model’s operation under distribution shift, adapting on-the-fly using only the incoming stream of unlabeled test data. Many extant TTA algorithms rely on batch statistics (e.g., for recalibration of BatchNorm), or entropy minimization on each batch. These methods tend to overfit, particularly when the data stream exhibits non-i.i.d. behavior—such as temporally correlated sequences common in real-world scenarios—resulting in class-imbalance and model bias toward transient majority classes. PBRS was introduced to address this by simulating an i.i.d. adaptation buffer through prediction- and time-balanced sample selection (Gong et al., 2022).

2. Algorithmic Structure of PBRS

PBRS maintains a fixed-capacity reservoir $M = \{(\mathbf x_i, \hat y_i)\}_{i=1}^N$ of size $N$, where $\mathbf x_i$ denotes a test sample and $\hat y_i$ its current model-predicted class. For each new test point $(\mathbf x_t)$ with predicted label $\hat y_t$, PBRS applies two interleaved update mechanisms:

• Prediction-Balanced Insertion: If $\hat y_t$ is underrepresented in $M$ relative to other predicted classes, PBRS uniformly selects an instance belonging to the majority (over-represented) class in $M$ and replaces it with $(\mathbf x_t, \hat y_t)$.
• Class-Conditioned Reservoir Sampling: If $N$0 is not underrepresented, reservoir sampling is performed within class $N$1. For class $N$2, the replacement probability for a newly seen sample is $N$3, where $N$4 is the count of class $N$5 in $N$6 and $N$7 is the cumulative count of class $N$8 encountered.

This buffer replacement is performed online with $N$9 per-sample overhead.

3. Mathematical Formulation and Buffer Dynamics

Let $\mathbf x_i$0 denote the running total of test samples with model-predicted class $\mathbf x_i$1 up to time $\mathbf x_i$2, and $\mathbf x_i$3 the count of class $\mathbf x_i$4 instances in $\mathbf x_i$5. The insertion rules are as follows:

• For an incoming test example with $\mathbf x_i$6:
  • Buffer filling phase: If $\mathbf x_i$7, append $\mathbf x_i$8.
  • Prediction-balanced replacement: If $\mathbf x_i$9 is a minority class in $\hat y_i$0, replace a randomly chosen sample with majority label.
  • Class-conditioned sampling: If not, with probability $\hat y_i$1, replace a randomly chosen buffer sample with label $\hat y_i$2.

Mathematically:

$\hat y_i$3

4. Integration with Continual Test-time Adaptation

PBRS operates in tandem with Instance-Aware Batch Normalization (IABN). After every $\hat y_i$4 insertions, the buffer $\hat y_i$5 is used to recompute BatchNorm statistics and update the affine parameters $\hat y_i$6 via a single backward adaptation pass. The global mean $\hat y_i$7 and variance $\hat y_i$8 are updated using exponential moving averages:

$\hat y_i$9

with momentum $(\mathbf x_t)$0. Here, $(\mathbf x_t)$1 and $(\mathbf x_t)$2 are calculated from activations of the $(\mathbf x_t)$3 buffered samples. Only $(\mathbf x_t)$4 are optimized through Adam (learning rate $(\mathbf x_t)$5). The buffer size $(\mathbf x_t)$6 matches a common mini-batch size.

5. Empirical Performance Evaluation

PBRS, solely in conjunction with IABN, was evaluated across multiple benchmarks reflecting severe temporal class-imbalance and real data streams. In the non-i.i.d. setting, mean error rates achieved by NOTE (IABN + PBRS) were:

Dataset	NOTE (IABN+PBRS)	Best Prior Baseline
CIFAR10-C	21.1%	36.2% (LAME)
CIFAR100-C	47.0%	63.3% (LAME)
ImageNet-C	80.6%	82.7%
KITTI-Rain	10.9%	11.3%
HARTH	51.0%	61.0%
ExtraSensory	45.4%	50.7%

NOTE outperforms all other TTA methods (BN-Stats, ONDA, PL, TENT, LAME, CoTTA) in non-i.i.d. streams, and matches or surpasses them when the i.i.d. assumption holds (Gong et al., 2022).

6. Theoretical and Empirical Properties

PBRS does not offer formal unbiasedness proofs but empirically maintains class frequencies in the buffer $(\mathbf x_t)$7 close to the long-term average as predicted by the model. The class-conditioned reservoir sampling ensures per-class time-uniform sampling, and the prediction-balanced policy prevents domination by any class under severe drift. Ablation studies demonstrate near-uniform class distribution in $(\mathbf x_t)$8 even under pronounced temporal skews, underpinning robust adaptation dynamics (Gong et al., 2022).

7. Implementation Details and Practical Considerations

Key parameters include buffer size $(\mathbf x_t)$9, BatchNorm EMA momentum $\hat y_t$0, and Adam learning rate $\hat y_t$1 for adaptation. Storage requirements are minimal: only $\hat y_t$2 pairs, with computational cost $\hat y_t$3 per sample for buffer management, and $\hat y_t$4 per adaptation step (triggered every $\hat y_t$5 samples). PBRS is always paired with IABN for maximum robustness in the non-i.i.d. paradigm (Gong et al., 2022).