TRAM: Bridging Trust Regions and Sharpness Aware Minimization

TL;DR

TRAM combines trust-region and sharpness-aware minimization techniques to enhance out-of-domain generalization in vision and language tasks by optimizing both parameter and function space curvature.

Contribution

It introduces a novel fine-tuning algorithm, TRAM, that unifies trust-region and SAM strategies to improve domain transfer and representation robustness.

Findings

TRAM outperforms existing SAM and trust-region methods across vision and language tasks.

TRAM achieves superior domain transfer, especially in challenging anticorrelated domain scenarios.

Minimal additional computation is required compared to previous sharpness-aware methods.

Abstract

Sharpness-aware minimization (SAM) reports improving domain generalization by reducing the loss surface curvature in the parameter space. However, generalization during fine-tuning is often more dependent on the transferability of representations in the function space. Trust-region methods (TR) target this goal by regularizing representation curvature to reduce catastrophic forgetting of pre-trained task-agnostic information while adopting task-specific skills. We consider unifying these strategies for low curvature in both parameter space and function space to improve out-of-domain (OOD) generalization. We propose Trust Region Aware Minimization (TRAM), a SAM algorithm fine-tuning for low parameter sharpness and smooth, informative representations preserving pre-trained structure. TRAM uses a trust region bound to inform the SAM adversarial neighborhood, introducing an awareness of function curvature within optimization for flatter minima. We empirically validate TRAM in vision (cross-dataset adaptation) and text (OOD language modeling, zero-shot cross-lingual transfer) tasks where robust domain transfer and representation generality are critical. TRAM outperforms SAM- and TR-based optimization across all tasks, notably surpassing competing methods for hard transfer between anticorrelated domains. TRAM establishes a novel standard in fine-tuning for domain-generalizable models with minimal additional computation over previous sharpness-aware methods.