Li Ju: Robust Learning from Distributed and Heterogeneous Data

Abstract

Modern machine learning is increasingly expanding beyond centralized, mono-modal training toward systems that must also learn from data across distributed edge devices and heterogeneous data modalities. This transition breaks the foundational identical and independent distribution (i.i.d.) assumptions of traditional models, making robustness a first-class requirement for real-world applications. This thesis studies the mechanisms and methodologies necessary to achieve algorithmic robustness across three intersecting dimensions: distributed optimization, geometry-aware uncertainty quantification, and simulation-based inference.

The first dimension addresses statistical heterogeneity in Federated Learning (FL), a distributed training framework in which multiple participants collaboratively train a shared model without exchanging their local data. In FL, the non-i.i.d. nature of distributed data often induces performance degradation, convergence issues and fairness problems. Through an empirical study on drug discovery and the development of new algorithms, this work demonstrates that adaptive optimization and dynamic hyperparameter adjustment can mitigate training instabilities. These methods ensure equitable performance across diverse data silos, preventing the global model from favoring specific participants.

The second dimension explores the structural challenges of multi-modal language models, which map data of heterogeneous modalities onto complex, non-Euclidean manifolds. This research models aleatoric and epistemic uncertainty with directional distributions via parametric models and Riemannian Flow Matching. This geometry-aware approach allows models to respect the intrinsic geometric structure of the embedding space, providing a mathematically grounded framework for models to quantify their ignorance when confronted with ambiguous or out-of-distribution inputs.

The final dimension addresses the robustness of a unified framework which supports both forward and inverse processes for Bayesian inference. The proposed framework utilizes a unified Flow Matching model to learn the joint distribution of parameters and observations. By employing randomized masking, this architecture robustly handles partially observed or noisy data, integrating forward and inverse processes into a single cohesive neural network without the need for specialized retraining.  Collectively, this thesis contributes theoretical analyses, novel algorithms, and empirical validations that advance the robustness of machine learning across federated optimization, multi-modal uncertainty quantification, and simulation-based inference, bridging the gap between idealized training assumptions and the demands of real-world applications.