Curvature-Aware Zeroth-Order Optimization for Memory-Efficient Test-Time Adaptation

AI summary

Overview

• Research area: Test-time adaptation (TTA) for computer vision models, specifically backpropagation-free (zeroth-order) optimization for on-device deployment.
• Technical level: Advanced. The paper combines empirical Hessian spectral analysis, a preconditioned zeroth-order gradient estimator, and a nonconvex convergence proof.
• Scope: The paper diagnoses why standard zeroth-order (ZO) gradient estimation is too noisy for TTA and proposes a curvature-aware variant, CAZO, that is evaluated on ImageNet-C, ImageNet-R, ImageNet-V2 and ImageNet-Sketch with a ViT-B/16 backbone.

What This Paper Is About

Test-time adaptation lets a pretrained model adjust to shifted, unlabeled test data, but most methods need backpropagation (BP), which is memory-heavy on edge devices. BP-free zeroth-order methods only need forward passes, but they are notorious for high-variance gradient estimates, so naively using them for TTA is slow and inaccurate. The paper's goal is to make ZO-based TTA both accurate and memory-efficient by shaping its random perturbations according to the curvature of the loss landscape.

Key Contributions

1. An empirical characterization of loss curvature during TTA. The authors compute the Hessian of the loss with respect to adapter parameters at several adaptation steps and show it has a persistent low-rank structure whose principal subspace changes slowly over time.
2. The CAZO method. A forward-pass-only optimizer that uses a sliding exponential moving average (EMA) of the element-wise squared ZO gradients as a diagonal curvature proxy, then samples perturbations from a Gaussian whose covariance is the inverse of that proxy — attenuating perturbations in sharp directions and amplifying them in flat ones.
3. Convergence guarantees. A bias/variance bound for the estimator (Lemma 1) and an O(1/sqrt(T)) nonconvex convergence rate (Theorem 1) whose constants depend explicitly on the curvature bounds.
4. Extensive evaluation. State-of-the-art results across ImageNet-C (severity 5, both with model reset and in continual adaptation), ImageNet-R/V2/Sketch, and 8-bit and 6-bit quantized models, with memory measured against BP-based baselines.

Main Findings

• Hessian is low-rank and persistent. For a pretrained ViT-B/16 adapted on ImageNet-C Gaussian corruption at severity level 5, measured at adaptation steps {0, 25, 50, 99} out of 100 total iterations, the top 20 eigenvalues account for more than 96% of the total variance, and the effective rank occupies only 0.22% of total parameter dimensionality.
• The principal curvature subspace drifts slowly. With a batch size of 512 and 25 adaptation steps, the projection ratio between adjacent Hessians stays around 0.9 across adaptation for ranks r ∈ {5, 10, 15, 20}. Larger batch sizes produce a smoother trajectory, which the authors take as evidence that the expected loss landscape also has low-rank, slow-varying structure.
• Standard ZO pays a variance penalty. The variance of the random gradient estimation (RGE) estimator scales as O(d/k), and naive ZO stochastic gradient descent may need O(d) more iterations than first-order methods to reach comparable accuracy.
• ImageNet-C (severity 5, ViT-B/16, model reset per corruption). CAZO reaches 69.0 ± 0.1% average accuracy. For comparison: NoAdapt 55.5%, LAME 54.1%, T3A 56.9%, FOA 65.8%, ZOA 67.5%, TENT 59.8%, CoTTA 61.9%, SAR 62.7%, DeYO 64.7%, EATA 66.8%, RoTTA 56.3%. That is +3.2% over FOA, +1.5% over ZOA, +6.3% over SAR and +7.1% over CoTTA, all without any backward pass.
• Continual adaptation (no model reset). CAZO again leads with 65.3% average accuracy, ahead of LCoTTA (62.3%), ETA (61.7%) and SAR (61.6%) by +3.0, +3.6 and +3.7 points respectively. Reported standard deviation for CAZO in this setting is ±1.0.
• Other domain-shift datasets. On ImageNet-R / V2 / Sketch with ViT-B/16, CAZO scores 64.4 / 75.3 / 50.9 for an average of 63.5, versus DeYO 63.8, CoTTA 63.0, FOA 62.7 and NoAdapt 59.9. The authors describe this as competitive rather than a clean sweep.
• Memory efficiency. Under full precision, TENT requires 6,404 MB of runtime CUDA memory while CAZO requires only 1,695 MB; CoTTA requires 17,773 MB and RoTTA 9,126 MB. The paper reports an overall 4–10× reduction versus BP-based methods and "roughly 70%" less memory, attributable to diagonal curvature estimation, lightweight adapter parameterization and memory management.
• Memory is nearly flat in the number of perturbations. Across k ∈ {2, 4, 8, 20}, CAZO memory stays at 1,693 / 1,693 / 1,695 / 1,695 MB.
• Accuracy and calibration improve with k. CAZO with k = 2, 4, 8, 20 gives 65.2%, 67.2%, 67.9% and 69.0% accuracy with ECE of 6.1, 5.4, 4.7 and 4.3. The plain ZO baseline at k = 20 reaches 62.9% with ECE 5.6.
• Runtime trade-off. CAZO with k = 20 achieves 69.0% in 3,127 seconds, versus FOA's 2,885 seconds at 65.8% accuracy. Measurements are on 50,000 ImageNet-C samples using a single NVIDIA H20 GPU.
• Quantized models. In the 8-bit setting CAZO reaches 67.8% average accuracy, above ZO (61.8%), FOA (63.3%), T3A (55.7%), NoAdapt (54.1%) and LAME (52.4%). In the 6-bit setting CAZO retains 61.2%, ahead of ZO (55.6%), FOA (54.8%), T3A (48.8%) and NoAdapt (47.6%).
• Ablation: adapter placement. Inserting the adapter at layer 3 of the 12-layer ViT-B/16 encoder gives peak accuracy; later layers perform worse. Layer 3 is the default.
• Ablation: adapter size. A down-sampling ratio of 384 works best; smaller ratios such as 96 and 128 add many more trainable parameters without gains and can hurt.
• Ablation: perturbation count and EMA. Going from k = 2 to k = 6 raises accuracy from 59.3% to 62.1% with a sharp ECE drop, after which gains flatten; k = 20 is the default. Single-point perturbation degraded performance, so symmetric multi-point perturbation is used. For the EMA coefficient, ν = 0.8 gives the best accuracy (69.0%) and calibration, while ν = 1.0 causes instability.

Methodology in Plain English

The authors start by asking a simple diagnostic question: when a model adapts at test time, what does the shape of its loss surface look like? They compute the Hessian — the matrix of second derivatives describing local curvature — for the small set of adapter parameters being updated, and find that only a handful of directions really matter and that those directions barely move as adaptation proceeds.

Standard ZO methods ignore this and perturb parameters in random, equally weighted directions. That wastes evaluations on nearly flat directions while under-sampling the sharp ones. CAZO instead builds a cheap diagonal approximation of the curvature, tracks it with an exponential moving average of squared gradient estimates (using the bias-corrected form D_t / (1 − (1 − ν)^t)), and samples perturbations from a Gaussian with that inverse diagonal as covariance. Concretely, each iteration samples K perturbations, evaluates the loss at θ + εu and θ − εu for each, forms the finite-difference gradient estimate, updates the diagonal curvature estimate, and takes an SGD step on the adapter parameters only. The pretrained backbone weights stay frozen. The training objective combines an unsupervised entropy loss on test data with an MSE feature-alignment loss that uses feature statistics from clean data. The default adapter sits at layer 3 of the ViT-B/16 encoder.

The theory side formalizes this: under L-smoothness, a data-variance bound σ², and the assumption that entries of the inverse curvature matrix lie in [β_l, β_u], Lemma 1 bounds the estimator's first and second moments, and Theorem 1 gives an O(1/sqrt(T)) rate with learning rate η = β_l / (2 L d(d+2) β_u sqrt(T)). The authors note that when the curvature proxy is well conditioned — for instance β_l² > β_u — the constant factor is smaller than for isotropic sampling.

Why This Matters

• Research impact. The paper reframes ZO variance reduction for TTA as a curvature-sampling problem, connecting Hessian spectrum analysis to perturbation design. It also reports that under a similar memory budget it outperforms the concurrent ZO TTA method ZOA, which targets quantized models and long-term domain-knowledge management rather than curvature-aware sampling.
• Real-world applications:
  • On-device vision on edge hardware with constrained GPU memory, where a 1,695 MB footprint versus 6,404 MB for a BP baseline changes what fits.
  • Low-bit deployed models, since CAZO retains 67.8% (8-bit) and 61.2% (6-bit) accuracy where other methods degrade sharply.
  • Continuously deployed sensors and cameras facing accumulating distribution shift, where the no-reset continual setting matters and CAZO reaches 65.3%.
  • Scenarios where backpropagation is unavailable or undesirable, such as hardware or software stacks that expose only forward inference.
• Industry relevance. Memory, not raw accuracy, is often the binding constraint on shipping adaptation to edge devices. The paper's framing of a 4–10× memory reduction with accuracy that beats BP-based alternatives speaks directly to deployment budgets, and the reported flat memory scaling with k gives practitioners a knob to trade runtime for accuracy without paying in memory.

Future Directions

• Beyond ViT-B/16. All main experiments use ViT-B/16; whether the low-rank, slow-varying Hessian observation holds for convolutional or other architectures is not reported.
• Beyond a single adapter layer. The ablation fixes the adapter at layer 3; adapting multiple layers or using different placement strategies is left open.
• Richer curvature proxies. The method uses a diagonal approximation of the inverse Hessian for tractability. Full-rank or block-diagonal approximations, and their memory/accuracy trade-offs, are not explored.
• Longer-horizon and broader shift regimes. The continual setting traverses the 15 ImageNet-C corruptions once; behavior under much longer horizons, and comparison against the concurrent ZOA method's domain-knowledge management, are natural extensions.

Target Audience

Researchers and practitioners working on test-time adaptation, backpropagation-free or zeroth-order optimization, and memory-constrained on-device deployment. It is also relevant to readers interested in loss-landscape geometry and how Hessian structure can be exploited for variance reduction. Readers without background in ZO gradient estimation or Hessian analysis will find the methodology section the most demanding part; the ablation and efficiency tables are accessible to a broader engineering audience.

Authors’ abstract

Test-time adaptation (TTA) aims to enhance the cross-domain performance of pre-trained models by adapting to unlabeled test data. While most existing TTA methods rely on backpropagation (BP) for finetuning, BP-free methods such as zeroth-order (ZO) methods are more desired in practical on-device scenarios. ZO methods rely only on forward computation, which can largely reduce the complexity and memory overhead of on-device deployment. However, ZO methods suffer from much higher variance compared with first-order methods in estimating the gradient. To address this, we propose an improved ZO method to substantially boost the performance of ZO optimization based TTA. First, we provide an observation to reveal the persistent low-rank Hessian structure of the loss during the adaptation process. Based on this insight, we then propose a loss-landscape curvature-aware zeroth-order (CAZO) method, which leverages a sliding-average estimation of the diagonal Hessian to construct a covariance matrix for anisotropic perturbation sampling. CAZO operates by freezing pretrained weights and optimizing minimal adapter parameters via forward-only passes based gradient estimation, which can substantially reduce the memory overhead compared to BP-based methods. Extensive experiments demonstrate that CAZO significantly outperforms existing TTA methods, achieving state-of-the-art performance while maintaining an excellent balance between accuracy and memory efficiency. Code is available at https://github.com/Hollyming/CAZO.