Multi-Expert Adaptive Selection: Task-Balancing for All-in-One Image Restoration

TL;DR

The paper tackles multi-task image restoration by introducing a multi-expert adaptive selection mechanism within an All-in-One framework. It combines a pixel-level, task-prompt-guided expert routing (STP-G-MESE) with a global, frequency-aware feature decomposition and ensemble (FD-MEE), enabling shared experts across tasks while respecting task-specific demands. The method employs task-specific prompts derived from degradation-aware cues and learns to balance expert utilization across tasks, achieving superior restoration accuracy and robust generalization across three and five degradation settings. Experimental results show consistent improvements over state-of-the-art All-in-One approaches, demonstrating practical potential for joint restoration under diverse degradation conditions.

Abstract

The use of a single image restoration framework to achieve multi-task image restoration has garnered significant attention from researchers. However, several practical challenges remain, including meeting the specific and simultaneous demands of different tasks, balancing relationships between tasks, and effectively utilizing task correlations in model design. To address these challenges, this paper explores a multi-expert adaptive selection mechanism. We begin by designing a feature representation method that accounts for both the pixel channel level and the global level, encompassing low-frequency and high-frequency components of the image. Based on this method, we construct a multi-expert selection and ensemble scheme. This scheme adaptively selects the most suitable expert from the expert library according to the content of the input image and the prompts of the current task. It not only meets the individualized needs of different tasks but also achieves balance and optimization across tasks. By sharing experts, our design promotes interconnections between different tasks, thereby enhancing overall performance and resource utilization. Additionally, the multi-expert mechanism effectively eliminates irrelevant experts, reducing interference from them and further improving the effectiveness and accuracy of image restoration. Experimental results demonstrate that our proposed method is both effective and superior to existing approaches, highlighting its potential for practical applications in multi-task image restoration.

Figures (9)

• Figure 1: Comparison between the traditional Mixture of Experts (MoE) 88 model and our proposed method. In MoE, the input is directly assigned to different experts for processing via a routing mechanism. The proposed method comprehensively considers the features at the pixel level as well as the global features composed of low-frequency and high-frequency components, and makes expert selection based on this.
• Figure 2: Overview of the proposed method: The input degraded image is first processed by TSPG, and task-related prompts are generated with the assistance of TRB Prompt. Subsequently, effective integration of task and content-related information is achieved through TCI. Based on the information integrated by TCI, we predict the importance of each expert in the expert database for the current sample pixel restoration task and select the top-$K$ most important experts to handle each pixel in the image. In the stage of feature decomposition and multi-expert ensemble, the FD separates the features output by the Transformer layer into high-frequency and low-frequency components, allowing specific experts to restore the image information as a whole for different frequency components.
• Figure 3: Framework structure diagram of TCI.
• Figure 4: Influence of different degradation types on image high- and low-frequency components.
• Figure 5: Framework structure diagram of FD.