Students Need More Attention: BERT-based AttentionModel for Small Data with Application to AutomaticPatient Message Triage

Small and imbalanced datasets commonly seen in healthcare represent a challenge whentraining classifiers based on deep learning models. So motivated, we propose a novel framework based on BioBERT (Bidirectional Encoder Representations from Transformers forBiomedical TextMining). Specifically, (i) we introduce Label Embeddings for Self-Attentionin each layer of BERT, which we call LESA-BERT, and (ii) by distilling LESA-BERTto smaller variants, we aim to reduce overfitting and model size when working on smalldatasets. As an application, our framework is utilized to build a model for patient portalmessage triage that classifies the urgency of a message into three categories: non-urgent,medium and urgent. Experiments demonstrate that our approach can outperform sev-eral strong baseline classifiers by a significant margin of 4.3% in terms of macro F1 score. The code for this project is publicly available at \url{https://github.com/shijing001/text_classifiers}.

Introduction
Online patients portals, e.g., MyChart by Epic, have become increasingly prevalent tools for communication between patients and healthcare providers (Ramsey et al., 2018). These portals have the potential to boost the productivity of providers, improve patient satisfaction, and reduce communication barriers (Goldzweig et al., 2013;Sieck et al., 2017). Despite these benefits, patient-provider communication tools have produced unintended consequences such as increased, often unpaid, workload for providers (Hefner et al., 2019). Further, due to the flood of non-urgent incoming patient messages, some, which may require a timely emergency response, can be delayed or effectively neglected.
To address these concerns, we consider the task of automated patient message classification to estimate the urgency of messages based on their content. Our dataset consists of 1, 756 messages collected from a University Hospital's portal. These messages, manually adjudicated by experienced healthcare providers, were grouped into three categories: non-urgent, medium, and urgent, as summarized in Table 1. With merely 170 urgent instances, our dataset is both small and imbalanced, posing a significant challenge in terms of properly training machine-learning-based classifiers. This challenge is common in many clinical datasets (Zhao et al., 2018;Para et al., 2019) due to the fact that manually labeling data is very laborious, time-consuming, expensive and oftentimes prohibitive. Further, certain labels are rare by nature, for instance, urgent electronic messages are far less common because patients tend to call the healthcare provider directly rather than using the portal.  Machine learning approaches have been previously applied to message classification tasks in the healthcare domain. For example, Cronin et al. (2015) utilized logistic regression and random forest algorithms to classify between patient health information needs such as symptom management and medication side effects, based on their messages.  studied the use of rule-based and random forest classifiers to classify patient portal messages into broad communication types, e.g., appointment rescheduling, examination enquiry, etc. Sulieman et al. (2017) showed that Convolutional Neural Networks (CNNs) outperform traditional classifiers in portal message classification, and Tafti et al. (2019) developed an ensemble of neural networks for text classification to categorize free-text patient portal messages as either containing active symptom descriptions or logistic requests. Chen et al. (2019) leveraged traditional machine learning methods such as Support Vector Machines (SVMs) to detect hypoglycemia incidents reported in patient messages. To the best of our knowledge, most existing classifiers employ either traditional machine learning methods such as SVMs or shallow networks for the classification of patient generated messages.
Recently, the field of natural language processing (NLP) has seen significant progress in large-scale pre-trained language models, which often considerably boost the performance of classifiers on text classification tasks (Devlin et al., 2018;Radford et al., 2019). Here, we introduce the use of Bi-directional Encoder Representations from Transformers (BERT) (Devlin et al., 2018) for patient message classification. BERT produces sequence representations using a multi-layered attention mechanism, which models associations between the tokens (words) of a sentence (Vaswani et al., 2017). We consider the challenging task of properly training such complex models on a small and imbalanced dataset, which is complicated because these models typically have hundreds of millions of parameters. Previous studies addressed this challenge by pre-training BERT on a large dataset and then finetuning it on the (small) dataset at hand (Devlin et al., 2018). However, the performance of this procedure is usually limited by the size and difficulty of the available small dataset.

Generalizable Insights from Machine Learning in the Context of Healthcare
In this paper we address this challenge with a framework which beyond the use of a pretrained model, has two novel components. First, we introduce a novel attention mechanism that utilizes label embeddings to better capture associations between the labels and the tokens of the input sequence. Importantly, the proposed attention mechanism can be incorporated into existing attention layers with minimum modification, allowing the use of pre-trained BERT or BioBert models (Devlin et al. (2018); Lee et al. (2020)). We term our method LESA-BERT, short for Label Embedding on Self-Attention in BERT. Second, a large deep learning model like BERT has millions of parameters and tends to overfit on small datasets. We hypothesize that a model with a reduced number of parameters may result in a smaller generalization error, thus likely improved performance. Accordingly, we employ the knowledge distillation technique to compress the LESA-BERT model, by training a smaller student model to reproduce the prediction ability of the teacher, a fine-tuned LESA-BERT model. From our framework, we devise distilled variants of LESA-BERT. We demonstrate LESA-BERT by building an automatic message triage classifier that can predict the urgency of messages based on their content. Experiments demonstrate that LESA-BERT and distilled variants result in improved performance compared to multiple baseline approaches. Specifically, LESA-BERT outperforms the baselines by a 2.5% margin in terms of macro F1 score, while the distilled LESA-BERT with 6 encoder layers (Distil-LESA-BERT-6) further outperforms LESA-BERT by 1.8%.

Related Work
BERT Widely used for natural language processing, BERT produces sentence representations through a multi-layered attention mechanism that encodes relations between the tokens that compose the sentence (Vaswani et al., 2017;Devlin et al., 2018). Due to its complex structure, which encompasses a large number of parameters, BERT is typically pretrained on a large dataset and then fine-tuned on the target dataset. Standard pre-training approaches include unsupervised tasks such as masked language modeling and next sentence prediction, which leverage massive unlabeled, general domain corpora like English Wikipedia (2,500M words) or BooksCorpus (800M words) (Zhu et al., 2015). Subsequently, the model is fine-tuned by adding one or more additional layers, the parameters of which are optimized using the target dataset, in a technique termed transfer learning. Recently BERT-based models have been leveraged to NLP tasks in healthcare. Huang et al. (2019) developed ClinicalBERT by pre-training BERT on a large set of clinical notes, and utilized it to predict patients' 30-day hospital readmission based on both discharge summaries and clinical notes. Lee et al. (2020) introduced BioBERT, which is essentially BERT pre-trained on large biomedical corpora that includes PubMed abstracts and PubMed Central full-text articles. They showed that pre-training the model directly on the domain of interest leads to improved performance on various biomedical NLP tasks.
Label Embedding Label embedding is a technique that embeds class labels along with the (text) data into a joint latent space, where the model can be trained to cross-attend the inputs and labels to boost the performance of deep learning models. Label embeddings were previously leveraged for image classification (Akata et al., 2015), multi-modal learning between images and text (Kiros et al., 2014), text recognition in images (Rodriguez-Serrano and Perronnin, 2015), zero-shot learning (Li et al., 2015;Ma et al., 2016) and text classification Zhang et al. (2017). Notably, Wang et al. (2018) proposed a framework named Label Embedding Attentive Model (LEAM), which jointly embeds the words and labels in a common latent space, and improves the performance on general text classification tasks. Inspired by LEAM, we consider the joint representation of the message and its corresponding class token, and propose to incorporate label embeddings to the self-attention mechanism inside BERT encoders, which improves the attention between the class token to the other tokens in the message.
Knowledge Distillation Knowledge distillation is a technique used to train a small model usually called student to reproduce the predictions, thus performance, of a larger model called a teacher (Hinton et al., 2015). It is typically used to compress models, which reduces their storage and computational costs, thus facilitating their deployment. Model distillation has been previously applied to compress complex models such as BERT and LSTMs (Sanh et al., 2019;Tang et al., 2019). When the training data size is small, we hypothesize that the teacher model may be prone to overfitting. In this situation, knowledge distillation is more likely to produce a student model that outperforms the larger, more complex teacher model, as we show in our experiments.

Background
BERT is an architecture composed of a stack of transformer encoder layers, each including two sub-layers: a multi-head self-attention module and a feed forward network. Each of these encoder layers, shown in Figure 1(b), is structured as a residual block with appropriate layer normalization and dropout Ba et al., 2016). In Figure 1(a) we show how to fine-tune BERT-based models on the message urgency dataset to build a classifier. First, each text sequence of length L is prepended with the special token [CLS]. The sum of token and position embeddings, [x [CLS] , x token 1 , . . . , x token L ], are represented as [e [CLS] , e token 1 , . . . , e token L ] and [e 0 , . . . , e L ] vectors, respectively. We use [h [CLS] , h token 1 , . . . , h token L ] vectors to denote the output, high-level representation of classification and input tokens, each of which has the same dimensionality as the input. After applying N (usually N = 12) different encoder layers (of the same structure in Figure 1(b) but different weight matrices), we obtain the high-level representations (output features) of the input sequence, in which each token representation contains information of other tokens. Finally, the first row of the output features, h CLS , is considered as the global sequence aggregator and thus fed through a softmax classification function composed of one or more fully connected layers. The multi-head self-attention module is an ensemble of multiple attention modules sharing the same formulation. Given a text sequence t represented as embedding matrix X ∈ R (L+1)×D , L represents the length of the text sequence and D the token and position embedding dimensions. Note that the first row in X corresponds to the [CLS] token, hence the L + 1 rows in X. For a single-attention head, tokens of [CLS] and the input sequence are first mapped into the key, query and value triplets, denoted as matrices K ∈ R (L+1)×d , Q ∈ R (L+1)×d and V ∈ R (L+1)×d , respectively, via: where {W K , W Q , W V } ∈ R D×d are learnable parameters for the key, query and value of self-attention. With K, Q and V, the attention mechanism can be formulated as where i = 1, . . . , h, h is the number of attention heads, softmax(·) is the softmax function applied row-wise. A is the attention score matrix representing the compatibility of K and Q obtained via inner products. The multi-head self-attention is defined by concatenating and projecting the representation of each head as where [·, ·] denotes column-wise concatenation and W ∈ R (d×h)×D is a learnable projection matrix. After the multi-head self-attention module, the position-wise (one token at the time) feed forward network module in Figure 1(b) composed of two fully connected layers is applied, where max(0, ·) is the standard ReLU activation function, {W 1 , W 2 , b 1 , b 2 } are learnable parameters, and u is the layer normalized residual block u = LayerNorm(x + o), where x (rows of X) and o (rows of O) are the inputs and outputs of the multi-head self-attention module in (1)-(3), respectively, and the LayerNorm(·) operator is implemented according to (Ba et al., 2016).

Multi-head Attention with Label Embedding
For the patient message triage task, the labeled patient text data is usually limited and hard to obtain. Specially, data for the urgent class is scarce but of the highest priority for the application. Therefore, we incorporate label embeddings as prior information into the self-attention modules, so that the model can more easily attend to class-representative keywords. The modified module is shown in Figure 2(a) and described below. We asked healthcare providers to select a set of keywords or short phrases associated with the classes medium and urgent. These are shown in Table 2. The label embeddings of each class is initialized as the average of the corresponding keyword embeddings. The label embedding for non-urgent class is initialized at randomly provided there are no immediately obvious keywords for this class. Below we describe how to incorporate label embeddings into the multi-head self-attention in each encoder layer of BERT. Note that from Figure 1(a), the text sequence t with L tokens, prepended with the token [CLS] as previously described, has embedding matrix X = [x [CLS] , X w ] ∈ R (L+1)×D , where the x CLS and X w are the input embedding vector for  [CLS] and matrix for all other input tokens in the sequence, respectively. Subsequently, the query, key and value matrices in each attention head can be represented equivalently to (1) as where [·; ·] denotes row-wise concatenation, and {q [CLS] , k [CLS] , v [CLS] } and {Q w , K w , V w } are the query, key and value triplets of the [CLS] token and the other input tokens, respectively. From (2), the original attention weights (before softmax) of the augmented sequence can be written as In A BERT , which is equivalent to (2), the cross-attention between the [CLS] token and all input tokens is denoted as S q T [CLS] K T w ∈ R 1×L . In LESA-BERT, we introduce label embeddings to boost the model's attention to keywords associated with different labels. We incorporate label embeddings to self-attention in three steps. First (Step 1), we compute the cross attention between the label embeddings and the message tokens where X l ∈ R 3×D is a matrix containing the three label embeddings, i.e., non-urgent, medium and urgent, which are encoded into label queries Q l via the same W Q as in (5). Next ( Step 2), we compute a modified cross-attention row vector S as where we concatenate S and A l by row and then keep the maximum value of each column. As a result, S represents the maximum attention score of a input token with both the [CLS] token and the label embeddings. Finally ( Step 3), we obtain the attention weights in LESA-BERT by replacing S by S in (6), thus obtaining In (10), when a token is highly relevant to one of the labels it will result in a larger attention score with [CLS] in S , thus the [CLS] embedding will be less affected by irrelevant information in the sequence, unlike (2) where only attention from the current [CLS] embedding is considered. The proposed attention layer is shown in Figure 2(b). The attention score matrix A in (2) is replaced as A LESA−BERT in (10). All other components are exactly the same as the original encoders in BERT as shown in (1)-(4). Note that by incorporating label embeddings to the self-attention layer in BERT model, (9) allows the model to account for label information across all layers. We share the same label embedding for all the layers. The label embedding is adapted to different layers via each layer's W Q s in the multi-head attention module. The label embeddings include trainable parameters, which are tuned simultaneously with other parameters in BERT. If keywords for labels are available, then label embeddings are initialized accordingly, otherwise, label embeddings could also be initialized at random (Wang et al., 2018). All other parameters can be initialized from the pre-trained Bert or BioBert model. Though LESA-BERT is motivated by our application, it can be employed for general text classification tasks.
To the best of the authors' knowledge, LESA-BERT is the first work to incorporate label embeddings to perform self-attention in BERT encoders. Note that this technique could also be applied to models that have the self-attention components such as GPT-2, and transformers (Radford et al., 2019).

Knowledge Transfer with Model Distillation
Knowledge distillation (Buciluǎ et al., 2006;Hinton et al., 2015) is a compression technique, in which a compact model, called the student model, is trained to reproduce the behavior of a larger teacher model, as in Figure 3. In our experiments, we use a loss function defined over the cross-entropy between the teacher and the student class probabilities, given by: where z t and z s are vectors of the output logits from the teacher and student classification networks, respectively, and T 0 is a parameter controlling the degree to which we focus on the class with the highest probability, usually set to 1. When the probabilities softmax(z t /T 0 ) from the teacher model are similar to those of the student model, the cross entropy loss in (11) is small. Therefore, the cross entropy loss measures the difference between output probabilities of teacher and student classifiers. The knowledge distillation technique is typically used to compress large models such as BERT in order to reduce their computational and memory costs; usually to allow their deployment (Buciluǎ et al., 2006;Hinton et al., 2015;Sanh et al., 2019;Tang et al., 2019;Strubell et al., 2019). Here, we hypothesize that, beyond these advantages, knowledge distillation can reduce overfitting in the case of small datasets as is our case. Specifically, we expect a smaller generalization error of the distilled model due to its smaller size. According to this hypothesis, we propose a distilled variant of the LESA-BERT: a small BERT model with fewer encoder layers (student) trained to reproduce the behavior of the finetuned LESA-BERT (teacher), which has 12 layers. In our experiments, we indeed show the improved performance of the distilled model.

Cohort
In this work, we utilized 1,756 web portal messages generated from 10/2014 to 08/2018 by adult patients (> 18 years old) of a large academic medical center. The Electronic Health Record (EHR) system (Epic Verona, WI, USA) with associated patient portal (My-Chart) was the source of all patient messages. A custom-built Application Programming Interface (API) securely made available the portal messages from the EHR enterprise data warehouse into a highly protected virtual network space offered by the medical center. Approved users were allowed access to work with the identifiable protected health information. These messages included free, unstructured plain text sent by patients to their healthcare team. Responses and messages sent from the clinician or health system to the patient were excluded from the analysis. Portal messages were manually labeled by experienced sub-specialty (cardiology) clinicians into three levels of priority: non-urgent, medium and urgent. Non-urgent labels include notes of appreciation (e.g., thank you). The Medium urgency class contains messages that could be reasonably responded to in 1-3 days. Urgent messages are those requiring an immediate phone call to the patient by the clinician. Conditions suggesting acute myocardial infarction, exacerbation of heart failure respiratory distress or possible stroke were labeled as urgent and would be inappropriate for an asynchronous patient portal.
As summarized in Table 1, the data set is imbalanced. Specifically, the total number of messages, 1,756, includes 631 non-urgent, 955 medium and 170 urgent messages. Urgent messages are scarce (∼ 10%), thus the culprit of the imbalance issue. Table 1 also includes a typical example message for each class. For example, an urgent message could be that a patient reports chest pain. In our experiments, the dataset is split into 80% training set (∼ 1.4K) and 20% test set (∼ 0.35K).

Baselines
To evaluate the performance of LESA-BERT, we compare it to strong baseline classifiers of three kinds: (i) traditional machine learning methods: SVM, (ii) shallow neural networks: text CNN and Bi-LSTM with attention layer, and (iii) pre-trained deep learning models: BERT and BioBERT.
SVM: We utilize an 2 regularized SVM classifier (Evgeniou et al., 2000), which optimizes the hinge loss function with a 2 penalty. SVMs are very effective for high-dimensional data (Ghaddar and Naoum-Sawaya, 2018), especially when the number of dimensions is greater than that of samples. SVMs have been widely used for text classification (Tong and Koller, 2001;Dadgar et al., 2016). In our experiments, we implement linear SVM classifiers with the Python module scikit-learn (Pedregosa et al., 2011), and train it via stochastic gradient descent (SGD) with 2 = 6 × 10 −4 penalty coefficient. The learning rate is set by Bottou (2010) and the weights are randomly initialized.
Bi-LSTM with Attention: LSTM, short for long short term memory, is one kind of neural network suited for sequential data. In NLP, bidirectional (forward and backward) LSTMs are usually used as feature extractors for sequences of tokens. On top of the Bi-LSTM layer, the attention layer is introduced to capture important words that drive the decisions of the document classification (Yang et al., 2016). The attention weights are further employed to compute the weighted sum of output vectors of Bi-LSTM as the hidden representation for each message. Then, this representation is fed into a fully connected layer that produces the logits for the three labels. The hidden dimension of Bi-LSTM is set to 60, and the max length of each message is set to 256 tokens. We implement Bi-LSTM with attention model in Pytorch (Paszke et al., 2019), and train it using Adam (Kingma and Ba, 2015) with batch size of 8. The learning rate is 0.01 and the weight matrices are initialized with method in He et al. (2015).
Text CNN: Details of text CNN model can be found in Kim (2014). We set the following parameters for our CNN model: the convolution kernel sizes: {1, 2, 3, 4, 5}, the filter numbers: {200, 300, 500, 500, 200}, dropout probability: 0.5, and maximum number of tokens in each message is 256. Text CNN also comprises maximum pooling layers, as well as the rectified linear unit (RELU) used as the non-linear activation function. The parameters of the network are learned in mini-batches of size 8 using Adam (Kingma and Ba, 2015). The learning rate is 0.01 and the weight matrices are initialized with method in He et al. (2015).
BERT/BioBERT: For BERT and BioBERT, the base uncased, i.e., all words are treated as lowercase, model is employed in this research. BERT base comprises 12 layers, 768 hidden units, 12 self-attention heads, and 110M parameters. These model are trained with warm-up Adam (Gotmare et al., 2019) setting a batch size of 8 and a learning rate of 3 × 10 −5 . The maximum sequence length is set at 256 and other parameters remain the same as the default BERT base configuration. The weight matrices in BERT and BioBERT are set to BERT base uncased (Devlin et al., 2018) and BioBERT v1.1 (Lee et al., 2020), respectively.

Proposed Models
We briefly cover the configuration of our models, which are implemented with Pytorch in a Python 3 environment 1 .
LESA-BERT: LESA-BERT has the same configuration of BioBERT except for its label embeddings, where embeddings of urgent and medium labels are initialized by the average embeddings of their keywords (see Table 2) while the embeddings of non-urgent labels is randomly initialized.
Distil-LESA-BERT: We distill the fine-tuned LESA-BERT model to two variants: one with 6 encoder layers (Distil-LESA-BERT-6), and the another one with 3 (Distil-LESA-BERT-3). The learning rate for these two variants is 2 × 10 −5 and they are initialized from the first 6 or 3 layers of the pre-trained BioBert model accordingly.

Feature Choices
We follow the literature to choose features of each message for the classification task (Sulieman et al., 2017). For traditional classifiers like SVM, TF-IDF features are commonly used. For shallow neural networks like TextCNN and Bi-LSTM with attention, word embeddings such as Word2vec or GloVe are utilized as input features. For BERT-based deep learning models, we just tokenize each message then the model initializes them from pre-training embeddings (Devlin et al., 2018).
For the SVM classifier, we use TF-IDF vectors (Salton and Buckley, 1988) of each message as features. TF represents the frequency of a certain token in a message and IDF represents the inverse of the number of messages in which this token has appeared. The product of TF and IDF for a token is a score that represents the importance of the token in a message. TF-IDF features have been shown to be very effective in text classification (Zhang et al., 2008).
For shallow neural network classifiers like CNN or LSTM models, instead of using TF-IDF, we adopt pre-trained GloVe word vectors (Pennington et al., 2014) as features for each message. We use the 100 dimensional GloVe word vectors pre-trained on 6 billion tokens from Wikipedia and Gigaword data sets. For the BERT-based models including BERT, BioBERT and LESA-BERT, we use the wordpiece tokenizer (Wu et al., 2016) and keep the models' pre-trained token embeddings as initialization, so that the input embeddings are compatible with the models' pre-trained parameters. We utilize the BERT base uncased model with 12 layer encoders and 12 attention heads and 768 dimensional embeddings, which was pre-trained on BookCorpus (Zhu et al., 2015) and English Wikipedia. BioBERT shares the same model configuration as BERT base, but was further pre-trained on biomedical corpora (PubMed abstracts and PMC fulltext articles). Parameters in LESA-BERT except the label embeddings are initialized from the pretrained BioBert model.

Evaluation Metrics
For classification tasks, commonly utilized metric criteria are precision, recall, F1 score and Area Under the Curve (AUC) score. Oftentimes a classifier has a good precision with a poor recall, thus F1 score provides a good balance of precision and recall. AUC is typically employed for binary classifiers. Since the message urgency dataset comprises multiple (three) classes and is imbalanced, we find that macro average F1 score as the most suitable evaluation criteria (Parambath et al., 2014). The F1 score is given by the harmonic mean of precision and recall. Accuracy is a poor metric in this situation because it encourages the model to focus on the majority class. Similarly, micro-level metrics like micro F1 score are not preferred because they are weighted by the number of cases in each label, which makes them favor the majority classes.
We also evaluate the advantages of the distillation process. We distill the fine-tuned LESA-BERT model to 6-layered and 3-layered versions, and evaluate the number of parameters and the inference time required for a full pass on the test set on a CPU with batch size of 1, which are commonly used evaluation criteria for distillation (Sanh et al., 2019). Table 3 presents quantitative results for various classifiers on the test set data. Mean and standard errors are computed based on 5 different random seeds for the SGD/Adam training algorithm. The F1 scores for shallow neural networks, i.e., TextCNN and Bi-LSTM attention, are slightly better than SVM, which might be owed to the benefits of semantic features captured by GloVe word vectors. BERT-based deep neural networks such as BERT and BioBERT only marginally outperforms Bi-LSTM attention and TextCNN, which is likely caused by the small size (∼ 1.4K) of the training data. The F1 score of BioBERT is 0.3% higher than that of BERT, which means that pre-training on large biomedical corpora transfers extra beneficial information to BERT model. In terms of macro average F1 score, our LESA-BERT and its distilled variants achieve higher F1 scores compared with other classifiers. The performance gains can be attributed to the addition of the label embeddings.

Quantitative Results
The model Distil-LESA-BERT-6 provides higher F1 score (0.807) than that of the full 12-layered LESA-BERT. Possibly, this implies that the full model overfits the training set, while the distilled model with 6-layers reduces the overfitting. Distil-LESA-BERT-3 model, which comprises only 3 layers provides F1 score of 0.78, slightly lower than the full LESA-BERT model. Table 4 presents the number of parameters in LESA-BERT and its distilled variants, as well as inference times for a full pass over the test set with batch size of 1. As the number of layers in LESA-BERT decreases, the number of parameters and inference time reduces dramatically. Distilling the original 12-layered LESA-BERT to 6 layers reduces the number of parameters by 40% and accelerates the inference by 2.7 times. Distil-LESA-BERT-3 has fewer parameters and faster inference speed than Distil-LESA-BERT-6. Taking into account their F1 scores, Distil-LESA-BERT-6 outperforms the LESA-BERT, and Distil-LESA-BERT-3 retains 98.9% of LESA-BERT performance. Distil-LESA-BERT-6 provides the best F1 score compared to all other methods. However, in the case that inference time is preferred, Distil-LESA-BERT-3 will be the top choice.

Qualitative Results
In Figure 4, we visualize how the obtained representation by LESA-BERT attend to different tokens in three example messages, one for each class. Tokens are shaded in terms of their importance for classification. Dark-colored tokens are more prominently weighted (attended) when constructing the message embedding. As seen in Figure 4, attention scores can identify relevant keywords in all three examples. We choose these three examples because they are representative in their respective categories.  messages of the other two classes. For the urgent and medium classes, their data points have significant overlap, meaning that urgent and medium messages, not surprising, have similar characteristics (features). This may be in part caused by the subjectiveness of the labeling process as the definitions of urgent and medium messages are not exactly clear.

Discussion
We have proposed a framework for building classifiers on small and imbalanced datasets, which is a commons scenario in healthcare. Our method builds upon a large deep learning model, BioBERT, which was pre-trained on huge general-domain and biomedical corpora. We developed LESA-BERT, which has a novel self-attention architecture that can incorporate label embeddings to boost the model's capacity for attention. We fine-tuned LESA-BERT on the target dataset by initializing its parameters from BioBERT. Subsequently, we distilled the fine-tuned 12-layered LESA-BERT to 6-layered or 3-layered variants, and found that the 6-layered Distil-LESA-BERT outperforms the 12-layered LESA-BERT. Therefore, knowledge distillation is not only a tool for model compression, but can also be used to reduce overfitting. We demonstrated the application of our framework on a real healthcare dataset -message urgency-and built a message triage classification model. Our methods outperformed baseline classifiers by a significant margin. Our technical solution can be easily applied to other clinical datasets.
Clinical Implications We built an automatic message triage model that can predict the priority of patient portal messages based on their content. This version of message triage could be further improved with a larger dataset to train our classifiers. The message triage system is of real impact to healthcare providers and potentially save valuable working hours by freeing them from reading patient portal messages in chronological order. Promptly accessing and realizing urgent messages can impact both patient safety as well as clinician workflow. After the message triage classifier produces the priority of messages, healthcare providers can take further actions based on their existing workflow. Further, they may want to utilize automatic responses or templates to answer non-urgent messages. Medium or time-sensitive messages may either be directed to emergency rooms, urgent care or telephone encounters.
Limitations Although our methods outperform a wide selection of baseline models, it can still produce some mistakes from a human perspective. For instance, Figure 5 shows that two non-urgent messages (purple dots) fall in the region of medium and urgent labels. In fact, we have investigated these very few messages that are close to medium or urgent messages, and one example reads "No chest pain at all". This message has very strong keywords, chest pain, which usually appears in urgent messages. Therefore, the LESA-BERT model maps this message to urgent class by ignoring the negation word "No". Many researchers have found that deep learning models are powerful but still easy to fool (Moosavi-Dezfooli et al., 2016;Heaven, 2019). As future work we plan on incorporating negation detection (Qian et al., 2016) into our classification framework.
Our current message triage system has a few limitations from the clinical perspective. Firstly, our message classifier only has three levels of urgency, which could be further refined into more granular categories. After passing patient portal messages through our urgency classifier, urgent messages will be treated with priority, but non-urgent and medium ones will probably need further automatic classifiers, such as message intention detector, etc. Secondly, our current work is limited to only utilizing the content of portal messages. One possible follow-up research could be to incorporate patient's demographic information or even comorbidities. For example, age can be quite important when deciding the priority of messages. Messages from patients aged over 65 should receive relatively more immediate attention than younger patients. Comorbidities could also be taken into account as patients with heart attacks, for example, usually get more intense care than patients with, for example, common respiratory infections.
Yang Zhao, Zoie Shui-Yee Wong, and Kwok Leung Tsui. A framework of rebalancing imbalanced healthcare data for rare events classification: a case of look-alike sound-alike mix-up incident detection. Journal of healthcare engineering, 2018, 2018.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.


s2fieldsofstudy:
Computer Science
Engineering

text:
A Fast Multi-tasking Solution: NMF-Theoretic Co-clustering for Gear Fault Diagnosis under Variable Working Conditions

Most gear fault diagnosis (GFD) approaches suffer from inefficiency when facing with multiple varying working conditions at the same time. In this paper, a non-negative matrix factorization (NMF)-theoretic co-clustering strategy is proposed specially to classify more than one task at the same time using the high dimension matrix, aiming to offer a fast multi-tasking solution. The short-time Fourier transform (STFT) is first used to obtain the time-frequency features from the gear vibration signal. Then, the optimal clustering numbers are estimated using the Bayesian information criterion (BIC) theory, which possesses the simultaneous assessment capability, compared with traditional validity indexes. Subsequently, the classical/modified NMF-based co-clustering methods are carried out to obtain the classification results in both row and column tasks. Finally, the parameters involved in BIC and NMF algorithms are determined using the gradient ascent (GA) strategy in order to achieve reliable diagnostic results. The Spectra Quest’s Drivetrain Dynamics Simulator gear data sets were analyzed to verify the effectiveness of the proposed approach.

Introduction
In those large-scale rotating machines, wear and tear always comes out in the teeth surface of driving gears if the pressure is not even or some extra impurities are mingled in the lubricating oil. Health monitoring technology of mechanical components has been proved to be effective at discovering early abrasion and reducing the failure rate [1][2][3][4][5]. As one of the major tasks in health monitoring, gear fault diagnosis (GFD) aims to assess the current gear state based on the obtained measurement data, then to inform users to take proper actions [6]. A GFD procedure generally consists of three main processes: (1) Data acquisition: data are collected from sensors to monitor the health status of gears; (2) Feature extraction: some feature extraction algorithms, such as wavelet transform (WT) [7] and least squares support vector machine (LSSVM) [8], are carried out based on the prior knowledge to provide recognizable features; (3) Fault recognition: classifiers are built to obtain gear faults with the analysis of the extracted features.
Clustering technology, one of unsupervised fault recognition approaches, has experienced long term development from partition based clustering to graph theory based clustering as listed in Table 1. Most of these algorithms were applied to fault diagnosis of rotating machinery. For instance, Yuwono et al. [9] combined particle clustering with a Hidden Markov Model (HMM) for bearing fault diagnosis; Pacheco et al. [10] classified gear fault severities using rough set theory. These researches have provided effective clustering applications related to machine fault diagnosis. However, they have a non-negligible limitation: each feature vector is treated as independent and uncorrelated unit in these clustering methods. In fact, strong correlation exists between machine working conditions and targeted diagnostic tasks. There is no doubt that clustering analysis in one dimension will lose significant information hidden in another dimension. To overcome this limitation, a co-clustering strategy for variable working conditions GFD is proposed in this paper. Coclustering was firstly utilized in the biology and medical domains since this concept was mentioned by Mirkin [11] in 1997. The joint clustering of genes shapes and locations promoted the discovery of genetic structure sequence [12]. Subsequently, co-clustering has been expanded to other fields such as text analysis [13] and search engines [14], etc.
In this study, co-clustering applications for gear fault diagnosis have been developed. Compared with previous GFD based on joint clustering methods, two highlights in this paper can be obtained: (1) when one varying working condition (such as rotating speed or load) and one diagnostic task (such as fault severity) are jointed in the same matrix, their correlations are extracted, and the classification accuracy of the latter can adjust with the range of the former; (2) when two diagnostic tasks (such as fault severity and fault type) are jointed in the same matrix, they can be classified at the same time, which improves the diagnosis speed compared with independent GFD strategy and offer a fast multi-tasking solution.
The remainder of this paper is organized as follows. Section 2 presents a brief summary about the applicability of co-clustering. It also describes the principle and basic framework of co-clustering to solve the GFD problem. Section 3 addresses the preparatory work of GFD, especially a short-time Fourier transform (STFT)-based feature extraction method. In Section 4, the co-clustering numbers are estimated based on the Bayesian information criterion (BIC). Then in Section 5, the traditional and modified NMF-theoretic co-clustering process is discussed in detail. To assess these algorithms, the gradient ascend algorithm is also implemented for parameters regulation in Section 6. Section 7 concentrates on the varying working condition GFD experiments using the Drivetrain Dynamics Simulator (DDS) system, which especially shows the superiority of co-clustering compared with classical clustering strategy such as X-means algorithm. Conclusions are drawn in Section 8 with discussions on the future GFD application based on joint clustering.

Co-clustering Framework of GFD
Traditional clustering can be defined as: dataset X exists in a limited data space, which can be represented with a n × d matrix, is composed of n elements: The purpose of general clustering is to segment dataset X into p categories: Different from general clustering, co-clustering could be defined as: dataset X exists in a limited data space, which can be represented with a m × n × d matrix, is composed of m × n elements: The purpose of co-clustering is to segment dataset X into p and q categories in horizontal and vertical axis, respectively: C h k (k = 1, 2, · · · , p) and C v k (k = 1, 2, · · · , q).
where i ∈ {1, 2, . . . , n} , j ∈ {1, 2, . . . , m}. Currently, one challenge for fault diagnosis of gears is that they are often operated under the varying working conditions. To explain the influence of varying working conditions on the diagnosis performance, a typical example is given in Figure 1 to compare the classification results between classical and co-clustering strategies. In the extracted feature matrix, different colors represent different feature value. C1-C5 represent different fault categories; D1-D5 represent different working conditions. R1-R9 represent the 1st clustering result; r1-r5 represent the 2nd clustering result. It can be seen that the classical clustering results may be distorted due to the fact that different working conditions have direct influence on the extracted features, and this makes difference between final clustering categories (R1-R9) and real fault categories (C1-C5). On the contrast, the co-clustering can reflect the actual diagnosis results (R1-R5), which shows the robustness of co-clustering in GFD under interference environment.
Notice that, co-clustering is not limited in the two dimensions. Theoretically, it can be generalized to higher dimension (n ≥ 3), thus giving an idea to solve
Graph theory-based ACODF [23], etc. the more complex problem such as the gear fault diagnosis problem under multiple working conditions. Figure 2 explains the mechanism of co-clustering GFD models under two working condition factors, such as varying rotating speed and varying load. In Model (a), the 2D classification matrix is structured in two dimensions: row for rotating speed and column for GFD task. In Model (b), the 3D classification matrix is structured in three dimensions: length for rotating speed, width for load and height for GFD task. Structured high dimension matrix is classified in each scale at the same time, which offers an idea for GFD under more than one working condition. According to the description above, a co-clustering framework is constructed for gear fault diagnosis, which is shown in Figure 3. The main process can be divided into four sub-frames.
• Feature extraction sub-frame: as the input of this model, the gear vibration signals are collected using several tri-axial accelerometers installed in monitored mechanical equipment, which may be operated in varying working condition environment.  Then, the feature vectors {F } are obtained using the short-time Fourier transform (STFT) approach, aiming to gain differentiable time-frequency features for effective co-classifier performance; • Clustering number estimation sub-frame: the Bayesian information criterion (BIC) strategy is adopted to characterize the distribution character of all feature vectors {F } and then estimate their co-clustering numbers k & l in row and column, respectively; • Co-clustering sub-frame: given co-clustering numbers, the conventional as well as modified NMFbased co-clustering classifiers are put into practice to build the varying working condition GFD models and get the classification results in various tasks; • Parameter regulation sub-frame: aiming to those adjustable parameters involved in BIC and the NMF algorithm, such as the weight factor and the transitional dimension d , the gradient ascent (GA) algorithm is implemented to find the optimal values, which reaches a reliable diagnostic accuracy.
More details of these four sub-frames will be given from Section 3 to Section 6.

STFT-based Feature Extraction
The performance of feature extraction algorithms depends critically on the characteristic of the input gear fault signals as well as the running environment of equipment. Because of the non-stationary property of the vibration signal, general algorithms, such as Fourier transform (FT), are not very useful in this regard [24]. As a kind of time-frequency analysis method, the main idea of short-time Fourier transform (STFT) is to execute Fourier transform in short sequential signals, cut by a sliding temporal window function γ (t) , such as a Hanning or Hamming window [25]. When analyzing the non-stationary vibration signal s(t), (t = 1, 2, . . . , t 0 ) , we supposed that it can be approximated as smooth among the window function γ (t) . Therefore, the STFT of s(t) is calculated by timefrequency units STFT (t, w) and is given by where τ means the position of window function γ (s) ; w represents the frequency parameter of STFT. Since the short time Fourier spectrum reflects a distribution of energies among all frequencies and temporal intervals, we have to seek those 'meaningful' values from the STFT graph. Without any known knowledge of the running equipment, an effective method is to find each maximum in the partitions of the STFT spectrogram to represent the feature of sub-window. Figure 5 gives an example of extracted features in the chipped gear in Figure 4, including 10 × 10, 40 × 40 and 80 × 80 features. Generally, the row & column divided dimensions depends on the non-stationary degree of signal in time and frequency domain, respectively. This figure indicates that more details are emerged in high dimension features compared with low dimension features. However, the increase of dimension will bring the expansion of computational load as well as time consuming. Therefore, a middle dimension is suitable if the definition and dimension are both acceptable, such as the 40 × 40 features in Figure 5 (2020) 33:16 where st ij represents coefficient at the location ( i, j ) int the sub-matrix of the STFT graph; The dimension N will be updated in the real GFD experiments in Section 7. Finally, the feature vectors from STFT are obtained and are given by

BIC-based Clustering Number Estimation
In most practical gear fault diagnosis applications, it is difficult to know the estimated number of clusters in advance, but the input clustering numbers always have direct influence in final clustering results. An optimal numbers estimation strategy based on the cluster validity indexes is found in historical literatures, including: the Calinski-Harabasz index [26], the Davies-Bouldin index [27], the weighted inter-intra index [28] and the in-group proportion index [29]. However, two drawbacks exist in the cluster validity indexes-based estimation strategy: (1) The estimation precision depends on selected clustering algorithm, dataset as well as the validity index. For instance, both using the IGP index, the affinity propagation (AP) method has more credible optimal clustering-number compared with the k-means strategy on account of the randomness of the latter [30]. (2) It is hard to apply the cluster validity indexesbased estimation strategy to 2D or higher dimension clustering like co-clustering because these validity indexes are mainly calculated according to the distance algorithms between different classes or within the same classes.
To find optimal numbers of co-clustering, an estimation algorithm based on Bayesian information criterion (BIC) is proposed. BIC is a statistical method which represents the descriptive power of a model to dataset [31], including: (1) the posterior likelihood of data estimation L ; (2) The model complexity |Θ| . The computational formula of BIC is given by where means the weight factor; N is the totality of samples. In clustering, the posterior likelihood of data estimation L is represented using the ratio of the mutual information entropy between after-clustering I(S * ; F * ) and before-clustering I(S; F ) . In Eq. (6), the entire meaning of L is that a good clustering must maintain original information entropy as possible as it can, But in 2D co-clustering, the BIC model requires to take row and column clustering into consideration at the same time. So we redefine the parameters as follows:

Direction
Sample length Sample size Clustering number According to these definitions, the BIC model complexity in co-clustering can be re-expressed as Substituting Eq. (7) and Eq. (8) into Eq. (6), we get Eq. (9) as follows: Further, this Bayesian information criterion can be extended to 3D field and is given by where p is the clustering number of the 3rd classification; q is the size of the 3rd dimension. Based on the description of theory above, the details of BIC algorithm is given as: Input: a. For single variable working condition GFD, the sample matrix C is structured by where x 1 · · · x 40 T represents the STFT values in continuous approximate one minute, which depends on the change speed of working condition; i ∈ {1, 2, . . . , n} represents the random samples collected from different fault type and n is the number of sample. b. For double variable working condition GFD, the sample matrix C is structured by  Process: a. Select an matrix construction method to build C matrix according to the number of running environments; b. Initialize the row clustering number k = 1 , the column clustering number l = 1 , the weight factor = 0.5; c. For k from 1 to m , l from 1 to n: where p s, f means the joint probability distribution between row and column; p(s) means the probability distribution in row vector; p f means the probability distribution in column vector. d. Search the max value of BIC during the whole k and l domain: Output: a. For single variable working condition GFD, it outputs the number of final clustering numbers k and l; b. For double variable working condition GFD, it outputs the number of final clustering numbers k , l , and p.

NMF Theory
Non-negative factorization (NMF) is a kind of efficient data compression strategy, aiming to describe the highdimensional data set using few base vectors with the help of the non-negative theory [32]. Different from the global (12) characteristics of vector quantization (VQ) and principle component analysis (PCA) theory, NMF offers a good description about the local features, so specializing in searching the small scale information. Non-negative factorization has been investigated for feature extraction and recognition of rolling element bearing fault [33]. However, the unique advantage of co-clustering has not been explored for 2D or higher dimensions application. The basic NMF problem is stated as the following equation: where C n×m is a n × m non-negative matrix; the basis matrix W and the gains matrix H are factorized from C n×m ; d < (n × m)/(n + m) represents the reduced rank. Therefore, C n×m can be linearly estimated by the sub-vectors W n×d and H d×m . In order to obtain matrix W and H , a large number of cost functions are defined to quantify the degree of approximation. In our strategy, the Euclidean distance is chosen as the cost function.
The purpose of NMF is to find the W and H , which possesses the smallest cost function: min where ⊗ is the element-wise multiplication, is the element-wise division; r represents the iteration.

Classical NMF-based Co-clustering
Recently, the clustering application based on NMF has attracted much attention. Particularly, KIM, etc., explored the effective combination between cluster and NMF [34]. This paper extends its application from single cluster to co-clustering, aiming to solve the varying work condition or multi-tasks problem. Rely on the computational W ∈ R n×r and H ∈ R r×m above, two objective functions J k & J l are defined as follows: where C r = [c 1 , c 2 , . . . , c k ] r ∈ R m×k and C 2 = [c 1 , c 2 , . . . , c l ] c ∈ R n×l represent the centroid matrix in row and column, respectively; The element c j , j ∈ [1, 2, . . . , k] of C r means the cluster centroid of the jth cluster in Task I (16)  (2020) 33:16 and the element c j , j ∈ [1, 2, . . . , l] of C c means the cluster centroid of the jth cluster in Task II; B r & B c denote clustering assignment in Task I and Task II respectively. In Task I, B r ij = 1 means the ith sample belongs to the jth cluster, otherwise B r ij = 0 , and so is Task II. The purpose of co-clustering is to find sparse matrix B r ij and B c , which has only one in each row, with others being zero. Taking Task I and Task II as example, we redefine C r and C c as Finally, a second order NMF is applied in J k and J l , aiming to factorize W to W (M r ) T and N r , to factorize H T to H T (M c ) T and N c . Therefore, the B r and B c matrix is obtained according to the second order NMF result. After that, the classifications in row and column are obtained from the B r and B c matrix.

Modified NMF-based Co-clustering
As described in Section 5.1, non-negative matrix C is factorized into two sub-matrices W and H in conventional non-negative factorization. Although the physical meanings of W and H are clear: they represent the (20) decomposition values in row and column respectively and promote the classification effect of co-clustering, the relation between two directions is still ill-defined. Hence, the typical NMF is improved and is given by where L k×l is named as 'the relation matrix', the value L ij represents the link between the ith cluster in Task I and the jth cluster in Task II; k is the clustering number in the row vector and l is the clustering number in the column vector. In modified NMF, the cost function and the update functions can be re-written as where ⊗ is the element-wise multiplication, is the element-wise division; r represents the iteration; W + is the generalized inverse of W : WW + W = W ; H + is the generalized inverse of H : HH + H = H. By introducing the matrix L , the W and H are not required to be orthogonal in modified NMF strategy. Therefore, it expands the optional range of W and H and improves the factorization performance, which will be proved in Section 7.

GA-based Parameter Regulator
Three parameters need to be designed in the co-clustering-based GFD strategy, including: (1) the feature-dimension N in STFT; (2) the weight factor in BIC algorithm; (3) the transitional dimension d in traditional NMFbased co-clustering, which are listed in Table 2.
Among these three parameters, the feature-dimension N in STFT can be decided from the GFD experiments. and d will be adjusted using the gradient ascent (GA) regulatory mechanism [35,36], whose fundamental is shown in Figure 6. The main idea of gradient ascent algorithm is to follow the fastest changing direction to find the maximum of diagnostic accuracy. In Figure 6, four different initial points of ( , d) are listed as example, where the approximation curves ① and ② reach the global optimum while ③ and ④ are limited in the local optimum.
Meanwhile, three concepts are included in the gradient ascent model: where y r i and y c i represents the label of the ith sample; zer(·) means the zero sign function. Notice that, the validity function can only be obtained in those training samples, whose classification labels are known.
The optimal and d/min(m, n) is gained according to the training samples and is used in others, called testing samples. It should be noted that in real GFD application, when performing GA algorithm: (1) If the step length is too large, the optimal parameter result might be skipped. But if the step length is too small, the iteration speed will be slow and cause too large computational load. (2) It is easy for the GA algorithm to be deep in the local optimum rather than the global optimum, which relies on the initial location of , d/min(m, n) . Therefore, it is necessary to take these two factors into consideration to balance the computational accuracy and the time consumption.

DDS Experimental System
The Spectra Quest's Drivetrain Dynamics Simulator (DDS) was used in this study for experimental verification, as shown in Figure 7. This system is composed of six units including: (1) speed regulator; (2) the driving motor; (3) the planetary gearbox; (4) the reduction gearbox; (5) brake device; (6) brake regulator. The faults occur in those gears in planetary & reduction gearboxes, under varying rotating speed and load conditions, which are adjusted using the speed regulator and the brake regulator, respectively. Four types of gear faults are studied: (1) root cracks; (2) missing teeth; (3) chipped teeth; (4) surface wear. The purpose of GFD is to classify these faults through 7 vibration sensors (3-axis for planetary gearbox; 3-axis for reduction gearbox; 1-axis for driving motor). In addition, in order to put the co-clustering methods into effect, we define different levels in four task sets listed in Table 3, including: (1) the fault type task (C1-C5); (2) the fault severity task (D1-D4); 3) the speed regulator task (E1-E5); (4) the load regulator task (F1-F5). Specially, the rotating speed curve and the load curve in Figure 8 was also conducted.

Experimental Setup
The varying rotating speed and load are designed using the regulator curves in Figure 8 for the experiment. In order to enlarge the data analytical ability of algorithms, 10 repeat collections were implemented to increase the sample points to 10 times (5120 × 50× 10) in each group, which are segmented by the 2.5 s sub-signals. Therefore, the sample number for fault type recognition can be gained in row clustering task (fault type) and column clustering tasks (rotating speed and load), and are listed in Table 4A and 4B.
In the fault type recognition experiments, several tests are compared using the models as follows: (1) X-means clustering; (2) Guassian Mixture Model (GMM)

Experimental Results and Discussion
Experiments are carried out, and the column task in Table 4 is considered to improve the effectiveness of row task in co-clustering method. In order to assess the quality of models, the Pr (C i ) and Re (C i ) indexes of clustering results are calculated and given by where the Pr (C i ) index reflects the ability that a model identifies correct samples while the Re (C i ) index reflects the ability that a model finds all correct samples. At first, we observed the fault type recognition results using the NMF-based co-clustering strategy under varying rotating speed and load environment, which are listed in Table 5. This table indicates that the diagnostic accuracy of co-clustering strategy in various load conditions (97.8%) is better than that in various rotating speeds (32) Pr Corrected classfied C i samples All classfied C i samples , Figure 7 The Drivetrain Dynamics Simulator (DDS) system: a the whole physical map of DDS system; b the internal structure image of DDS system; c transmission ratio of reduction gearbox Torque regulator (N·m) < 1.83 1.83-5.49 5.49-9.14 9.14-12.80 > 12.80 (2020) 33:16 (96.3%). This can be explained by two possible reasons: (1) The classify boundary of the latter is more clear than the former since that the rotating speed presents a gradual change characteristic from 0 Hz to 40 Hz while the brake load jumps from 0 to 14.63 N·m; (2) Changing the rotating speed has a stronger interference effect than just changing the load to collected vibration signal, which have been verified in our previous study. Secondly, to illustrate clustering performance in the varying rotating speed model, the NMF-based co-clustering results are listed in Table 6 (k = 6; l = 6) and Figure 9, where 200 samples are tested in each category. Some details can be seen here: (1) the misclassification cases always appear between 'Health' and 'Surface' or between 'Root' , 'Missing' and 'Chipped' because the time domain features of the former are more similar but the frequency domain features of the latter are alike. For example, the 'Chipped' type is easy to be classified as the 'Missing' type if the crack of chipped tooth is large enough; (2) The 6th category (R2 and R3) occurs in the C2 type when the number of clustering is set as 6, which means the discrete ability and the inconsistency exists inside the 'Root' samples. Interestingly, although the differences exist in local precision and recall index of different categories, the total precision and recall are very nearly the same in these two tables.   Especially, we reclassify the fault types of testing samples in Table 4 using the modified NMF-based co-clustering strategy. Here the normalized relation matrix L k×l (k = l = 5) can be gained and given by Notice that, the element l ij represents the relevance between the ith category in row task and the jth category in column task. In the L k×l matrix, the l ij rises up to approximate 7 in 'Missing' & 'Chipped' fault type with the increase of rotating speed. The effect of speed regulator on different fault types presents the following rank: Miss-ing≈Chipped>Root>Health≈Surface. However, there is a less link between load and fault types, seeing from the L k×l matrix under changing load environment. In order to further verify the necessity for modified NMF, Table 7 gives the evaluation indexes of this approach. It shows that an improved precision occurs in the varying rotating speed dataset but makes no difference in the varying load dataset. That is because the L k×l matrix in modified strategy cuts off the link between W and H in the former, which promotes the separation ability between row and column task. As can be seen in Table 7, the middle dimension of W and H is not limited in a single value, like traditional NMF method does, thus improving both the flexibility of selected dimension and the GFD precision (97.0% and 97.8%). Also, it can be known that the recognition performance of various loads (100%) is superior to rotating speed (95.3%) on account of the continuity of speed regulator.

Table 4 Sample number for (A) rotating speed fault recognition and (B) load fault recognition
Finally, concentrating on the varying working condition GFD, the traditional clustering strategies, including the K-means and the GMM methods, and co-clustering approaches were compared using two selected indexes: the precision and the time consumption. The performance comparison results are listed in Table 8. According to this table, although the time complexity of single algorithm of 1D clustering is smaller than joint strategy (4.923 s < 7.991 s), the accumulation of computational load for two tasks is larger than co-clustering proposed (4.923 s + 4.856 s > 7.991 s). Meanwhile, the co-clustering have an apparent precision increase in varying working condition GFD, about 12.51% in varying rotating speed and 7.00% in varying load. Therefore, this experiment proves the superiority of co-clustering in gear fault diagnosis under variable working conditions.

STFT Dimension Adjustment Experiments
During the STFT dimension adjustment experiments, we adjusted the feature dimension N from 10 to 80 one by one, and then the diagnostic precisions of Task I as well as the time consumptions of co-clustering model were observed and were drawn in Figure 10. It can be seen that the diagnostic accuracy increases from 78.76% to 97.54% when the feature dimension N increases from 10 to 80, meanwhile, the computational load indicates an exponential increase from approximately 16.78 s to 48.45 s. It can be seen from Figure 10 that N = 42 is considered as an appropriate dimension, in which the diagnosis accuracy is satisfactory enough (97.02%), while the time consumption keeps at a low level (25.12 s). Although the precision will continue to improve up to 97.54% if we  continuously increase the feature dimension, but computational cost will also increase quickly.

BIC Algorithm Experiments
Based on the dataset from Drivetrain Dynamics Simulator (DDS) system, we tried to adjust the number of cluster in row and column, respectively. Generally, the search range of clustering number is from 2 to 10: 2 ≤ k ≤ 10; 2 ≤ l ≤ 10. Figure 11 illustrates an example of the BIC results in varying load GFD experiments when the clustering number changes from 2 to 10. It can be seen that the peak of BIC value (− 8124) exists at the point (5,5), which means that the optimal co-clustering results happens when the row and the column clustering numbers both equal to five. According to the real dataset and the standard labels in Table 3, the BIC-based estimation algorithm satisfies the requirements of practical gear fault diagnosis applications. For further study on the BIC estimation algorithm, the BIC method was compared with a kind of traditional selfadapting classification estimation algorithm: X-means, which is an improved strategy of K-means. The estimations of clustering number in BIC as well as X-means algorithm are listed in Table 9. On one hand, for the BIC algorithm, with the increase of clustering number k, the diagnostic precision presents an increasing trend from 60.0% to 96.9% first, then it begins to slightly decrease when the clustering number k is larger than 5. On the other hand, by comparing the estimation results between X-means and BIC algorithm, we find the estimation of clustering number using the X-means is much larger than the normal (k = 12) and also its diagnostic precision is only 84.1%, which proves the superiority of BIC-based clustering number estimation.

GA Parameter Regulator Experiments
As shown in Figure 4, the gradient ascent algorithm was used in the varying working condition datasets to obtain the weight factor and the transitional dimension d . The regulator results of the gradient ascent algorithm are listed in Table 10. Here the final ( , d ), the number of iterations as well as the precision of row task are observed using 4 initial ( , d ) values: (0.5, 500), (0.5, 200), (0.7, 500), (0.7, 200). After gradient ascent, it could be seen that four initial values all reached the global optimum in varying load experiments. But in varying rotating speed experiments, only two points reach the global optimum, which achieves 99.3% diagnostic accuracy. In addition, it can be found that the less iterations occur when the distance between initial and final ( , d ) point is short. Therefore, these tests prove that the performance of GA algorithm depends on the selected initial point to a great extent.

Conclusions
A NMF-theoretic co-clustering strategy is presented in this paper to offer a fast multi-tasking solution to solve the gear fault diagnosis problem under variable working conditions. Here the time-frequency features are extracted from the STFT spectrogram, and are utilized to structure the 2D matrix for joint clustering. Experiments indicate that 97.02% diagnostic precision can be achieved when the STFT dimension is set as 42. Meanwhile, seeing from the results of the BIC-based optimal clustering number estimation, they are close to the practical categories, no matter in varying rotating speed or varying load dataset. After NMF, row and column clustering task can be identified at the same time, with approximately 10% improved accuracy and less time cost compared with those single task clustering algorithms, such as X-means and GMM algorithm. There is an internal connection in most of gear failure signals. The proposed co-clustering strategy has better performance than independent clustering strategy because the modified NMF helps to provide a relation matrix, which shows a strong correlation between different rotating speeds and fault types. Therefore, the NMFbased co-clustering has a good potential to apply in the gear fault diagnosis of large-scale rotating machines under varying working conditions. In the future, co-clustering with higher dimension will probably apply in the more complex working conditions or more diagnostic tasks to improve the gear fault diagnosis performance.

Authors' Contributions
RY designed the experiment, FS and CC analyzed the data. All authors read and approved the final manuscript.


extfieldsofstudy:
Biology
Computer Science

text:
An Overview of the Application of Deep Learning in Short Read Sequence Classification

Advances in sequencing technology have led to an ever increasing amount of available short read sequencing data. This has, consequently, exacerbated the need for efficient and precise classification tools that can be used in the analysis of this data. As it stands, recent years have shown that massive leaps in performance can be achieved when it comes to approaches that are based in heuristics, and alongside these improvements there has been an ever increasing interest in applying deep learning techniques to revolutionize this classification task. We attempt to gather up these approaches and to evaluate their performance in a reproducible fashion to get a better perspective on the current state of deep learning based methods when it comes to the classification of short read sequencing data.

I. INTRODUCTION
Individual genome and metagenome sequencing has become increasingly more affordable over the past years [1] [2], which has lead to an explosion in the amount of data that is available for analysis. This has, in turn, spawned a need for more affordable tools to analyze this data. In this study we will look at recently developed taxonomic classification tools and evaluate their performance when working with metagenome sequencing data. A state of the art tool widely used for tasks such as this is Kraken2 [3], which has shown that there is still a lot of room for improvement in currently used methods, with Kraken2 reducing its memory footprint to just 15% and its runtime to just 20% its predecessor Kraken [3].
A number of deep learning and machine learning based tools have been produced in recent years with a goal of alleviating different, resource intensive, aspects of current methodsmemory requirements, classification time, classification precision or disk space requirements [4][5] [6] [7]. A major issue of evaluating the current state of these tools is in the fact that there are no published results showing their relative performance when the same reference data is used and when they're applied to non-synthetic samples. This is an issue that is characteristic to many bioinformatics tools [8]. There is also a known issue with self reported performance metrics, which are the only results available for a lot of toolsthey tend to have bias problems [9].
There have been some attempts to standardize the way tools are benchmarked [8], but thus far no standard protocols nor benchmark data sets have been adopted and most tools still report their results using bespoke approaches. This makes gauging the performance of these tools in real world scenarios quite tricky as generated samples can be less complex than real world ones [16].
Our goal is to evaluate the performance of a set of freely available published deep learning based tools by using the exact same reference data: metagenome data from a set of samples that have been sequenced by MGI DNBSEQ-T7 sequencers [10], which include 2 reference samples with a predefined composition (Zymo samples [11]referred to as C1 and C2) and 3 human gut (fecal) metagenome samples that were produced within the ERDF project 'Optimisation of H.pylori eradication therapy for population-based gastric cancer prevention' (referred to as S1, S97 and S104). The non-control samples were added to the experiment to assess how the tools perform if the reference database does not contain all of the organisms present in the sample and therefore there are no exact matches to the sequences in the sample.

II. METHODS
We selected suitable deep learning based tools for genomic classification by searching the SCOPUS and Arxiv databases, excluding the tools that were designed for tasks other than shotgun sequence classification. This way we picked 4 tools: • MetaVWembedding based tool, leverages Vowpal Wabbit [12] as its backbone [4], • fastDNA -fastText [18] based embedding tool [5], • GeNeta classification tool that uses a convolutional model [6], • DeepMicrobesa classification tool that leverages an attention model [7]. In addition, we elected to use the Kraken2 results as our 'ground truth' for samples that we did not have a theoretical taxonomic distribution for because it is a widely recognized and commonly used tool [3]. We found this necessesary due to the way the experiment was set up: we included real-life samples with an unknown composition, so an estimate of their composition was established using Kraken2 as a reliable stateof-art tool.
The expirements were performed on a workstation running a 12 core procesor (Intel i7-8700k), 64 GB of RAM, a Nvidia GTX 1080 Ti graphics card (GPU) and 500 GB of PCIe storage. We make a point of noting down the storage solution due to Vowpal Wabbit being heavily bound by storage speed [12] which will impact the performance of MetaVW [4].
We evaluated all tools using the same 5 samples (read counts and average read length, which provide insight into data complexity, are given in Table 1) and used the small database from MetaVW [13] (mostly due to practical concernsit contained all of the info needed for all of the tools to train) to create the reference databases for all of the tools (containing 1565 sequences related to 193 species) [4].
The two control samples are ZymoBIOMICS™ Microbial Community Standard samples and contain data from 10 different species [11]. Reference databases were created and models were trained using the default setting wherever possible. The only deviation from this being DeepMicrobes where we used 11-mers instead of 12-mers due to not being able to train the 12-mer model within 11 GB of video RAM. The commands we used have been made available on GitHub [14]. These runs were timed and resource usage was monitored during training. A general outline of the database creation process is presented in Figure  1.
After training the neural network models, they were used to classify the previously described samples. The results were evaluated using both the necessary computing resources, which were quantified as runtime and resource usage data (memory and graphics card usage), and the classification precision, based on the 'ground truth': the theorethical composition of the reference control samples or composition identified by Kraken2 for the gut microbiome samples. We calculated the overall precision as the proportion of reads that the current tool has classified the same as the 'ground truth': We also evaluated the overall coverage, which was achieved by each of the tools. The coverage is a relevant metric due to its ability to clearly show if a tools precision could be an artefact of it either under or over reporting when compared to other tools and was calculated as:  (3) Set up paths (2) Build database Preprocess metadata and sequences (1) Preprocess metadata (5) Split up reference data (4) Build database Set up paths (6) Split up reference data (7) Simulate reference reads (8) Trim reads and convert files Build database Fig. 1. An overall outline of the database creation process. Some practical items that we would like to draw additional attention to include: fastDNA (1) requires all sequences to be uppercase and needs custom meta tags for the sequences, MetaVW (2) requires some external libraries that the tool does not ship with, GeNet (3) needs the reference data to be split up on a taxon-by-taxon basis and (4) requires custom meta tags for the sequences, DeepMicrobes (5) ships with all of the required tools, but they're split up between a multitude of subdirectories, (6) it also requires the reference data to be split up into separate files similarily to GeNet and (7) the training set has to be processed with a read simulator [18]. The scripts that this figure represents are available on GitHub [14].
Classification results (precision and coverage) were evaluated both in raw form, where the resulting taxonomic rank of the organism in the taxonomy tree was defined by the tool, and normalized in 2 different ways to determine the suitability of produced results for practical applications: selecting the results classified at the phylum rank or at the genus rank. This was done by traversing taxonomic tree data that was aquired from NCBI [16].

III. RESULTS AND DISCUSSION
Overall training process showed a high dispersion in terms of time to train/build. Kraken2 and fastDNA showed the best results ( Table 2). In all probability fastDNA outperformed the rest of the tools that leverage neural networks by such a large margin due to it being based on the highly efficient fastText [18]. Interestingly, there is a significant difference between GeNet and DeepMicbrobes and between MetaVW and fastDNA, even though the two pairs use similar methods when approaching the problem of learning the contents of the reference set. The disparity between MetaVW and fastDNA could be explained by the multithreaded capabilities of the tools that they are based on [12] [17]. While the disparity between GeNet and DeepMicrobes comes down to termination conditionswith DeepMicrobes terminating after 1 epoch by default [19] while the implementation of GeNet implies that it should terminate either at 400 epochs or when reaching 99.5% precision, but it did notwe stopped the training process by hand once it had reached 99.5% precision (at epoch 418) [19].
It should be noted that a signicifant portion of the training time for DeepMicrobes was spent on preprocessing the samples for further use. This along with the time that was necesary for Kraken2 to prepare for database being built has been separated out to give a better idea of the time these two tools spends on the build/train process itself (Table 2). Choosing Kraken2 as the baseline when data on sample composition were not available carries a caveat: accuracies for the non-control samples are not necessarily indicative of the best performancethey are indicative of the performance most similar to Kraken2.
Some concessions had to be made when performing the classification tasks. We excluded GeNet from the classification tasks since it had no readily made classification scripts and producing such scripts was out of the scope of this article.
MetaVW and fastDNA required single input files (we found functions that implied that paired end data could be fed into fastDNA, but could not find a way of accessing it without modifying the tool [20]), which meant that samples were merged using bbmerge. Additionally, DeepMicrobes could process controls as is, but required additional utility scripts to be produced to facilitate working with non-control samples due to the large amount of storage space that preprocessing each sample required (a 9 GB gzipped sample expanded to 127 GB in temporary files that resulted in a 86 GB tfrec file). Therefore we batched the samples into random 10% subsamples whose results were then joined; as a result the reported runtimes are based solely on the controls.
Just like when training the model preprocessing samples for DeepMicrobes took a significant amount of time. It is indicated separately in the table in brackets (Table 3). When it comes to the amount of reads that each of the tools reported as classified, DeepMicrobes and MetaVW differentiated themselves from the rest by not reporting any results as unclassified, but always reporting some taxon ID. FastDNA fell on the other end of the spectrum classifying around 8% of the reads. We found Kraken2 to be the only tool where the number of reads that the tool was able to classify was related to the complexity of the sample at hand: the results for control samples composed of few different species were much higher than for the diverse human samples (Table 4). We found that MetaVW outperforms every other tool (even Kraken2) when examining the precision metric of the baseline results (Table 5). It should be noted that these results are skewed as both Kraken2 and DeepMicrobes traverse the taxonomic tree when they are not confident enough to classify something at a certain taxonomic rank of the said tree [3] [7]. But it is still noteworthy that MetaVW can be reasonably used to gather insight from simple samples without the need for normalization scripts. When data is normalized to genus (that is, all of the items above genus are brough down to genus rankall others are marked as unclassified), we see some interesting patterns emerge: the reported percentages of most tools drop, with Kraken2 dropping around half of the classified reads for the non-control samples and DeepMicrobes dropping around 25% of the reads. This increases the overall classification precision of Kraken2 up to around the performance of MetaVW with the other two tools underperforming by more than 60 percentage points when it comes to the controls. When looking at the non-control samples, interestingly, MetaVW drops below both DeepMicrobes and fasDNA in terms of precision. This can mainly be explained by the large amount of unclassified reads present in the Kraken2 results. It is interesting to note that fastDNA performs very closely to Kraken2 when looking at the non-control samples (Table 7). Normalizing to phylum helps most of the tools retain a larger number of reads as expected since phylum is a rank that includes genus [21]. The reduction in coverage is less than one percentage point for all tools except DeepMicrobes ( Table 8).
The pattern that the tools fall into does not differ from the results that we saw when normalizing for genus: precision improves significantly for less diverse samples and there is some improvement from baseline for other samples, although the increase is not as significant as in genus data (Table 9). The coverage of DeepMicrobes diminished significantly during normalization and this pattern was not observed in other tools. Therefore we investigated the taxonomic division of the results for all of the tools. fastDNA provided species results, MetaVW mostly provided species results as well as sometimes responded with taxon IDs that corresponded to genotype (less than 1% of the reads). Kraken2 traversed most of the tree with all returned taxids being valid: Kraken2 classified around 37% of the control sample reads and around 1.4% of the non-control sample reads as species. We found that unlike Kraken2 DeepMicrobes can generate arbitrary taxon IDs that cannot be found in the NCBI taxonomy data: overall DeepMicrobes returned taxon IDs that correspond to species for around 50% of the reads and classified around 20% of the control reads, while 30% of the non-control reads had been assigned nonexistant taxon IDs [14]. The prevalence of these non-existant taxon IDs strongly implies that the training set we used is too small for DeepMicrobes to function properly.

IV. CONCLUSIONS
We found that tools that leverage deep learning still have some catching up to do when comparing their performance to Kraken2 in terms of both memory usage and runtime. With Kraken2 using less than a fifth of the memory of the most memory intensive tool and undercutting some of the tools more than 500 times when looking at runtimes at the scale that we tested at. A notable exception when looking at the runtimes is fastDNAits performance was comparable to Kraken2. It was interesting to see that embedding based tools perform so similarily to Kraken2 when evaluating their coverage and classification precision, even if this similarity in performance is limited to specific types of samples depending on the thresholding strategy that the tool employs. While the performance of DeepMicrobes did not meet the prior expectations (based on results reported in the literature) and we did not manage to fully benchmark GeNet, we still see a lot of potential in using such tools in environments where system memory is limited and researchers are willing to wait for a longer amount of time to obtain the results. It is important to note that this is a statement that heavily hinges on the comparative performance of these tools when they are trained with a suitably large reference set.
When looking at the results, we can see that the most useful and operable tool is still Kraken2, but there is a disadvantageusing Kraken2 with a larger database still requires a significant amount of RAM (around 200 GB [22]). This is where the deep learning based tools should readily outscale Kraken2.
We believe that a different thresholding mechanism for fastDNA could possibly make it competitive with Kraken2 even when using a small reference database such as the one that was used in this study. And that exploring the behaviour of these tools with significantly larger reference databases could yield more insights into how well these tools scale with the amount of available data.
Overall we would like to conclude that a lot of promising progress has been made when it comes to applying deep learning to bioinformatics and we are eager to see what the future will bring. We are hopeful that this paper draws more interest to these tools and applying deep learning to bioinformatics. Additionally we hope for someone with more complex control samples to evaluate the performance of these tools.
An interesting byproduct of performing this study is the realization that these tools are harder to use than we had initially expected. We attempt to make replicating this study easier by outlining the setup process in the repository that contains the scripts that were used to obtain these results. This has been made publicly available on GitHub [14].