PubSwap: A Practical Recipe for Federated RLVR with Public-Data Coordination

Executive Summary

TL;DR. This paper targets a very timely problem: how to do reasoning post-training with RLVR when the useful data is split across multiple organizations and cannot be shared. The authors propose PubSwap, which combines LoRA-based federated training with periodic public-data coordination, so clients can exchange response-level learning signals without exposing private prompts or answers. The key design choice is a Balanced response swap rule that only injects globally correct responses when a client is locally failing, reducing off-policy mismatch while still anchoring clients to a shared objective. Empirically, the method beats FedAvg-GRPO and FedProx across math and medical reasoning, especially when the number of local steps is large.

Positioning. This is not a foundation-model-scale breakthrough, nor a new RL algorithm in the pure optimization sense. Its contribution is more important from a systems-and-training perspective: it identifies a clean, low-overhead mechanism that makes federated RLVR actually workable under realistic privacy and communication constraints. In that sense, it is best read as a pragmatic training recipe with strong empirical value, rather than a fully closed theoretical solution.

---

1. Problem & Motivation

RLVR has become a core ingredient in reasoning model post-training because it optimizes directly against outcome-level rewards on tasks where correctness is verifiable. That is a strong fit for domains like:

• mathematics,
• coding,
• scientific reasoning,
• medical QA,
• finance.

But those are also exactly the domains where the training data is often institutionally siloed.

Why centralized RLVR breaks in these settings

In the centralized setting, the trainer can:

• collect all prompts centrally,
• sample many trajectories,
• update one shared model,
• keep the optimization loop tightly synchronized.

In the federated setting, none of this is straightforward.

The main failure modes are:

1. Communication cost

   • Full fine-tuning of LLMs requires transmitting huge parameter updates.
   • RL training often involves repeated synchronization, which makes this even more expensive.

2. Client drift under heterogeneity

   • If each client performs many local GRPO steps on its own private distribution, the local policies diverge.
   • In autoregressive models, even small parameter differences can produce different token distributions, which then create different rollouts, rewards, and gradient targets.

3. Parameter-space proximity is not policy-space proximity

   • This is an important insight in the paper.
   • Methods like FedProx constrain local parameters to stay close to the global model, but for sequence generation that does not guarantee similar sampled trajectories.

So the paper’s motivating question is:

> Can we keep the communication benefits of federated PEFT while introducing just enough shared signal to reduce drift, without sharing private data?

Their answer is: yes, by coordinating on a small public dataset at the response level rather than the raw private-data level.

---

2. The Core Idea in One Sentence

PubSwap alternates local private GRPO with public-data steps where clients generate responses on the same shared prompts, and then selectively reuse globally correct responses to anchor training across organizations.

That sentence contains almost the whole paper.

The elegance of the method is that it does not require:

• sharing private prompts,
• sharing private rewards,
• synchronizing full dense models,
• distilling logits into a server model.

Instead, it uses a small public dataset as a coordination substrate.

---

3. Methodology - The Core

3.1 GRPO recap: what is being optimized?

The paper builds on GRPO. For each prompt, the model samples a group of responses and computes a normalized relative advantage:

[ A_k(x, y_k) = \frac{r(x,y_k)-\bar r(x)}{\sigma(x)}. ]

Intuitively:

• responses are judged relative to one another,
• the model gets pushed toward responses that outperform the group average,
• no separate critic is needed.

The clipped GRPO objective then plays the same stabilizing role as PPO-style clipping:

• reuse old rollouts for a few updates,
• limit destructive policy jumps,
• optimize token-level likelihood ratios weighted by relative advantage.

In practice, a GRPO step has two stages:

1. sample responses using the old policy,
2. update parameters via gradient ascent on the GRPO objective.

That standard pipeline is retained. PubSwap modifies how the response set is constructed during certain steps.

3.2 LoRA for federated efficiency

To make federated RLVR feasible, the authors use LoRA instead of full fine-tuning. Each weight matrix is parameterized as:

[ W_p = W_0 + B_p A_p ]

where:

• (W_0) is frozen,
• (A_p, B_p) are low-rank trainable factors,
• only LoRA parameters are communicated.

The aggregation rule is simple averaging of LoRA factors across clients:

[ B_{p+1} = \frac{1}{N}\sum_{n=1}^{N} B_{p+1}^{(n)}, \qquad A_{p+1} = \frac{1}{N}\sum_{n=1}^{N} A_{p+1}^{(n)}. ]

The communication reduction is substantial:

• dense update cost per layer: (2md)
• LoRA update cost per layer: (2r(m+d))

with (r \ll \min(m,d)).

This is the first half of the solution: cheap synchronization.

3.3 PubSwap: public-data off-policy coordination

The second half is what makes the paper novel.

During a local training round of length ( au ):

• most steps are standard private-data GRPO,
• every ( au_{swap} )-th step is a public-data step.

On a public step:

1. the server samples a batch of public prompts,
2. all clients generate (K) responses per prompt,
3. the server pools these responses,
4. it sends back a selected subset of responses to each client,
5. each client performs its GRPO update using that response set.

The important point is that the exchanged object is not gradients and not private data. It is only public-prompt generations.

3.4 Random vs Balanced response aggregation

The paper studies two ways to build the public-step response set.

Random

The server randomly samples (K) responses from the global pool of (NK) candidate responses per prompt.

This is simple but highly off-policy:

• the local client may train on trajectories produced mostly by other clients,
• as client drift increases, this mismatch can become harmful.

Balanced

This is the main contribution.

For a client (n), let (c_n) be the number of correct responses among its own (K) generated responses for a prompt.

• If (c_n \ge K/2): keep the client’s own responses.
• If (c_n < K/2): replace up to (K/2 - c_n) incorrect responses with correct responses drawn from other clients.

This design is subtle and quite smart.

Why Balanced should work

It has three interacting effects:

1. It only intervenes when needed

   • If the local client is already doing reasonably well, no swap happens.
   • So easy prompts remain mostly on-policy.

2. It preserves local policy identity

   • Only incorrect responses are replaced.
   • Correct local responses are retained.
   • That keeps the learning signal closer to what the client itself would have seen.

3. It maximizes useful reward variance

   • If all responses are correct or all are wrong, GRPO gives near-zero learning signal because relative advantages collapse.
   • A half-correct / half-incorrect set creates informative variance in rewards.
   • Balanced explicitly pushes toward that regime, but only up to the threshold.

This is the conceptual center of the paper: PubSwap is not simply exchanging good answers; it is shaping the response set so GRPO has a better training signal while limiting off-policy damage.

---

4. The Intuition Behind the Theory

The appendix gives a simplified drift analysis. It is not a full convergence theorem, but it is useful.

For private local steps, drift between two clients grows with:

• the smoothness constant over private objectives,
• gradient heterogeneity between client datasets,
• sampling noise.

In shorthand, private training drift depends on the fact that each client is following a different local gradient field.

For Balanced public steps, the drift decomposes differently:

• all clients are anchored to the same public minibatch,
• residual mismatch appears through an off-policy distortion term,
• that distortion is controlled by:
  • how many responses must be replaced,
  • how sensitive the GRPO gradient is to one replacement.

This gives a useful mental model:

• Private steps: drift is driven directly by heterogeneous data.
• Balanced public steps: drift is driven only by the mismatch introduced by swapped responses.

If the public data is representative enough and clients are not too far apart, then the public step should create less divergence than another private step. That is precisely why the method helps most when local-step counts are large.

---

5. Experiments & Results

The evaluation spans:

• Math reasoning: MATH and DeepMath
• Medical reasoning: MedQA + MedMCQA mixture
• Models: Qwen2.5-Math-1.5B, Qwen3-1.7B, Qwen3-4B-Instruct, Llama-3.2-3B-Instruct
• Federated setup: 4 clients with heterogeneous splits
• PEFT: LoRA rank 32, alpha 64

The comparison baselines are:

• FedAvg-GRPO
• FedProx-GRPO
• FedAvg-PubSwap

5.1 Main math results

The strongest story appears on DeepMath, especially at high local-step counts.

Key numbers:

• Qwen3-1.7B, DeepMath, τ=120

  • FedAvg-GRPO: 50.7
  • FedProx-GRPO: 47.7
  • PubSwap: 55.8

  This is:

  • +5.1 points over FedAvg-GRPO
  • +8.1 points over FedProx

• Qwen3-1.7B, DeepMath, τ=90

  • FedAvg-GRPO: 50.4
  • PubSwap: 53.3
  • gain: +2.9

• Qwen2.5-Math-1.5B, DeepMath, τ=120

  • FedAvg-GRPO: 49.3
  • FedProx-GRPO: 49.4
  • PubSwap: 53.1
  • gain: +3.8 vs FedAvg

On MATH, the gains are smaller and sometimes mixed. That is actually believable:

• MATH is easier than DeepMath,
• the need for cross-client coordination is weaker when the task is easier,
• public-data correction matters more in harder reasoning regimes.

5.2 Qwen3-4B-Instruct results

The 4B model on DeepMath also supports the trend:

• τ=10: 66.9 → 70.0 (+3.1)
• τ=90: 67.0 → 70.2 (+3.2)
• τ=120: 67.4 → 70.5 (+3.1)

That consistency across scale is a good sign that the method is not an artifact of one specific model family.

5.3 Medical reasoning

The medical results are especially interesting because they represent a domain where federated learning is genuinely realistic.

For Llama-3.2-3B-Instruct:

• τ=40: 52.0 → 53.0 (+1.0)
• τ=90: 51.6 → 56.7 (+5.1)
• τ=120: 51.3 → 54.9 (+3.6)

Again, the method helps most when local computation is heavy.

5.4 FedProx underperforms: why that matters

One of the paper’s most useful empirical observations is that FedProx is not a good fix here.

In classical federated supervised learning, a proximal penalty often helps because it discourages local models from moving too far in parameter space. But in RLVR for autoregressive LLMs:

• the update target depends on sampled trajectories,
• sampled trajectories are sensitive to token-level shifts,
• small parameter changes can induce large rollout changes.

So a Euclidean penalty on parameters is a fairly weak proxy for the thing we actually care about: policy similarity.

This is an important negative result. It suggests that federated RLVR needs trajectory-aware or response-aware coordination, not just optimizer regularization.

---

6. Ablation: Why the Swap Period Matters

The paper varies the PubSwap period ( au_{swap} ), which controls how often public-data coordination occurs.

The tradeoff is intuitive:

• More frequent public steps

  • better synchronization,
  • stronger anchor against drift,
  • but more off-policy contamination.

• Less frequent public steps

  • less anchor,
  • but also less mismatch.

The data shows there is no universal “more public is always better” rule. That is exactly what you would expect in off-policy RL. Once models drift sufficiently far apart, importing many external responses can destabilize training.

So the paper’s recommendation is moderate rather than maximal coordination:

• swap periods like 2, 4, or 8 are usually good,
• the best choice depends on the local-step regime and model.

This is one of the paper’s strengths: it does not oversell the public-data intervention as free lunch.

---

7. What Is Actually Novel Here?

On the surface, the components are familiar:

• federated learning,
• LoRA,
• GRPO,
• public data,
• off-policy response reuse.

The novelty is in how these parts are assembled.

The paper’s real contribution is a new training pattern:

1. Use PEFT to make federated RLVR affordable.
2. Use shared public prompts as a coordination anchor.
3. Exchange response-level supervision instead of private data or full gradients.
4. Do selective correction rather than indiscriminate mixing.

That fourth point is the most important. Balanced is not just a heuristic; it encodes a principled bias:

• preserve locality when the local model is competent,
• borrow signal only where the local rollout set is weak.

This is why the method feels well matched to RLVR, where the composition of the sampled response set directly controls learning signal quality.

---

8. Critical Analysis

Strengths

• Practical relevance
  • The problem setting is real: multi-organization post-training with privacy constraints.
• Low implementation complexity
  • PubSwap adds a lightweight coordination protocol without redesigning GRPO itself.
• Strong empirical focus on the right regime
  • Gains are largest when τ is large, which is exactly when communication efficiency matters.
• Good negative result on FedProx
  • Useful for the community because it clarifies what does not work well.

Limitations

• Needs a public dataset

  • The whole method assumes a shared public reference set exists and is useful.
  • In some domains, finding even a small representative public dataset is nontrivial.

• Off-policy bias is only partially controlled

  • Balanced is smarter than Random, but it is still off-policy.
  • The paper acknowledges that stronger off-policy correction methods could help.

• Theory is suggestive, not definitive

  • The appendix provides drift intuition rather than a full theorem for realistic GRPO with AdamW, clipping, and multiple updates.

• Scale remains modest

  • The experiments are on 1.5B–4B class models with 4 clients.
  • It is still unclear how the method behaves at larger federation scales or much larger model sizes.

• Public prompts may favor global over local optima

  • If the public dataset is weakly aligned with some clients’ domains, the anchor could become biased rather than helpful.

A subtle conceptual limitation

The method improves coordination by sharing correct responses. But correctness in RLVR is only part of what matters for reasoning models:

• response style,
• chain-of-thought diversity,
• calibration,
• robustness under distribution shift.

PubSwap does not directly address whether the imported correct responses are the most educational responses for a given client policy. Future work could make this smarter by selecting responses based on informativeness, not just correctness.

---

9. Future Work

The paper already hints at good directions, and several seem especially promising.

9.1 Adaptive public-step scheduling

Instead of a fixed swap period:

• coordinate more when client drift is high,
• coordinate less when policies are already aligned.

9.2 Better donor-response selection

The server could choose not just any correct response, but the one that is:

• closest to the client’s policy,
• most informative for the current error mode,
• most stable for off-policy reuse.

9.3 Off-policy correction on top of PubSwap

Importance weighting, clipped corrections, or policy-mismatch-aware filtering could reduce bias further.

9.4 Beyond text-only reasoning

The same idea should transfer naturally to:

• code generation,
• multimodal reasoning,
• medical vision-language post-training,
• tool-using agents.

9.5 Personalized federated RLVR

A particularly interesting extension is to combine public coordination with personalized heads or client-specific adapters, allowing shared reasoning gains without collapsing local specialization.

---

10. Final Takeaway

PubSwap is a strong paper because it solves the right problem with the right amount of machinery. It does not claim to reinvent RLVR. Instead, it identifies a practical gap in current training pipelines: federated reasoning post-training breaks when clients drift too far, but dense synchronization is too expensive.

The solution is conceptually clean:

• LoRA keeps communication cheap.
• Public prompts create a shared coordination surface.
• Balanced response swapping gives clients external correction only when they need it.
• The result is better federated RLVR, especially in the high-local-step regime that matters most in deployment.

If centralized RLVR was the first wave of reasoning post-training, this paper sketches part of the recipe for the next one: privacy-aware, decentralized, communication-efficient RLVR with response-level coordination.