# Capsule Network Improved Multi-Head Attention for Word Sense Disambiguation
## 3 Methodology
### Model Details  
In this subsection, we present the details of our model. An overview architecture is depicted in Figure 1. The model encodes the context and sense glosses of an ambiguous word separately, then scores each sense for the target word. The score is calculated by the dot product of contextual embedding from the context module and sense embedding from the sense glosses module. In fact, the encoders of the two modules are exactly the same. In other words, we generate the context embedding and sense embedding in the same way. Inspired by [33; 34], we employ multi-head attention with the capsule network as our encoders. As illustrated in Figure 2, for an input sentence sequence, we initialize each word with BERT word embeddings \(E=(e_{1},e_{2},\dots,e_{L})\in\mathbb{R}^{L\times D}\) as inputs of our model, where \(L\) represents the length of the sequence and \(D\) denotes the word embedding dimension. Figure 1: Overview architecture of BiCapAtt. Figure 2: The architecture of each module in BiCapAtt. Multi-head attention used in transformer [29] focuses on information from various perspectives through splitting the model into multiple subspaces. The attention mechanism first projects the query \(Q\), key \(K\), and value \(V\) to \(H\) different subspaces with linear matrices as shown in Equation (1), where initially \(Q=K=V=E\). Then, scaled dot-product for each head is calculated by Equation (2):  
\[Q_{h},K_{h},V_{h}=QW_{h}^{Q},KW_{h}^{K},VW_{h}^{V};h\in[1,H], \tag{1}\]  
\[\frac{Q_{h}K_{h}^{T}}{\sqrt{d_{k}}}=\left[ \begin{array}{c}e_{1h}\\ \ldots\\ e_{1h}\end{array} \right];h\in[1,H]. \tag{2}\]  
In the original capsule network, the input vector \(\mathbf{u}_{i}\) is multiplied by a pose matrix \(\mathbf{W}_{ij}\), which represents the spatial relationship between low-level features and high-level features. The result of multiplication is \(\mathbf{u}_{ji}\) which indicates the high-level features derived from low-level features. The dynamic routing then is used to better determine the information added to the high-level capsules in the low-level capsules. In multi-head attention, the multiple attention heads can represent different partial information of the input sequence. We therefore treat them as low-level capsules, namely \(I\times J\ \mathbf{u}_{j|i}\) that already contain the spatial relationship. Figure 3 depicts the architecture of the capsule network. The dynamic routing algorithm (DR) we used is described in Algorithm 1. ```
Input:\(I\times J\) vectors \(\mathbf{u}_{j|i}\), iteration times \(t\) Process: \(1\). \(\forall\ i,j:b_{ij}\gets 0\) for\(t\)do \(\forall i:\ c_{i}\gets softmax(b_{i})\) softmax computes Equation (3) \(\forall j:\ s_{j}\leftarrow\sum_{i}c_{i}\mu_{j|i}\) \(\forall j:\ y_{j}\leftarrow\mathrm{squash}(\mathbf{s}_{j})\) squash computes Equation (4) \(\forall i,j:b_{ij}\gets b_{ij}+\mathbf{u}_{j|i}\cdot\mathbf{v}_{j}\) endfor Output: / output vectors \(\mathbf{v}_{j}\), weights \(b_{ij}\)
```  
**Algorithm 1** Dynamic Routing (DR). There are \(I\) input capsules and \(J\) output capsules in Algorithm 1. Each input capsule will generate \(J\) vectors and each vector will be assigned a weight value \(b_{ij}\). First, the input  
Figure 3: The architecture of the capsule network. vectors \(\mathbf{u}_{j\mid i}\) as the shallow capsules are weighted and added to get \(s\) (Algorithm 1, Line 4). The weight \(c_{ij}\) is computed by:  
\[c_{ij}=softmax(b_{i})=\frac{\exp\left(b_{ij}\right)}{\sum_{k}\exp(b_{ik})}, \tag{3}\]  
where \(b_{ij}\) is initialized to zero and is updated as Algorithm 1, Lines 6. The softmax function can map multiple scalars into a probability distribution. Therefore, for each low-level capsule, its weight \(c_{ij}\) defines the probability distribution of the output belonging to each high-level capsule. Then, the non-linear function called squash is used to obtain the deep capsule vector, which is formulated as Equation (4). The squash function is mainly to make the length of \(v\) not exceed \(1\), and keep \(v\) and \(s\) in the same direction.