When Prompts Interact: Assessing Prompt Arithmetic for Deconfounding under Distribution Shift

TL;DR

This paper introduces Hybrid Prompt Arithmetic (HyPA), a parameter-efficient method that combines prompt tuning and task arithmetic to improve model robustness against confounding shifts in classification tasks.

Contribution

It proposes HyPA, a novel approach that enhances robustness to distribution shifts by combining prompt and confounder prompts, with analysis of its effects on model representations.

Findings

HyPA consistently improves robustness-performance trade-off under distribution shift.

HyPA reduces the influence of confounder signals in model predictions.

HyPA is a parameter-efficient method that mitigates reliance on spurious features.

Abstract

In classification tasks, models may rely on confounding variables to achieve strong in-distribution performance, capturing spurious features that fail under distribution shift. This shortcut behavior leads to substantial degradation in out-of-distribution settings. Task arithmetic offers a potential solution by removing unwanted signals via subtraction of secondary model updates, but it typically requires full fine-tuning, which is computationally expensive. Prompt tuning provides a parameter-efficient alternative by adapting models through a small set of trainable virtual tokens. Task arithmetic on the resulting prompts presents an appealing alternative to operations on entire models, but the extent to which this approach can limit reliance on spurious features remains to be established. In this work, we study whether composing soft prompts through task arithmetic improves robustness to confounding shifts. We propose Hybrid Prompt Arithmetic (HyPA), which combines task prompts with linearized confounder prompts to counteract spurious correlations. Across multiple benchmarks, HyPA consistently improves the robustness-performance trade-off relative to prompt-arithmetic baselines under distribution shift. We further analyze how HyPA affects hidden representations and find evidence consistent with it mitigating confounding either by reducing the influence of confounder signals on predictions or by suppressing them in the representation. These results establish HyPA as a parameter-efficient and promising approach for improving robustness under confounding shifts in the evaluated setting.