Matching Rates and Optimal Allocation for Federated Probe-Logit Distillation under Heterogeneous Bandwidth Budgets
Abstract
In federated language modeling, nodes each hold samples but cannot pool data or exchange full-precision gradients or weights. We study the minimax rate at which a conditional distribution over tokens can be estimated when each node may upload at most 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 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 under non-degeneracy, pinning the bandwidth-axis rate at . -round sequential refinement with nested/scaled residual quantizers achieves ; vanilla FPLD's -independent bandwidth term is suboptimal for every . Second, we establish a heterogeneous-bandwidth upper bound for per-node budgets , paired with a closed-form optimal allocation , 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 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.
Cite
@article{arxiv.2605.29642,
title = {Matching Rates and Optimal Allocation for Federated Probe-Logit Distillation under Heterogeneous Bandwidth Budgets},
author = {Prasanjit Dubey and Xiaoming Huo},
journal= {arXiv preprint arXiv:2605.29642},
year = {2026}
}