Optimization of 3D U-Net Using Attention Mechanism for Accurate Protein Object Identification in Cryo-Electron Tomography

Abstract

Object segmentation in 3D tomograms is a key problem in Cryo-Electron Tomography (Cryo-ET) analysis. In this work, performance for the 3D U-Net architecture, and its variants with three attention mechanisms (Attention Gate (AG), Squeeze-and-excitation (SE), and Convolutional Block Attention Module (CBAM)) was evaluated. Experiments were conducted on a publicly available Cryo-ET dataset comprising two tomogram samples using a combination of (16,32,64,128) and (32,64,128,256) channel configurations, and with patch sizes of (32,64,64), (32,96,96), and (32,128,128) respectively. Model performance was evaluated with the F-Beta Score metric. The results of the analysis show that larger patch sizes significantly improve performance, and deep channel configurations do not always lead to better performance. Compared to the baseline 3D U-Net, which achieved a best score of 0.670, 3D UNet + SE led to the best model performance with the highest F-Beta Scores at 0.718, representing an improvement of 0.048. 3D U-Net + CBAM was second with F-Beta Scores at 0.707, improving by 0.037 over the baseline, while 3D U-Net + AG exhibited prediction inconsistency, with its accuracy falling below the baseline in multiple settings. Overall, these results show that incorporating either SE or CBAM, is a better approach to improve segmentation accuracy for 3D tomogram analysis.