Bilevel Learning of Task-Adapted Regularizers

Abstract

In scientific imaging, reconstructing an unknown image from noisy, incomplete measurements is often necessary for downstream tasks (e.g., classification, segmentation). We address this using a task-adapted framework that learns a data-driven regularizer tailored to the specific analysis task. We propose a bilevel learning approach: solving a lower-level variational optimization problem with a parameterized learnable regularizer, while optimizing its parameters by minimizing the upper-level task loss on training data. This facilitates learning a bespoke regularizer for a pre-trained task operator, or jointly learning the parameters of both the regularizer and a deep neural network-based task operator. We introduce a principled algorithm that accommodates inexact lower-level reconstructions (using an iterative first-order solver) and demonstrate its efficacy on challenging inverse problems (image deblurring, CT, and MRI) across classification and segmentation tasks. Our method achieves competitive end-task performance compared to existing methods that use deep networks for both reconstruction and task operators, but does so with a significantly parsimonious regularizer parameterization. Furthermore, our approach requires only raw measurements and task label pairs, eliminating the need for ground-truth reference images.