White-Box Prompt Transformers: Variationally Grounded Prompt–Attention Coupling for Unified Image Restoration

ICLR 2026 Conference Submission16419 Authors

Keywords: Interpretable Image Restoration

Abstract:

Can soft prompts in vision Transformers be made explainable?
Prompt-based models have achieved remarkable success in image restoration, yet they remain largely opaque: the underlying Transformer operations and the mechanism by which prompts modulate attention are poorly understood. This work revisits guided image restoration, where an auxiliary modality \(A\) assists in restoring a target modality \(B\). We interpret \(A\) as a prompt and formulate a tailored structure-tensor total variation (STV) model, whose gradient suggests a white-box correspondence to prompt--attention interactions. This provides a principled bridge between prompts and attention. In scenarios where \(A\) is unavailable, we abstract its role into learnable soft prompts, enabling end-to-end training within standard Transformer pipelines. By unrolling the gradient flow of the STV variational problem, we derive the White-Box Prompt Transformer (WBPT), a cascaded architecture that embeds interpretability directly into attention operations. Extensive experiments on multiple benchmarks demonstrate that WBPT achieves state-of-the-art restoration performance while offering interpretable, controllable, and robust prompt--attention dynamics.

Primary Area: interpretability and explainable AI

Submission Number: 16419