# Multiple Instance Learning with Trainable Soft Decision Tree Ensembles
## 6 Attention and the Whole Scheme of STE-MIL  
After training, the output of each neural network corresponding to the \(k\)-th tree is the embedding \(\mathbf{v}_{j,k}^{(i)}\) of length \(E\), where \(i\) and \(j\) are indices of the corresponding bag and instance in the bag, respectively. This implies that we obtain \(T\) embeddings \(\mathbf{v}_{j,1}^{(i)},\ldots,\mathbf{v}_{j,T}^{(i)}\) for the \(j\)-th instance from the \(i\)-th bag, \(j=1,\ldots,n\), \(i=1,\ldots,N\), under an assumption of an identical number of trees in all RFs. It should be noted that numbers of trees in RFs can be different. However, we considered the same numbers for simplicity.  
Embeddings \(\mathbf{v}_{i,1},\ldots,\mathbf{v}_{i,T}\) are aggregated by using, for example, the averaging operation, resulting in vectors \(\mathbf{e}_{j}^{(i)}\), \(j=1,\ldots,n\), corresponding to the \(i\)-th bag. Then, aggregated embeddings \(\mathbf{e}_{1}^{(i)},\ldots,\mathbf{e}_{n}^{(i)}\) are attended to in order to obtain a final representation of the \(i\)-th bag in the form of vector \(\mathbf{a}_{i}\), which is classified. This motivated us to replace the class probability distributions at the tree leaves with the embeddings \(\mathbf{v}\) defined above. We can define several ways to construct embeddings from class probability distributions. However, we selected a simple procedure, which has demonstrated its efficiency from the point of view of accuracy and computationally.  
Hence, the second idea behind STE-MIL is to aggregate the embeddings over all bags by using the attention mechanism and to calculate the prediction logits by the linear projection of the aggregated embedding to the one-dimensional space. This idea is also motivated by the Attention-MIL approach proposed in [23], and by the Multi-attention multiple instance learning model proposed in [25], which may help to train a better bag-level classifier. A scheme of the whole STE-MIL model is shown in Figure 2. It can be seen from Figure 2 that each instance (\(\mathbf{x}_{j}^{(i)}\)) from the \(i\)-th bag learns the corresponding RF such that embeddings are combined to the aggregated vector \(\mathbf{e}_{j}^{(i)}\). Vectors \(\mathbf{e}_{j}^{(i)}\), \(j=1,\ldots,n_{i}\), can be regarded as keys in terms of the attention mechanism. They attend to and produce vector \(\mathbf{a}_{i}\), which is the input of the classifier. The whole system is trained on all instances from all bags.  
The attention module produces a new aggregate embedding \(\mathbf{a}_{k}\) corresponding to the \(k\)-th bag, which is computed as follows:  
\[\mathbf{a}_{k}=\sum_{i=1}^{n}\mathbf{\beta}_{i}^{(k)}\mathbf{e}_{i}^{(k)},\ k=1, \ldots,N, \tag{20}\]  
where  
\[\mathbf{\beta}_{i}^{(k)}=\text{softmax}\Big{(}\mathbf{q}^{\text{T}}\mathbf{k }_{i}\Big{)}, \tag{21}\]  
\[\mathbf{q}=\mathbf{V}_{q}\mathbf{g},\ \mathbf{k}_{i}{=}\mathbf{V}_{k}\mathbf{e}_{i }^{(k)}. \tag{22}\]  
Here \(\mathbf{V}_{k}\) and \(\mathbf{V}_{q}\) are the trainable weight matrices for \(\mathbf{e}_{i}^{(k)}\) (keys) and the template vector \(\mathbf{g}\) (query), respectively.  
Figure 2: A scheme of the ensembled STE-MIL.