FedMOA: Federated GRPO for Personalized Reasoning LLMs under Heterogeneous Rewards

TL;DR

FedMOA extends Group Relative Policy Optimization to a federated setting to achieve personalized reasoning in LLMs under heterogeneous rewards. It introduces client-side adaptive hypergradient weighting to dynamically balance multiple objectives and a server-side task-aware aggregation that prioritizes updates demonstrating progress on the primary accuracy objective. Across math and code benchmarks, FedMOA consistently outperforms naive FedGRPO, delivering improvements in global accuracy and multi-objective balance, especially under non-IID data and reward heterogeneity. This work advances privacy-preserving, on-device personalization of reasoning capabilities in LLMs by enabling stable, multi-objective optimization with limited communication rounds.

Abstract

Group Relative Policy Optimization (GRPO) has recently emerged as an effective approach for improving the reasoning capabilities of large language models through online multi-objective reinforcement learning. While personalization on private data is increasingly vital, traditional Reinforcement Learning (RL) alignment is often memory-prohibitive for on-device federated learning due to the overhead of maintaining a separate critic network. GRPO's critic-free architecture enables feasible on-device training, yet transitioning to a federated setting introduces systemic challenges: heterogeneous reward definitions, imbalanced multi-objective optimization, and high training costs. We propose FedMOA, a federated GRPO framework for multi-objective alignment under heterogeneous rewards. FedMOA stabilizes local training through an online adaptive weighting mechanism via hypergradient descent, which prioritizes primary reasoning as auxiliary objectives saturate. On the server side, it utilizes a task- and accuracy-aware aggregation strategy to prioritize high-quality updates. Experiments on mathematical reasoning and code generation benchmarks demonstrate that FedMOA consistently outperforms federated averaging, achieving accuracy gains of up to 2.2% while improving global performance, personalization, and multi-objective balance.

Figures (5)

• Figure 1: Training curves of multi-objective rewards in federated GRPO. Format-related rewards converge rapidly, while accuracy-related rewards remain low and exhibit substantial fluctuations.
• Figure 2: Overview of FedMOA. Each client performs multi-objective GRPO with a local hypergradient-based procedure to adapt objective weights, prioritizing accuracy while balancing auxiliary objectives. The adapted weights are sent to the server together with model updates and used for task-aware grouping and accuracy-aware aggregation.
• Figure 3: Trade-offs among reward components on GSM8K(upper) and MATH(lower). Pareto frontiers are shown for format reward vs. accuracy (left), tag count reward vs. accuracy (middle), and tag count reward vs. format reward (right). Gray dots denote client-level checkpoints, while colored lines indicate server-level trajectories under different training settings. Insets zoom into the final-stage region, highlighting subtle yet consistent advantages of FedMOA over FedGRPO in both homogeneous and heterogeneous settings.
• Figure 4: Reward evolution during MATH training. This figure illustrates the evolution of reward signals for accuracy, format, and tag count, together with the average response length (right y-axis). The top row reports results obtained with FedGRPO, while the bottom row corresponds to FedMOA.
• Figure 5: Dynamic reward weight adaptation on GSM8K. We illustrate the evolution of reward weights for accuracy, format, and tag count throughout training. The top panel shows the Homo+Homo setting, whereas the bottom panel shows the Heter+Heter setting.