FIRM: Federated In-client Regularized Multi-objective Alignment for Large Language Models

TL;DR

FIRM is a federated learning algorithm that efficiently aligns large language models with human values by reducing communication overhead and balancing multiple objectives through in-client regularization.

Contribution

FIRM introduces a novel in-client regularized multi-objective optimization method that eliminates multi-gradient communication, ensuring convergence and improved trade-offs in federated LLM alignment.

Findings

Reduces communication by transmitting only adapted parameters

Achieves convergence to Pareto-stationary points with finite-time guarantees

Improves training stability and objective trade-offs in experiments

Abstract

Aligning Large Language Models (LLMs) with human values often involves balancing multiple, conflicting objectives such as helpfulness and harmlessness. Training these models is computationally intensive, and centralizing the process raises significant data privacy concerns. Federated Learning (FL) offers a compelling alternative, but existing Federated Multi-Objective Optimization (FMOO) methods face severe communication bottlenecks as their reliance on transmitting multiple gradients to a server is unscalable for large models. We introduce FIRM (Federated In-client Regularized Multi-objective alignment), a novel algorithm that achieves both client disagreement drift mitigation and communication efficiency. In FIRM, each client locally solves a regularized multi-objective optimization problem. By directly mitigating client disagreement drift through in-client regularization, our method eliminates the need for the multi-gradient transmissions common in prior works. Consequently, clients need only to transmit a single set of adapted parameters, maintaining high communication efficiency. We prove that our algorithm converges to Pareto-stationary points and, to our knowledge, provide the first finite-time convergence guarantees for this federated multi-objective alignment setting. Empirically, we show that FIRM leads to smoother training dynamics, reduced client disagreement drift, and improved reward trade-offs compared to baselines. We further propose a method to incorporate a preference over the objectives and report empirical Pareto plots, demonstrating that FIRM can smoothly adapt trade-offs between objectives in response to specified preferences.