TallyTrain: 通信高效的联邦蒸馏

# TallyTrain：通信高效的联邦蒸馏
来源：https://arxiv.org/html/2607.00173
\jmlryear
2026\jmlrworkshop
\NameRadhakrishna Achanta\Emailrachanta@cisco\.com \NameWill Reed\Emailwilreed@cisco\.com \addrCisco Systems Inc

###### 摘要

联邦学习在两个正交轴上受限于带宽：模型大小，这限制了参数平均方法可以负担的合并频率；以及类别数量，这使得在大词汇量场景下基于每个探针的软标签蒸馏变得不可行。随着现代系统规模的扩大，这两个瓶颈都在收紧。我们通过仅传输每个对等体的\arg\max类别索引（其中C是输出类别的数量），将每个探针的类别数量轴压缩到⌈log₂C⌉位。由此产生的协议TallyTrain不仅仅是压缩的：在非IID训练下，它可能*优于*软标签蒸馏，因为训练不足的对等体往往会自信地犯错，而多数投票可以过滤这种噪声，但软标签平均会放大它。在标准基准测试中，TallyTrain以最多三个数量级的通信成本，匹配或击败了软标签蒸馏。我们还放宽了模型大小轴：将廉价的硬标签共识与稀疏参数合并相结合，得到一种带宽桥接变体，该变体在Pareto上支配了标准FedAvg、FedProx和FedDF基线的所有测试工作点。

###### 关键词：

联邦学习；知识蒸馏；通信效率；硬标签共识；去中心化训练；非IID

## 1 引言

联邦学习（FL）在不对原始数据进行集中化的情况下，从多个对等体持有的数据分片中训练模型。两种主要范式在它们*通信什么*上有所不同。

**参数空间方法**交换模型参数（FedAvgmcmahan2017communication）或每轮权重增量，这些增量充当外部循环伪梯度（DiLoCodouillard2023diloco）的角色。两者都有每轮带宽Θ(‖W‖)，其中‖W‖是模型参数的数量，这对于在边缘或移动平台上处理十亿参数模型是不切实际的。

**函数空间方法**（FedMDfedmd2019）在共享的公共探针集上交换每个示例的预测。每个对等体每轮的带宽为Θ(C·‖D_pub‖)，其中C是输出类别的数量——独立于模型大小，但随C线性增长，这对于大词汇量任务（BPE分词的语言模型通常使用C ∈ [2,048, 50,000]）变得不可行。

这些范式共享一个结构假设：带宽沿着*频率*轴减少，通过降低通信频率，同时每个消息的大小被视为固定（完整权重，或完整的C维softmax）。本文则采用正交的*大小*轴。我们保持通信频繁，但使每条消息非常小：每个探针仅传输\arg\max类别索引，C ≤ 256时为一个字节，C ≤ 65,536时为两个字节。Argmax投票充当噪声滤波器——当对等体达成一致时，共识是可靠的；当它们不一致时，投票平均会分摊个体错误而不是放大它（这正是软标签期望在多个对等体同时训练不足且自信时所做的事情）。我们将由此产生的协议称为**TallyTrain**——每轮，对等体*tally*它们在公共探针集上的argmax投票到一个共识直方图中，然后*train*针对它——并表明，结合稀疏参数合并，它在两个模态和三个类别数量的联邦学习带宽-准确率Pareto前沿上占据主导地位。我们的贡献如下：

(1) **一种硬标签通信原语。** 我们介绍了TallyTrain，并论证N个对等体在公共探针集上的argmax预测的投票直方图是一种有效——且出奇强大——的蒸馏共识分布（§4.1 (https://arxiv.org/html/2607.00173#S4.SS1)）。

(2) **硬标签在各种模态下匹配或击败软标签。** \arg\max信道携带了软标签几乎所有的信号，其每个探针的带宽比随C线性增长（§4.1 (https://arxiv.org/html/2607.00173#S4.SS1), §4.4 (https://arxiv.org/html/2607.00173#S4.SS4), §4.6 (https://arxiv.org/html/2607.00173#S4.SS6)）。

(3) **两种操作模式。** 同一原语支持*纯函数空间*模式（当探针带有标签时α > 0，当使用KL衰减时α = 0；§4.1 (https://arxiv.org/html/2607.00173#S4.SS1), §G (https://arxiv.org/html/2607.00173#A7)）和*参数空间平均的函数空间稳定器*——带宽桥接变体TallyTrain+faMM，它交错进行稀疏的FedAvg合并和廉价的硬标签信道（§4.2 (https://arxiv.org/html/2607.00173#S4.SS2), §4.6 (https://arxiv.org/html/2607.00173#S4.SS6)）。相对于mcmahan2017communicationCIFAR默认设置（E ≈ 5, fa ≈ 400），桥接变体将参数信道带宽减半并*提高*了准确率；相对于fa = 1，它将其降低了两个数量级。

(4) **硬标签蒸馏的收缩理论：** 一个函数空间收缩引理，一个关于投票top-1的Condorcet界，以及一个关于蒸馏梯度的方差缩减界（§3 (https://arxiv.org/html/2607.00173#S3)）。

#### 范围：

带宽是约束资源。我们不对训练计算进行摊销：每个对等体持有一个完整模型并运行独立的本地SGD，因此总FLOP随N线性增长。我们在N ≤ 10的全网格拓扑下进行评估，这在轨迹上等同于集中式中继（§3 (https://arxiv.org/html/2607.00173#S3)）；扩展到数千个对等体的稀疏八卦扩展在此讨论但未进行实验评估。

## 2 相关工作

分布式和联邦学习方法kairouz2019advances不仅在拓扑（集中式与点对点）上有所不同，还在**通信对象**上有所不同。经典的去中心化联邦学习（DFL）在相邻客户端之间交换参数、梯度或压缩后的模型增量。基于蒸馏的方法则交换**预测空间对象**，如logits、软化目标、原型、质心或函数值，从而将协作与精确的参数对齐解耦。这种区别在异构和去中心化环境中尤其重要，因为在这些环境中，客户端在架构、容量、数据分布、可用性和通信预算上可能有差异。下面，我们通过这个视角回顾最相关的文献，从蒸馏谱系开始，然后将其与最近的优化中心型DFL方法进行对比。

### 2.1 知识蒸馏作为合作原语

知识蒸馏（KD）最初是一种教师-学生范式，其中学生从更强的教师的软化预测中学习hinton2015distilling。对分布式学习而言，相关的含义是协作可以发生在**函数空间**：共享对象是预测分布而不是权重向量，这可以容纳异构模型，并且如果只交换紧凑的预测摘要，原则上可以减少通信。*联邦蒸馏*（FD）jeong2018communication是这一思想的早期形式，在联邦范围内交换平均后的每类logits。*FedMD*fedmd2019形式化了公共探针集版本：客户端持有私有数据并且可能有不同的架构，但共享一个小型公共参考集，反复在其上评估模型，并使用聚合的logits作为蒸馏目标。*FedDF*lin2020ensemble将其扩展到集成蒸馏：服务器端学生在客户端在未标记的公共集上的平均软预测上训练，但每个探针的有效载荷仍然是一个长度为C的浮点向量。*Cronus*chang2019cronus在精神上与我们的原语更接近：它将黑盒知识转移与公共集上的硬预测配对，并指出了由此产生的带宽和隐私优势。*FedKD*fedkd2021保留了KD观点，但将通信对象压缩为紧凑的student模型产物。这些方法是服务器协调的；它们共同确立了蒸馏可以桥接异构性并减少通信，但没有一个给出完全去中心化的协议，其每个探针的有效载荷随log₂C而不是C缩放。

### 2.2 异构学习的去中心化蒸馏

多线工作将FedMD的直觉移植到点对点设置中。*DFML*dfml2024是无服务器的，并且避免了公共辅助数据：相邻客户端直接相互蒸馏，采用循环的监督/蒸馏调度，但代价是对邻域质量和本地数据不平衡敏感。*DeSA*desa2024通过**合成**锚点重建共享的对齐基底，为异构模型之间的蒸馏提供了一种没有真实公共数据的通用介质，同时付出了合成成本并继承了锚点质量风险。*DFedCAD*dfedcad2025用类质心摘要替换了完整预测，当带宽或掉队者同步占主导时具有吸引力，但传输的不确定性信息较为粗糙。因此，现代去中心化蒸馏的设计空间按共享什么（真实公共锚点、对等体预测、合成锚点、质心或函数空间一致性项）分解，每种选择都带有关于辅助数据、异构性、拓扑和通信的不同假设。

### 2.3 优化中心型的去中心化联邦学习

与蒸馏文献并行，一条强大的工作线通过更好的优化、拓扑设计、个性化和隐私改进了DFL，同时主要保持在参数空间内。DSpodFLdspodfl2025, DPFLdpfl2025, NTK-DFLntkdfl2025, FedSPDfedspd2025, GFlatgflat2026, PaMEpame2026, 和 Spod-GTspodgt2026分别贡献了偶发通信收敛保证、学习式协作图、基于NTK的稳定化、软聚类个性化、平坦极小泛化、部分消息带宽减少和在有向图上的梯度跟踪。当客户端可以共享一个共同的模型家族且主要挑战是系统效率时，它们优于蒸馏，但它们的参数空间消息预设了传输的坐标在客户端之间具有可比意义，这限制了跨架构的迁移。互补的线索将DFL的范围扩展到准确率和收敛之外：用于隐私-效用权衡的*f-DP for DFL*fdpdfl2025，用于策略操纵的*Competitive Advantage Attacks*compadvdfl2025，以及用于卫星网络拓扑的*Learning in Orbit*orbitdfl2026。

#### 定位TallyTrain。

去中心化硬标签联邦蒸馏（TallyTrain）针对上面留下的工程瓶颈。与FedAvg或DiLoCo不同，它完全在函数空间中操作，绕过了非IID数据上的权重合并障碍。与FedMD/FedDF/Cronus不同，它每个探针传输一个单一的1字节\arg\max索引，而不是长度为C的logits。与DFML和其他点对点相互蒸馏协议相比，对等体间的有效载荷*仅*是1字节的硬预测；没有软信号、锚点或参数通过线路传输。这使得该协议在每当‖W‖或C足够大以至于浮点向量或参数交换主导通信预算时特别有吸引力（§4.6 (https://arxiv.org/html/2607.00173#S4.SS6), §4.4 (https://arxiv.org/html/2607.00173#S4.SS4)）。

## 3 方法

FedAvg: 权重 Θ(‖W‖) 字节  
DiLoCo: 伪梯度 Θ(‖W‖) 字节  
FedMD: 软标签 4C 字节/探针  
TallyTrain: argmax 投票 1 字节/探针

图1：联邦学习范式按每轮有效载荷排序。箭头颜色编码交换的对象：权重（黑色，Θ(‖W‖)），伪梯度（红色，Θ(‖W‖)），软标签（橙色，4C B/探针），argmax 投票（绿色，1 B/探针）。所有四个面板都绘制在相同的星形拓扑上，以使视觉比较仅关于有效载荷；拓扑与有效载荷选择正交，并且四种协议中的任何一种都可以在全网格、环或八卦上运行而无需算法更改（§3 (https://arxiv.org/html/2607.00173#S3)）。

TallyTrain 围绕一个单一的设计选择构建——在共享的公共探针集上进行 argmax 投票——并有两种操作变体与之组合：一种将锚定到真实标签的带标签公共集混合体，以及一种交替进行周期性 FedAvg 参数合并的带宽桥接变体。图1 (https://arxiv.org/html/2607.00173#S3.F1) 通过每轮有效载荷将原语与现有的联邦学习范式进行了对比。在全网格和集中式中继下，该原语在轨迹上是等价的。

### 3.1 函数空间对齐

设 D_n 表示对等体 n 的私有非IID分片，D_pub 是所有对等体共享的公共探针集，每个对等体持有 f_n(·; W_n): X → Δ^{C-1}。权重平均 W_g = (1/N) ∑_n W_n 继承了非线性模型在非IID数据上常见的损失曲面不匹配。TallyTrain 转而将*functions*在 D_pub 上对齐：两个在 D_pub 中所有 x 上达成一致的模型，在由 D_pub 主导的任何测试分布上实现了相同的预测器。

### 3.2 通信原语：argmax 投票

对于每个 x ∈ D_pub，对等体 n 仅广播 top-1 预测

y_n(x) = argmax f_n(x; W_n) ∈ {1, ..., C}. \hspace{1em}(1)

一个类别索引需要 ⌈log₂C⌉ 位，按字节对齐传输（C ≤ 256 时 b_C = 1 B，C ≤ 65,536 时 2 B）。相对于 4C 字节的 32 位软标签，实现比率为 ρ_C = 4C / b_C（ρ_10 = 40, ρ_100 = 400, ρ_2048 = 4096，图5 (https://arxiv.org/html/2607.00173#S4.F5)）；位打包的下界为 ρ_C^⋆ = 32C / ⌈log₂C⌉。节省量随 C 线性增长。

跨对等体，硬标签预测形成经验投票直方图

\=H(x) = (1/N) ∑_{n=1}^N e_{y_n(x)} ∈ Δ^{C-1}, \hspace{1em}(2)

其中 e_c 是第 c 个基向量。\=H(x) 是一个有效的概率分布，并作为蒸馏的共识目标。

#### 去中心化 ≡ 集中式在轨迹上。

在全网格下，每个对等体本地平均 {y_m(x)} 以获得与集中式中继会计算的相同 \=H(x)，因此每个对等体的蒸馏梯度相同，两种拓扑共享其优化轨迹；它们仅在连线模式（全对全 vs. 星形）上有所不同。因此，§4 (https://arxiv.org/html/2607.00173#S4) 中的所有数字对两种解释同时有效。每个对等体每个探针的带宽，对于去中心化硬标签为 2(N-1) 字节，而对于集中式软标签星形为 8C 字节；因此只要 N ≤ 4C+1（例如，C=10 时 N ≤ 41，C=1000 时 N ≤ 4001），硬标签就更便宜；超过此范围，需要稀疏八卦拓扑（§5.3 (https://arxiv.org/html/2607.00173#S5.SS3)）。

### 3.3 理论分析

三个结果解释了 TallyTrain 的性能：