Trapped by Their Own Rollouts: Understanding Aggregation--Rollout Feedback in Federated On-Policy Distillation
Organizations: School of Software Engineering, Xi’an Jiaotong University
Abstract
On-policy distillation (OPD) is a promising approach to language-model adaptation, aligning teacher supervision with the student's own generated trajectories. When adaptation prompts are distributed across clients, can this process benefit from federated collaboration? We study federated OPD and find that substantial collaboration gains can be obscured by learning-rate sensitivity: FedAvg can perform no better than independent local training at a small learning rate, yet recover a clear advantage at a larger rate. We explain this phenomenon through the student's dual role as learner and generator of future training data. An aggregation-induced optimization lag can delay access to useful teacher supervision, which in turn slows subsequent learning. Our theory establishes this aggregation--rollout feedback in a solvable model with a common optimum and stable updates, and identifies two coupled roles of learning rate: learning from current supervision and reaching future supervision. Guided by this analysis, we propose FedTOPS (Federated Teacher-guided On-Policy Scaling), which reuses teacher feedback on current trajectories to adapt the FedAvg update magnitude under clientwise predictive-change constraints. Across six mathematical reasoning benchmarks, FedTOPS improves macro Avg@8 over FedAvg by 4.56--14.57 percentage points across the evaluated student models and local learning rates.
Figures & tables
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
| Symbol | Meaning |
|---|---|
| Number of clients, aggregation weight, and fixed prompt distribution | |
| Fixed teacher, receiving learner, and trajectory generator | |
| Maximum horizon and uniformly position-sampled prefix distribution | |
| Fixed-source loss and its learner gradient | |
| Learner response and generator response of the update vector field | |
| Local gradient steps per round, step size, and communication rounds |
| Setting | Value |
|---|---|
| Participation / rounds | All 3 clients / 50 rounds |
| Rollouts per round | 32 per client, 1 response per prompt |
| Mini-batch / micro-batch | 8 / 4; one optimization epoch |
| Local learning rate | 0.6B: or ; 1.7B: |
| Optimizer | AdamW, , weight decay 0.01, gradient clipping 1 |
| Scheduler | Cosine; no warmup |
Federated Knowledge Distillation
Federated knowledge distillation transfers predictive behavior among data-isolated clients to train useful models without sharing raw data.
See all 36 papers