Model Optimization for Federated Learning and Flow-Based Image Restoration

Abstract

This thesis explores model optimization strategies for two fundamental areas in machine learning: Federated Learning (FL) and Image Restoration (IR), both of which must address challenges posed by data heterogeneity and distribution shifts. We present three contributions aimed at improving robustness, adaptability, and performance in these settings.

The first chapter introduces a gradient-based client selection method for FL. We propose a novel $\ell_4$-norm cosine similarity metric that captures higher-order gradient structures, allowing the server to prioritize clients whose updates are more aligned and informative. This approach accelerates convergence and improves the final model quality compared to random or traditional $\ell_2$-based selection strategies, especially under non-i.i.d. client distributions.

The second chapter presents MoiréXNet, a multi-scale image restoration network designed to remove complex visual distortions such as moiré patterns. Our framework integrates linear attention modules for efficient feature aggregation, test-time training for adaptation to unseen degradations, and a truncated flow matching prior to enforce structural consistency. MoiréXNet achieves state-of-the-art performance across several real-world benchmarks.

The third chapter addresses the "Last Mile" of image restoration through a rectified flow-based refinement process. We design a two-stage restoration framework: a coarse estimate is first optimized for distortion-oriented metrics, followed by refinement using generative model-based methods that learn efficient distribution mappings to enhance perceptual fidelity. This strategy balances distortion reduction and perceptual quality, producing visually realistic results even under severe degradation.

Collectively, this thesis advances gradient-based optimization for federated systems and flow-guided adaptive restoration methods, contributing to the development of AI models that are robust, efficient, and capable of handling messy, unpredictable real-world data.