🤖 AI Summary
This work addresses the lack of effective aggregation methods for memory-augmented autoencoders (MemAEs) with attention mechanisms in federated learning, which hinders efficient unsupervised anomaly detection under non-IID and resource-constrained settings. To this end, we propose an attention-layer-oriented guided aggregation strategy—the first federated aggregation mechanism specifically designed for MemAEs. Our approach significantly enhances model robustness and detection performance across multiple edge nodes with imbalanced data distributions while maintaining a lightweight architecture. Experimental results demonstrate that the proposed method outperforms existing baselines on various non-IID datasets, achieving higher anomaly detection accuracy and validating the feasibility of shallow MemAEs in edge computing scenarios.
📝 Abstract
Attention layers are the backbone of today's most powerful and impactful models. Models with multi-million and billion parameters rely on contextual knowledge provided by attention layers. However, their use goes well beyond just being the core component of large language models. One particularly interesting application is in Memory Augmented Autoencoders (MemAE), specifically for unsupervised representation learning in outlier detection tasks. It was shown that attention helps these models be more effective in centralized learning scenarios. Our work aims to address the lack of specialized aggregation techniques in Federated Learning (FL) when it comes to MemAE models. In this paper we analyze the intricacies of the architecture behind Memory Augmented Autoencoders, and propose novel, guided approaches to effectively aggregate these models in federated scenarios. We demonstrate our approach on non-IID datasets and show that these novel aggregation schemes are more robust when dealing with numerous edge nodes in environments with unbalanced datasets, specifically for unsupervised anomaly detection scenarios. This approach improves the performance of even very shallow autoencoders, allowing them to be used in resource constrained environments.