FlatLand: 通过定制洛伦兹空间的个性化图联邦学习

# 基于定制洛伦兹空间的个性化图联邦学习  
来源：https://arxiv.org/html/2608.21096  
**作者**：刘嘉宏  
**单位**：香港中文大学计算机科学与工程系，中国香港特别行政区  
**联系邮箱**：[[email protected]](mailto:[email protected])  
傅新宇  
华为技术有限公司，中国香港特别行政区  
杨梦琳  
人工智能学部，香港科技大学（广州），中国广东省广州市  
张伟熙  
华为技术有限公司，中国香港特别行政区  
Rex Ying  
计算机科学系，美国耶鲁大学，美国康涅狄格州纽黑文  
Irwin King  
计算机科学与工程系，香港中文大学，中国香港特别行政区  
**联系邮箱**：[[email protected]](mailto:[email protected])  

###### 摘要  
联邦学习支持隐私保护的协同训练，但高度异构的客户端数据仍是难题，尤其在客户端持有结构多样化图的图联邦学习中尤为突出。现有个性化联邦学习方法忽略了不同图结构的内在几何特性。我们提出 **FlatLand**，一种新颖的个性化联邦学习方法，将不同客户端的数据嵌入到双曲几何中的**定制洛伦兹空间**。我们的核心洞见是：双曲几何天然适配现实世界图中普遍存在的负曲率特性，而洛伦兹空间中的**类时维度**为编码客户端特定异构性提供了一种合理途径。我们开发了**参数解耦策略**，将捕获异构性的类时参数与保留共性知识的类空参数分离，从而无需客户端相似性估计或额外计算模块即可直接聚合。在多种联邦图学习任务上的实证结果表明，**FlatLand** 尤其在低维设定下取得了更优的性能。代码已开源：[GitHub仓库](https://github.com/HUBERILT/FlatLand_ICML)。  

###### 关键词  
联邦学习，图神经网络，双曲几何，个性化  

## 1 引言  
**图示**：图1：示例：(a) 来自三个CiteSeer客户端的度分布KDE (11)，(b) 对应不同洛伦兹尺度参数 \(K\) 的二维洛伦兹空间。  
联邦学习是一种在多客户端间实现协同机器学习并保护数据隐私的范式。传统联邦学习面临数据异构性挑战，单一模型难以满足多样化的本地需求 (57)。这一挑战在图联邦学习中更为突出，复杂的拓扑结构导致客户端间存在显著的结构异构性 (61；58)。严重时，联邦学习甚至可能低于本地训练性能 (5)。为解决异构性，个性化联邦学习通过共享通用模型知识并允许客户端特定适配来应对。当前针对图数据的PFL方法主要在聚合时采用三种策略处理异构性：(1) 参数解耦：将模型拆分为共享与个性化部分 (45；58)；(2) 客户端相似性估计：分析权重或梯度以评估客户端相似度 (61)；以及 (3) 辅助模块计算：引入额外模块以区分全局有益参数与客户端特定参数 (5)。尽管这些方法有效，但现有PFL方法受限于欧几里得空间，默认所有客户端数据分布具有统一平坦几何。这一假设需要设计复杂机制来解决异构性。简单的参数解耦常失效，而更先进的技术（如客户端相似性估计或辅助模块集成）虽提升性能，但计算开销显著。因此，我们通过 **Ricci曲率** 的几何视角重新审视PFL (16；56)，该曲率刻画了图结构的内在特性：其符号指示双曲（负）、平坦（零）或球形（正）几何，其大小反映空间弯曲程度。  

**观察**。我们在多个真实数据集上的实证分析揭示了两个关键观察（图3和图2）：(1) 客户端图普遍呈现**负**Ricci曲率，表明其固有双曲结构；(2) 不同客户端的曲率值**差异显著**，揭示了超越简单统计差异的内在几何异构性。这些发现表明，现有方法中统一的欧几里得几何假设与图数据的真实几何本质存在根本性偏差 (48；1；28；58；22)，导致表征次优且异构性建模复杂化。  

**图2**：跨数据集（Cora、ogbn-arxiv 和 Amazon-Photo）的平均Forman-Ricci曲率。柱状图越高，表明数据集的非欧几里得特征越显著。  

为突破单一欧几里得几何，我们从理论上论证了**洛伦兹几何**对PFL的优势（第4.1节），表明其天然适用性源于两点：  
首先，**增强的表征能力**：洛伦兹空间能以低失真方式表征估计的内在图属性 (48；50；4；63)。当客户端图呈现不同Ricci曲率时，为每个客户端分配适当的双曲曲率支持更忠实的建模（定理4.1）。如图1(a)所示，分布呈长尾且偏度不同。特别是，客户端1更“陡峭”，受益于具有较大曲率幅度（更小 \(K\)）的洛伦兹空间，该空间提供了更“宽敞”的嵌入环境以分离尾节点。  
其次，**自然的异构性编码**：洛伦兹空间中额外的类时维度提供了载体，用于捕获客户端间的内在几何异构性（定理4.3）。在示例图1(b)¹中，诸如“尾节点与头节点间不平衡程度如何？”的异构属性可通过类时维度 \(x_t\) 在洛伦兹空间自然区分，而共性信息（如“星形是尾节点”）则保留在类空维度 \(\mathbf{x}_s\)（“平坦地带”）中作为共享节点表征。  
基于上述洞见，我们提出 **FlatLand**，一种探索性PFL框架，将客户端数据嵌入定制洛伦兹空间以忠实捕获其内在几何（第5节）。然而，仅嵌入本身无法解决参数聚合中的异构性挑战。因此，我们进一步利用类时维度的特性，开发了一种**理论驱动的参数解耦策略**，将异构性相关参数指定为个性化参数，仅聚合承载共享信息的参数。该设计无需辅助模块或客户端相似性估计，有效缓解异构性，同时保持洛伦兹几何的有效性。据我们所知，这是首次以原则性方式将双曲几何与PFL结合以解决客户端异构性的研究，提供了利用几何特性的新颖**简洁**且**高效**的视角。实验结果表明，**FlatLand** 在低维设定（对通信高效的联邦学习至关重要）下的性能优于其欧几里得对应方法。  

## 2 相关工作  
##### 图上的个性化联邦学习  
个性化联邦学习通过学习客户端适配模型（而非单一全局模型）来解决统计异构性 (26；15；25；10；7)。对于图数据，现有的个性化联邦图学习方法通常通过梯度聚类客户端 (61)、引入额外个性化模块 (58) 或估计客户端相似性以进行定制聚合 (5)。这些方法有效但通常在欧几里得空间运行，依赖客户端相似性估计或辅助组件处理异构性。此类设计未显式建模现实世界图中广泛存在的无标度和层次结构 (1；30)，促使我们采用几何感知的个性化图联邦学习方法。  

##### 双曲联邦学习  
双曲表征为层次化和幂律数据提供了自然几何 (48；6)。近期联邦方法利用双曲距离进行知识蒸馏 (2)、双曲原型进行非独立同分布学习 (38)，或在FedAvg式图联邦学习流程中使用双曲图神经网络 (14)。然而，这些方法未对底层客户端几何进行个性化，或在聚合时未将客户端特定几何信息与共享知识分离。**FlatLand** 的不同之处在于为每个客户端分配定制洛伦兹空间，并将类时个性化参数与类空共享参数解耦，无需客户端聚类或额外相似性估计即可直接聚合。更详细的讨论见附录A。  

## 3 预备知识  
##### 双曲几何的洛伦兹模型  
洛伦兹模型（又称双曲面模型）是 **黎曼** 双曲空间在平坦闵可夫斯基环境空间 \(\mathbb{R}^{d+1}\) 中的表示 (48；9)。给定一个具有常负曲率 \(-1/K\ (K>0)\) 的 \(d\) 维洛伦兹流形 \(\mathcal{L}_{K}^{d}\)，设点 \(\mathbf{x} \in \mathcal{L}_{K}^{d}\)，其形式为 \(\mathbf{x} = \begin{bmatrix} x_{t} & \mathbf{x}_{s} \end{bmatrix}^{\top} \in \mathbb{R}^{d+1}\)，其中第一维 \(x_{t} \in \mathbb{R}\) 称为**类时维度**，其余维 \(\mathbf{x}_{s} \in \mathbb{R}^{d}\) 称为**类空维度**。它满足以下条件：\(\langle \mathbf{x}, \mathbf{x} \rangle_{\mathcal{L}} = -K\) 且 \(x_{t} > 0\)，其中 \(\langle \mathbf{x}, \mathbf{y} \rangle_{\mathcal{L}} = -x_{t}y_{t} + \mathbf{x}_{s}^{\top}\mathbf{y}_{s}\) 是洛伦兹内积。这里 \(K\) 是正的洛伦兹尺度参数；不同的 \(K\) 值诱导不同的双曲几何，其截面曲率为 \(-1/K\)。形式化定义见附录B.1。  

通常，输入位于欧几里得空间，需要映射到双曲空间。将欧几里得空间数据 \(\mathbf{v}^{E} \in \mathbb{R}^{d}\) 投影到洛伦兹空间 \(\mathbf{x} \in \mathcal{L}_{K}^{d}\) 的方式可简化为 ²  
² 为清晰起见，所有洛伦兹空间嵌入记为 \(\cdot^{H}\)。具体地，若洛伦兹尺度参数 \(K\) 已知，记为 \(\cdot^{K}\)。欧几里得空间嵌入记为 \(\cdot^{E}\)。所有向量 \(\mathbf{x}\)，若无上标，均假设位于洛伦兹空间。  
\[
\mathbf{x}^{K} = \exp_{\mathbf{o}}^{K}\left(\mathbf{v}^{E}\right) = \exp_{\mathbf{o}}^{K}\left(\left[0, \mathbf{v}^{E}\right]\right) = \left[x_{t}, \mathbf{x}_{s}\right]^{\top},
\]  
\[
x_{t} = \sqrt{K} \cosh\left(\frac{\|\mathbf{v}^{E}\|_{2}}{\sqrt{K}}\right), \quad \mathbf{x}_{s} = \sqrt{K} \sinh\left(\frac{\|\mathbf{v}^{E}\|_{2}}{\sqrt{K}}\right) \frac{\mathbf{v}^{E}}{\|\mathbf{v}^{E}\|_{2}}.
\]  
不同的 \(K\) 将 \(\mathbf{v}^{E}\) 映射到不同的洛伦兹曲面。  

##### 全洛伦兹神经网络  
全洛伦兹网络 (9) 因减少了空间投影需求而被视为PFL的理想选择，提升了计算效率。这些网络还结合了洛伦兹变换（推移与旋转），改善了数据异构性处理和参数可解释性（附录B.3）。给定输入向量 \(\mathbf{x} \in \mathcal{L}_{K}^{n}\) 和待优化的线性层矩阵 \(\mathbf{W} \in \mathbb{R}^{m \times (n+1)}\)，全洛伦兹线性层可通式表示为 \(\mathrm{LT}\)：  
\[
\mathrm{LT}\left(\mathbf{x};f;\mathbf{W}\right) := \left(\sqrt{\|f(\mathbf{W}\mathbf{x})\|^{2} + K}, f(\mathbf{W}\mathbf{x})\right)^{T},
\]  
其中 \(f\) 是如激活、丢弃和偏置等函数。  

##### 问题陈述  
给定客户端集合 \(\mathcal{C} = \{1,2,\ldots,C\}\)，每个客户端拥有数据集 \(\mathcal{D}_{c} = \{(\mathbf{x}_{i}^{c},y_{i}^{c})\}_{i=1}^{N_{c}}\) 和分布 \(p_{c}(\mathbf{x},y)\)，当任意客户端对 \(i \neq j\) 满足 \(p_{i}(\mathbf{x},y) \neq p_{j}(\mathbf{x},y)\) 时，个性化联邦学习面临异构性，导致性能下降。PFL的目标是优化个性化模型 \(f_{c}\left(\cdot; \bm{\theta}_{c}, \bm{\theta}_{s}\right)\)。