基于粒子群优化的注意力自适应U-Net架构搜索

Attention-Adaptive U-Net Architecture Search via Particle Swarm Optimization

摘要: U-Net通过跳跃连接对多尺度特征的有效融合，已成为医学图像分割的主流方法。然而，其架构超参数(卷积核大小、通道数)和注意力机制选择高度依赖人工设计，缺乏自动化优化手段。若将注意力选择纳入搜索空间，会导致维度爆炸和计算开销激增。针对此问题，文章提出PSO-AttUNet，通过粒子群优化(PSO)搜索U-Net架构超参数，并结合基于计分的注意力机制为每层独立选择最合适的注意力类型(SE、CBAM、AG)，实现架构搜索与注意力选择的并行优化。在BUSI乳腺超声数据集上的实验表明，PSO-AttUNet在Dice、IoU、HD95等指标上优于U-Net、Attention U-Net等基线模型，同时参数量降低至U-Net的1/47，验证了方法的有效性。

Abstract: U-Net has become the mainstream method for medical image segmentation due to its effective fusion of multi-scale features through skip connections. However, its architectural hyperparameters (kernel sizes, channel numbers) and attention mechanism selection heavily rely on manual design, lacking automated optimization methods. Incorporating attention selection into the search space would lead to dimensional explosion and significant computational overhead. To address this issue, this paper proposes PSO-AttUNet, which searches U-Net architectural hyperparameters through Particle Swarm Optimization (PSO) and independently selects the optimal attention type (SE, CBAM, AG) for each layer using a score-based attention evaluation mechanism, achieving parallel optimization of architecture search and attention selection. Experiments on the BUSI breast ultrasound dataset demonstrate that PSO-AttUNet outperforms baseline models such as U-Net and Attention U-Net in metrics including Dice, IoU, and HD95, while reducing the parameter count to 1/47 of U-Net, validating the effectiveness of the proposed method.