Learning Many Tasks via Weight-Space Geometry

Deep neural networks excel due to their ability to learn from large datasets. As their deployment expands, challenges arise in adapting to new objectives, controlling conflicting behaviors, and enabling modular reuse without retraining. This thesis explores these issues from the perspective of the geometry of the weight space, developing methods for multi-objective optimization, efficient model merging, and editing.

We begin our investigation by exploring the problem of learning Pareto-optimal solutions in settings with multiple competing objectives. Traditional multi-task learning methods yield a single solution, failing to capture the spectrum of optimal trade-offs inherent in multi-objective problems. To overcome this, we extend the concept of linear mode connectivity to the multi-task setting and propose a method that models the Pareto front as an ensemble of single-task models, enabling the learning of diverse solutions conditioned on task preferences. Building on this foundation, we mitigate the scalability constraints and convergence issues associated with weight ensembles when dealing with a large number of tasks or high-capacity models by introducing task-specific low-rank adapters integrated into a shared backbone network, facilitating the learning of a continuous Pareto front within the convex hull of these adapters.

The thesis then transitions from modeling the Pareto Front to post-training model editing in the weight space. We explore the performance degradation caused by task arithmetic and propose a training-free algorithm that constructs a binary mask per task, localizing the important information on the multi-task vector and allowing near-complete recovery of the original performance. We then employ the masks to (i) heavily compress the original checkpoints and (ii) improve model merging algorithms by studying the mask agreement profile across tasks. Extending our exploration of post-training interventions, we examine the weight-space structure induced by fine-tuning and its role in catastrophic forgetting. We edit the fine-tuned model by rescaling the parameter updates as a linear function of layer depth: updates in shallow layers are shrunk toward their pre-trained values, while deeper layers retain more of their task-specific changes. This strategy not only recovers forgotten capabilities but also improves compatibility in model merging scenarios, both in single and multi-task settings as well as merging LLM policies aligned with different reward signalsâ a setting analogous to Pareto front learning. We then expand model merging to scenarios where several tasks operate on the same input and produce distinct outputs, departing from the standard setting with a one-to-one mapping from input to task. This regime introduces a new challenge: a misalignment between the merged encoder and the task-specific decoders. We address this via post-merging adjustments that realign the encoder's outputs with the expectations of each decode