A novel hybrid transformer-CNN architecture for environmental microorganism classification

Abstract

The success of vision transformers (ViTs) has given rise to their application in classification tasks of small environmental microorganism (EM) datasets. However, due to the lack of multi-scale feature maps and local feature extraction capabilities, the pure transformer architecture cannot achieve good results on small EM datasets. In this work, a novel hybrid model is proposed by combining the transformer with a convolution neural network (CNN). Compared to traditional ViTs and CNNs, the proposed model achieves state-of-the-art performance when trained on small EM datasets. This is accomplished in two ways. 1) Instead of the original fixed-size feature maps of the transformer-based designs, a hierarchical structure is adopted to obtain multi-scale feature maps. 2) Two new blocks are introduced to the transformer's two core sections, namely the convolutional parameter sharing multi-head attention block and the local feed-forward network block. The ways allow the model to extract more local features compared to traditional transformers. In particular, for classification on the sixth version of the EM dataset (EMDS-6), the proposed model outperforms the baseline Xception by 6.7 percentage points, while being 60 times smaller in parameter size. In addition, the proposed model also generalizes well on the WHOI dataset (accuracy of 99%) and constitutes a fresh approach to the use of transformers for visual classification tasks based on small EM datasets.

Fig 1. Overall architecture of HTEM.

First, the inputs are fed to a Convolutional Token Embedding (CTE) to obtain patches. Second, the feature maps are processed using a feature embedding layer and repeated transformer-encoder blocks consisting of CPSA and LFFN blocks in each stage. Finally, a global average pooling block is employed to obtain the class token.

Fig 2

(A) Multi-Head Attention (MHA) block in ViT [30]. (B) Convolutional Parameters Sharing multi-head Attention (CPSA) block in HTEM.

Fig 3. Local Feed-Forward Network (LFFN) in HTEM.

DW Conv denotes depth-wise convolution.

Fig 4. The accuracy and loss curves of the proposed HTEM model.

Fig 5. Confusion matrix of HTEM model on the test set after data augmentation.

In the confusion matrix, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 represent Actinophrys, Arcella, Aspidisca, Codosiga, Colpoda, Epistylis, Euglypha, Paramecium, Rotifera, Vorticella, Noctiluca, Ceratium, Stentor, Siprostomum, K. Quadrala, Euglena, Gymnodinium, Gymlyano, Phacus, Stylongchia, Synchaeta.

Fig 6. Images of Paramecium, Codosiga, and K. Quadrala.

Fig 7. Images of Epistylis, and Vorticella.