Rapid Network Adaptation: Learning to Adapt Neural Networks Using Test-Time Feedback

TL;DR

This paper introduces Rapid Network Adaptation (RNA), a fast, flexible test-time adaptation method for neural networks that uses a learned feedback loop to handle distribution shifts effectively.

Contribution

The paper presents RNA, a novel test-time adaptation approach using a learned optimizer, significantly improving speed and flexibility over existing methods.

Findings

RNA is faster than baseline adaptation methods.

RNA effectively adapts across diverse datasets and tasks.

RNA handles various distribution shifts with promising results.

Abstract

We propose a method for adapting neural networks to distribution shifts at test-time. In contrast to training-time robustness mechanisms that attempt to anticipate and counter the shift, we create a closed-loop system and make use of a test-time feedback signal to adapt a network on the fly. We show that this loop can be effectively implemented using a learning-based function, which realizes an amortized optimizer for the network. This leads to an adaptation method, named Rapid Network Adaptation (RNA), that is notably more flexible and orders of magnitude faster than the baselines. Through a broad set of experiments using various adaptation signals and target tasks, we study the efficiency and flexibility of this method. We perform the evaluations using various datasets (Taskonomy, Replica, ScanNet, Hypersim, COCO, ImageNet), tasks (depth, optical flow, semantic segmentation, classification), and distribution shifts (Cross-datasets, 2D and 3D Common Corruptions) with promising results. We end with a discussion on general formulations for handling distribution shifts and our observations from comparing with similar approaches from other domains.