Clustering Federated Learning for Wafer Defects Classification on Statistical Heterogeneous Data

Abstract

Data-driven deep learning techniques for wafer defect image classification provide wafer manufacturers with a tool to rapidly identify surface defects. However, the defect data and computational capabilities of a single wafer manufacturer are often insufficient to support the training of deep learning models. In response, we introduce federated learning (FL), a paradigm that leverages the data and computational capabilities of various wafer manufacturers, all while ensuring that the original data from different manufacturers remain unexposed to each other. Due to variations in manufacturing processes and image acquisition equipment, identical wafer defects can exhibit different features in different manufacturing settings, leading to statistically heterogeneous datasets. This heterogeneity can reduce model convergence speed and accuracy. To counteract this issue, we propose a personalized FL approach with clustering. In the personalization phase, we train distinct network layers for each client’s local model, capitalizing on the feature extraction capability of the global model’s shallow network, while also achieving commendable performance on each client’s unique dataset. During the clustering phase, we provide a theoretical analysis, demonstrating that the divergence of weights between two models is bounded above, laying a theoretical foundation for the clustering operation. We then enhance a density-based clustering method, enabling the clustering of clients with similar data features without the need to specify the number of cluster centers, thus mitigating the problem of global model oscillation. We have conducted experiments under various data heterogeneity scenarios. The experiments show that our method can achieve a 2.8% accuracy improvement average versus the compared state-of-the-art federated methods with a faster convergence rate.