FoundIR-v2: Optimizing Pre-Training Data Mixtures for Image Restoration Foundation Model
This work addresses the challenge of improving all-in-one image restoration models for real-world scenarios by optimizing data mixtures, representing an incremental advancement in foundation model training.
The paper tackles the problem of optimizing pre-training data mixtures for image restoration foundation models, finding that data mixture proportions from different tasks critically affect performance, and proposes FoundIR-v2 with a data equilibrium scheduling paradigm to dynamically optimize these proportions, achieving favorable performance against state-of-the-art approaches across over 50 sub-tasks.
Recent studies have witnessed significant advances in image restoration foundation models driven by improvements in the scale and quality of pre-training data. In this work, we find that the data mixture proportions from different restoration tasks are also a critical factor directly determining the overall performance of all-in-one image restoration models. To this end, we propose a high-capacity diffusion-based image restoration foundation model, FoundIR-v2, which adopts a data equilibrium scheduling paradigm to dynamically optimize the proportions of mixed training datasets from different tasks. By leveraging the data mixing law, our method ensures a balanced dataset composition, enabling the model to achieve consistent generalization and comprehensive performance across diverse tasks. Furthermore, we introduce an effective Mixture-of-Experts (MoE)-driven scheduler into generative pre-training to flexibly allocate task-adaptive diffusion priors for each restoration task, accounting for the distinct degradation forms and levels exhibited by different tasks. Extensive experiments demonstrate that our method can address over 50 sub-tasks across a broader scope of real-world scenarios and achieves favorable performance against state-of-the-art approaches.