FedKD-hybrid: Federated Hybrid Knowledge Distillation for Lithography Hotspot Detection

📄 arXiv: 2501.04066v1 📥 PDF

作者: Yuqi Li, Xingyou Lin, Kai Zhang, Chuanguang Yang, Zhongliang Guo, Jianping Gou, Yanli Li

分类: cs.LG, cs.AR

发布日期: 2025-01-07

🔗 代码/项目: GITHUB

💡 一句话要点

提出FedKD-hybrid，一种用于光刻热点检测的联邦混合知识蒸馏方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 联邦学习 知识蒸馏 光刻热点检测 隐私保护 模型聚合

📋 核心要点

  1. 现有联邦学习光刻热点检测方法未能充分利用和转移学习到的信息，限制了性能提升。
  2. FedKD-hybrid通过共享模型层和公共数据集，结合参数聚合与知识蒸馏，实现更有效的全局共识。
  3. 实验表明，FedKD-hybrid在ICCAD-2012和真实FAB数据集上优于现有联邦学习方法。

📝 摘要（中文）

本文提出FedKD-hybrid，一种用于分布式隐私保护设置下，基于机器学习的光刻热点检测(LHD)的联邦学习(FL)方法。目前，主要有两种方法来聚合本地模型并达成全局共识，包括基于参数的和基于非参数的（也称为知识蒸馏，即KD）。虽然这两种方法在特定场景下都显示出有效性，但我们注意到它们没有充分利用和转移学习到的信息，使得基于FL的LHD的潜力仍未被充分挖掘。因此，我们提出了FedKD-hybrid来弥补这一研究空白。具体来说，FedKD-hybrid客户端在所有参与者之间共享几个相同的层，并使用公共数据集来实现全局共识。在训练过程中，训练好的本地模型将在公共数据集上进行评估，生成的logits将与相同的层参数一起上传。聚合的信息随后被用于通过公共数据集作为媒介来更新本地模型。我们将提出的FedKD-hybrid与ICCAD-2012和FAB（真实世界收集）数据集下的几种最先进的(SOTA) FL方法在不同设置下进行了比较；实验结果表明了FedKD-hybrid算法的优越性能。我们的代码可在https://github.com/itsnotacie/NN-FedKD-hybrid 获取。

🔬 方法详解

问题定义：论文旨在解决联邦学习在光刻热点检测中的信息利用不足问题。现有的联邦学习方法，如基于参数平均和知识蒸馏的方法，在聚合本地模型时，未能充分利用各个客户端学习到的知识，导致全局模型的性能提升受限。尤其是在数据异构性较强的场景下，这个问题更加突出。

核心思路：论文的核心思路是结合参数聚合和知识蒸馏的优点，提出一种混合的联邦学习框架。通过在客户端之间共享部分模型层，并利用公共数据集进行知识蒸馏，实现更有效的知识转移和模型聚合。共享层用于参数聚合，公共数据集用于知识蒸馏，二者相互补充，从而提升全局模型的性能。

技术框架：FedKD-hybrid的整体框架包含以下几个主要阶段：1) 客户端选择：服务器选择参与当前轮训练的客户端。2) 本地训练：每个客户端使用本地数据训练模型，模型包含共享层和私有层。3) logits生成：客户端使用训练好的本地模型在公共数据集上生成logits。4) 参数和logits上传：客户端将共享层的参数和生成的logits上传到服务器。5) 全局聚合：服务器聚合客户端上传的参数和logits。6) 知识蒸馏：服务器使用聚合后的信息更新公共数据集上的logits，并将其作为知识蒸馏的目标。7) 本地模型更新：客户端使用公共数据集和蒸馏后的logits更新本地模型。

关键创新：FedKD-hybrid的关键创新在于混合了参数聚合和知识蒸馏两种联邦学习方法。传统的参数聚合方法只关注模型参数的平均，忽略了各个客户端学习到的知识。知识蒸馏方法虽然可以传递知识，但需要大量的公共数据。FedKD-hybrid通过共享模型层进行参数聚合，并利用公共数据集进行知识蒸馏，从而实现了更有效的知识转移和模型聚合。

关键设计：FedKD-hybrid的关键设计包括：1) 共享层的选择：选择哪些层作为共享层需要根据具体任务和数据集进行调整。2) 公共数据集的选择：公共数据集需要具有代表性，能够反映各个客户端数据的分布。3) 知识蒸馏的损失函数：可以使用KL散度等损失函数来衡量本地模型logits和蒸馏后的logits之间的差异。4) 客户端选择策略：可以使用随机选择或基于客户端数据质量的选择策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，FedKD-hybrid在ICCAD-2012和真实FAB数据集上均优于现有的联邦学习方法。例如，在某个实验设置下，FedKD-hybrid相比于FedAvg，在热点检测精度上提升了超过3%。这证明了FedKD-hybrid在光刻热点检测任务中的有效性。

🎯 应用场景

FedKD-hybrid可应用于光刻热点检测，降低芯片制造缺陷率，提高良率。该方法也可推广到其他需要隐私保护的分布式机器学习任务，如医疗影像分析、金融风控等，具有重要的实际应用价值和广阔的应用前景。

📄 摘要（原文）

Federated Learning (FL) provides novel solutions for machine learning (ML)-based lithography hotspot detection (LHD) under distributed privacy-preserving settings. Currently, two research pipelines have been investigated to aggregate local models and achieve global consensus, including parameter/nonparameter based (also known as knowledge distillation, namely KD). While these two kinds of methods show effectiveness in specific scenarios, we note they have not fully utilized and transferred the information learned, leaving the potential of FL-based LDH remains unexplored. Thus, we propose FedKDhybrid in this study to mitigate the research gap. Specifically, FedKD-hybrid clients agree on several identical layers across all participants and a public dataset for achieving global consensus. During training, the trained local model will be evaluated on the public dataset, and the generated logits will be uploaded along with the identical layer parameters. The aggregated information is consequently used to update local models via the public dataset as a medium. We compare our proposed FedKD-hybrid with several state-of-the-art (SOTA) FL methods under ICCAD-2012 and FAB (real-world collected) datasets with different settings; the experimental results demonstrate the superior performance of the FedKD-hybrid algorithm. Our code is available at https://github.com/itsnotacie/NN-FedKD-hybrid