# Feature Augmentation Based on Pixel-Wise Attention for Rail Defect Detection
## 3 Materials and Methods
### Feature Augmentation  
To utilize the support defect images and improve class-wise defect detection ability, we perform the feature augmentation framework to aggregate input and support features, as illustrated in Figure 3. The framework mainly consists of a query, key, and value embed  
Figure 1: Samples of the rail defect dataset. From left to right, first row contains damage, dirt, unknown and gap, and second row contains dent, crush, scratch and slant. Figure 2: The architecture of the rail defect detection model. ding, and a feature fusion operation based on pixel-wise attention. Although there exist alternative embedding and attention methods, we adopt the transformer-style attention [14] because it has been widely used and proved to be generally effective for various tasks. Query, Key and Value Embedding. The sample and support features extracted by backbone feature extractor are encoded separately to reduce their dimension for saving computation cost. Specifically, sample feature is encoded into a query feature map and a value feature map, while support features are encoded into a concatenated key feature map and a concatenated value feature map. Query and key feature maps contain semantic information about relations, while value maps carry richer detailed information about the original feature maps. The two encoders share the same structure, i.e., two \(3\times 3\) convolution layers, a batch normalization, and a leaky ReLU activation function, but have different weights. The query encoder takes sample feature of size \(H\times W\times C\) as input, and output query feature map \(Q\) and value feature map \(V_{q}\) as:  
\[Q\in\mathbb{R}^{HW\times C_{k}}, \tag{1}\] \[V_{q}\in\mathbb{R}^{HW\times C/2},\]  
where \(H\) is feature height, \(W\) is feature width, \(C\) is channel dimension, and \(C_{k}\) is the user-defined query dimension (e.g., \(C_{k}=C/4\)). On the other hand, the support encoder takes \(N\) support features as input, and output a series of key feature maps and value feature maps. Support key feature maps are horizontally concatenated, while support value maps are vertically concatenated:  
\[K =concat[K_{1}^{T},\cdots,K_{N}^{T}]\in\mathbb{R}^{C_{k}\times NHW}, \tag{2}\] \[V_{s} =concat[V_{s1},\cdots,V_{sN}]^{T}\in\mathbb{R}^{NHW\times C/2}.\]  
Pixel-wise Attention. After acquiring query, key, and value feature maps, pixel-wise attention is performed to activate co-existing patterns among sample and support features. The key of pixel-wise attention lies in a learnable pixel-level similarity metric. To calculate the metric, two different linear transformation \(\phi,\phi^{\prime}\) are first applied to query and key feature maps:  
\[Q^{*}=\phi(Q),\quad K^{*}=\phi^{\prime}(K). \tag{3}\]  
Then, linear transformed query and key feature are fused by inner-production to obtain the pixel-level similarity metric:  
\[S=Q^{*}\cdot K^{*}=[S^{1},S^{2},\cdots,S^{N}]\in\mathbb{R}^{HW\times NHW}. \tag{4}\]  
Figure 3: Feature augmentation framework. The sample and support features are first encoded into query (Q), key (K), and value (V) embedding. Two linear transformations are performed on both Q and K feature maps. Pixel-wise attention is performed to fuse the Q and K feature maps to form the final learnable similarity metric, thus retrieving co-existing information from support features. This similarity metric concerning \(\phi,\phi^{\prime}\) can be dynamically learned through back-propagation during fine-tuning. After obtaining the fused similarity metric, global softmax normalization is performed to obtain the final attention weight:  
\[\omega=\left[ \begin{array}{cccccc}\omega_{1,1}^{1}&\cdots&\omega_{1,1}^{2}& \cdots&\cdots&\omega_{1,HW}^{N}\\ \vdots&\ddots&\vdots&\ddots&\omega_{j,k}^{i}&\vdots\\ \omega_{HW,1}^{1}&\cdots&\omega_{HW,1}^{2}&\cdots&\cdots&\omega_{HW,HW}^{N} \end{array} \right], \tag{5}\]  
\[\omega_{j,k}^{i}=\frac{exp(S_{j,k}^{i})}{\sum_{NHW}exp(S_{j,k}^{i})},\]  
where \(i=1,2,\cdots,N\) is the support index, \(j,k=1,2,\cdots,HW\) is the pixel index of the similarity metric.