AutoMO-Mixer: An automated multi-objective multi-layer perspecton Mixer model for medical image based diagnosis

1 INTRODUCTION . With the development of modern medicine , medical image has become an essential way to carry out personalized and accurate diagnosis . Due to the outstanding image analysis ability , deep learning has been widely used in medical image based diagnosis and has achieved great success ( Zhang & An , 2017 ) ( Shen et al. , 2017 ) in the past years . As there are so many parameters to train in current deep learning model such as convolutional neural network ( CNN ) ( Chua , 1998 ) , a large scale dataset is always required to feed the model . In particular , the parameters in convolutional layer accounted a majority of the parameters . However , it is very difficult to collect sufficient medical image samples for a particular clinical problem in practice . Recently , a new deep learning model termed MLP-Mixer ( Mixer ) ( Tolstikhin et al. , 2021 ) which was developed based on multi-layer perception ( MLP ) was proposed . Compared with CNN , the convolutional layer is removed from Mixer , leading to decreasing the architecture parameter scale sharply and Mixer can be trained sufficiently through small scale dataset . On the other hand , Mixer can achieve similar performance with CNN ( Tolstikhin et al. , 2021 ) . In summary , Mixer is a better architecture for building medical image based diagnostic model . Furthermore , building the balanced model between sensitivity and specificity is necessary ( Banerjee et al. , 2018 ) ( Mazurowski et al. , 2008 ) in clinical diagnosis . Sensitivity indicates the proportion of people who are actually sick that are diagnosed as sick , whilst specificity indicates the proportion of people who are truly healthy who are diagnosed as healthy ( Reitsma et al. , 2005 ) . As shown in Table 1 , although the accuracy is high , sensitivity and specificity imbalance , resulting in higher rate of missed diagnosis . Therefore , a more balanced model is necessary and a multi-objective model which considers sensitivity and specificity simultaneously is needed . As such , a new automated multi-objective Mixer ( AutoMO-Mixer ) model is developed in this study . In AutoMO-Mixer , both sensitivity and specificity are considered as the objective functions simultaneously and a Pareto-optimal model set can be obtained through the multi-objective optimization ( Zhou et al. , 2017 ) in training stage . In addition , since there are several hyperparameters which may affect the model performance in AutoMO-Mixer , Bayesian optimization ( Pelikan , 2005 ) is used to train the hyperparameters . In testing stage , the Pareto-optimal models with balanced sensitivity and specificity are chosen so as to improve model diversity and stability . To obtain more reliable result , evidential reasoning ( ER ) ( Yang & Xu , 2002 ) approach is used to fuse the output of selected Pareto-optimal models to obtain final outcome . The experimental studies on two public medical image datasets demonstrated that AutoMO-Mixer can outperform Mixer and CNN , and more balanced results can be achieved as well . 2 METHOD . 2.1 FRAMEWORK . The framework of AutoMO-Mixer is shown in figure 1 , which consists of training stage and testing stage . In training stage , medical images are fed into multiple Mixer models with randomly initialization . Then sensitivity and specificity are considered as the multi-objective functions simultaneously , and The Pareto-optimal Mixer model set can be obtained through multi-objective optimization . Meanwhile , since there are several hyperparameters which may affect the model performance , Bayesian optimization is used to optimize the hyperparameters . In testing stage , models with balanced sensitivity and specificity are fused using the ER approach to obtain more reliable final prediction model . 2.2 TRAINING STAGE . Assume α = { α1 , ... , αK } represents the hyperparameters , where K is the number of hyperparameter and β denotes the parameters in Mixer models . M = { m1 , ... , mq } represents the Mixer model , where q is the number of models . The MLP-Mixer network is a visual-oriented , all-MLP architecture as shown in figure 2 . It consists consisting of four parts ( Tolstikhin et al. , 2021 ) : splitting a picture into multiple patches shaped as a “ patches×channels ” table as an input ; converting all patches into feature embedding with a set of fully connected networks ; refining feature information through N mixer layers ; and classifying them with a full-connected layer after global average pooling . Mixer is divided into token-mixer ( feature refining along the columns direction of the table , i.e . spatial locations ) and channel-mixer ( feature refining along the rows direction of the table , i.e . channels ) , all using MLP for feature extraction , each MLP consists of two layers of full-connected and an activation function GELU . Since there are two groups ’ parameters to train , the training stage is a bilevel iterative optimization process . In other words , Mixers ’ model parameters and hyperparameters in AutoMO-Mixer are updated iteratively . The training stage is shown in figure 3 . In the beginning , all the parameters are initialized randomly . When the hyperparameters are determined , the Mixers ’ model parameters are trained first . To obtain the balanced model , sensitivity and specificity are considered as objective functions simultaneously , they are : fsen = TP TP + FN ( 1 ) fspe = TN TN + FP ( 2 ) where TP is true positive , TN is true negative , FP is false positive , FN is false negative . The goal is to maximize fsen , fspe simultaneously , that is : f = max β ( fsen , fspe ) ( 3 ) To optimize the above function , an iterative multi-objective immune algorithm ( IMIA ) ( Zhou et al. , 2017 ) ( Gong et al. , 2008 ) is used . IMIA consists of six steps : initialization , cloning , mutation , deletion , update , and termination . First , the initial solution set denoted by D ( t ) = { M1 , ... , MN } is generated , where Mi = { mi1 , ... , miq } , i = 1 , 2 , ... , N . Then the models with higher Mi = { mi1 , ... , miq } , i = 1 , 2 , ... , N will be replicated using the proportional cloning method . In the third step , a probability of mutation is randomly generated for each model , and the model performs mutation when its probability is larger than the mutation probability . After the mutation , the new models are generated . If the models have same sensitivity and specificity , only one model is remained and the all the other models are removed . Then the model size is kept through AUC based non-dominated sorting strategy . This training process will not stop until the maximum number of iterations is reached . Then the Pareto-optimal Mixer model set is generated , where the model set size is J . Now the optimal parameters for Mixer are obtained and hyperparameters are optimized . Since there are hyperparameters in Pareto-optimal model selection and fusion phase , the weight for each Pareto-optimal model should be estimated , which is denoted by wj . As the balanced model between sensitivity and specificity is desired , the ratio between them is considered , which is fsenfspe or fspe fsen . When the ratio is less than 0.5 or greater than 1 , the model is considered as extreme imbalance , setting wj as 0 . The expression of wj is as follows : wj = λ fjsen fjspe + ( 1− λ ) AUC , when 0.5 ≤ f j sen fjspe ≤ 1 λ fjspe fjsen + ( 1− λ ) AUC , when 0.5 ≤ f j spe fjsen ≤ 1 , j = 1 , 2 , ... , J , 0 Other situations ( 4 ) where λ indicates the importance of balance , and 1-λ indicates the importance of AUC . After calculating the wj for each model , the weights are normalized , that is : J∑ j=1 wj = 1 , 0 ≤ wj ≤ 1 ( 5 ) Then the Pareto-optimal models are fused through ER approach . Assume there are J models in Pareto-optimal model set , and the output probability ( normal and abnormal ) for each model is denoted by in Pj = { P 1j , P 2j } , j = 1 , 2 , ... , J , where P 1j is the output probability of abnormal , and P 2j is the output probability of normal . The final output probability P cfin , c = 1 , 2 is obtained through the ER fusion strategy ( Wang et al. , 2006 ) , that is : P cfin = ER ( P c j , wj ) , j = 1 , 2 , ... , J , c = 1 , 2 ( 6 ) where ER is : P cfin = µ× [ J∏ j=1 ( wjP c j + 1− wj ) − J∏ j=1 ( 1− wj ) ] 1− µ× [ J∏ j=1 ( 1− wj ) ] , c = 1 , 2 ( 7 ) The normalized factor µ is : µ = [ 2∑ c=1 J∏ j=1 ( wjP c j + 1− wj ) − J∏ j=1 ( 1− wj ) ] −1 ( 8 ) The hyperparameter α can be optimized through Bayesian optimization ( Pelikan , 2005 ) ( Springenberg et al. , 2016 ) , where the objective function is : fH = min α ( 1−AUC ) ( 9 ) Eqs . ( 3 ) and ( 9 ) are optimized iteratively until it reaches the termination critation . 2.3 TESTING STAGE . An illustration of testing stage is shown in figure 4 . The medical images are fed into the trained Pareto-optimal Mixer models and the output probabilities from each model can be obtained . Then ER approach is used to fuse the outputs of all the Mixer models and final output can be obtained . 3 EXPERIMENTS . 3.1 MATERIALS AND SETUP . The two public databases used in this paper are MRI images of brain tumors ( Panigrahi , 2021 ) and CT images of COVID-19 ( LuisBlanche , 2020 ) . The first dataset aims to help people for detecting brain tumors . It contains MRI scans of the brain , and each type of brain has 1500 images . We selected 900 images for modeling . The next dataset is collected from COVID19-related papers from medRxiv , bioRxiv , NEJM , JAMA , Lancet , etc . It contains 349 CT scans that are positive for COVID-19 , This dataset can be used to perform classification and automatically detect COVID-19 on CT scans . The distribution of the two types of data is shown in Table 2 . The examples of two datasets are shown in figure 5 . Before feeding the model , all the images are resized into 224 x 224 images , and each image is divided into 14 x 14 patches , where each patch size is 16 x 16 . There are two hyperparameters , Table 2 : Distribution of two types of data Train Train Test Test Abnormal Normal Abnormal normal Brain Tumor 360 360 90 90 COVID-19 280 318 69 79 Figure 5 : The examples of two datasets . MP and λ , which the range is set between 0 and 1 . The network structure contains five parameters , they are : the number of Mixer layers , the number of channels ( Hidden size C ) , the number of hidden layer neurons of token-mixing ( MLP dimension Ds ) , and the number of hidden layer neurons of channel-mixing ( MLP dimension Dc ) . The settings are shown in table 3 . In comparative study , Mixer and CNN model are evaluated . For fairness of comparison , the number of layers of the two models is required to be equal . We make the number of layers five . The four parameters of the Mixer network are set to 5 , 256 , 392 , 1024 , respectively . AUC ( Zhao et al. , 2011 ) is an ideal indicator to evaluate the reliability of the prediction results , and the larger the AUC shows that the better the classification of the model . Accuracy ( ACC ) represents the correct rate of diagnosis and is also an important indicator for evaluating the model . Therefore , sensitivity ( SEN ) , specificity ( SPE ) , AUC , ACC are used for evaluation . All the experiments are performed five times and the average results are shown .

Summary: The authors proposed a new method for learning imbalanced medical image datasets. They propose to use the recently propose MLP-Mixer instead of CNN  to reduce the number of trainable parameters and thus can be trained with fewer images as compared to CNNs. To handle the class imbalance in the datasets, they use multi-objective loss functions to balance between sensitivity and specificity. Bayesian optimation is used for finding the optimal hyperparameters. The proposed method is evaluated on two medical image datasets, showing improvements over CNN and Mixer baselines.