# Matching Rates and Optimal Allocation for Federated Probe-Logit Distillation under Heterogeneous Bandwidth Budgets

arXiv: https://arxiv.org/abs/2605.29642 · HTML: https://mlsanity.com/paper/2605.29642/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Federated, robust and private distillation, one of 12 directions in Knowledge distillation.
- Builds on 6 of its 7 references that are in the map.
- First posted on arXiv 28 May 2026.

## Placed in

- Knowledge distillation / [Federated, robust and private distillation](https://mlsanity.com/distill/federated-robust-and-private-distillation/)

## Builds on (most cited first, 6 of 6)

- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Inference under Information Constraints I: Lower Bounds from Chi-Square Contraction](https://mlsanity.com/paper/1812.11476/) (2018)
- [On Communication Cost of Distributed Statistical Estimation and Dimensionality](https://mlsanity.com/paper/1405.1665/) (2014)
- [A Distributed One-Step Estimator](https://mlsanity.com/paper/1511.01443/) (2015)
- [Feature Distillation is the Better Choice for Model-Heterogeneous Federated Learning](https://mlsanity.com/paper/2507.10348/) (2025)
- [Federated Language Models Under Bandwidth Budgets: Distillation Rates and Conformal Coverage](https://mlsanity.com/paper/2605.09986/) (2026)

## Abstract (arXiv)

In federated language modeling, K nodes each hold n samples but cannot pool data or exchange full-precision gradients or weights. We study the minimax rate at which a conditional distribution over V tokens can be estimated when each node may upload at most B bits per query in a public probe set. In federated probe-logit distillation (FPLD), each node transmits a scalar-quantized logit vector on the probe set, and an aggregator distills a global parametric student. Prior work (Dubey and Huo, 2026) establishes a high-probability KL rate O(d/(Kn) + ρ√V log V / m + K-1 · 2-2B/V) plus optimization slack, with the bandwidth term in its trace-sharpened form. Whether this bandwidth-term rate is tight, and how the upper bound generalizes to heterogeneous per-node bandwidths, are left open. We close both gaps. First, the dithered FPLD construction has a matching single-round lower bound Omega(K-1 · 2-2B/V) under non-degeneracy, pinning the bandwidth-axis rate at Theta(K-1 · 2-2B/V). T-round sequential refinement with nested/scaled residual quantizers achieves O(K-1 · 2-2TB/V); vanilla FPLD's T-independent bandwidth term is suboptimal for every T > 1. Second, we establish a heterogeneous-bandwidth upper bound for per-node budgets Bi, paired with a closed-form optimal allocation Bi^* = Btot/K + (V/2) log2(wi / barwg), a log-tilted water-filling rule that is the per-node analogue of reverse water-filling for distortion-rate optimization. A plug-in adaptive variant estimates the weights from a short warm-up phase and attains 1 + O(√log(K/δ)/(m T0)) relative suboptimality. Synthetic n-gram simulations confirm that empirical KL is bracketed by the upper and lower bounds and that the optimal allocation strictly dominates uniform and inverse-weighted baselines under heterogeneous clipping.