解决异构带宽下联邦模型蒸馏的最优传输速率与分配问题
Matching Rates and Optimal Allocation for Federated Probe-Logit Distillation under Heterogeneous Bandwidth Budgets
- 提出带抖动的FPLD构造，证明带宽项下界为Θ(K⁻¹·2⁻²ᴮ⁄ⱽ)
- 多轮迭代量化可实现2⁻²ᵀᴮ⁄ⱽ的更优速率，单轮次非最优
- 给出异构带宽下的最优比特分配公式，适配资源不均场景
在联邦语言建模中，K个节点各持有n个样本，但无法共享数据或全精度梯度/权重。研究在每节点每查询最多上传B比特条件下，对V个词元的条件分布估计的极小极大率。在联邦探针-逻辑值蒸馏（FPLD）中，各节点向聚合器发送探针集上的标量量化逻辑值向量，聚合器训练全局学生模型。先前工作（Dubey and Huo, 2026）给出高概率KL率O(d/(Kn) + ρ√(V log V / m) + K⁻¹·2⁻²ᴮ⁄ⱽ)，含优化松弛项，且带宽项以迹锐化形式呈现。该带宽项是否紧致，以及上界能否推广至异构节点带宽，仍待解决。本文填补两个空白：第一，带抖动的FPLD构造在非退化条件下有匹配的单轮下界Ω(K⁻¹·2⁻²ᴮ⁄ⱽ)，确定带宽轴率为Θ(K⁻¹·2⁻²ᴮ⁄ⱽ)；多轮序列精炼结合嵌套/缩放残差量化器可达O(K⁻¹·2⁻²ᵀᴮ⁄ⱽ)，表明原版FPLD的带宽项在T > 1时非最优。第二，推导出异构带宽上界，对应闭式最优分配方案B_i* = B_tot/K + (V/2) log₂(w_i / w̄_g)，即对数倾斜的水填规则，是失真-速率优化中反向水填的节点级类比；插件自适应变体通过短预热阶段估计权重，达到1 + O(√(log(K/δ)/(m T₀)))的相对次优性。合成n-gram模拟验证了经验KL被上下界夹住，且最优分配在异构裁剪下严格优于均匀和逆权基线。
原文摘要 · Abstract (English)
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) + ρ\sqrt{V \log V / m} + K^{-1} \cdot 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 $Ω(K^{-1} \cdot 2^{-2B/V})$ under non-degeneracy, pinning the bandwidth-axis rate at $Θ(K^{-1} \cdot 2^{-2B/V})$. $T$-round sequential refinement with nested/scaled residual quantizers achieves $O(K^{-1} \cdot 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 $B_i$, paired with a closed-form optimal allocation $B_i^* = B_{\mathrm{tot}}/K + (V/2) \log_2(w_i / \bar{w}_g)$, 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(\sqrt{\log(K/δ)/(m T_0)})$ 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。