EdgeNAT: Transformer for Efficient Edge Detection

Below is a brief introduction on DiNAT, encoder of our network, following the work presented in [17].

To begin with, DiNAT employs two 3 × 3 convolutional layers with a stride of 2 as a tokenizer to obtain a feature map with a resolution of one-fourth of the input image. Additionally, DiNAT utilizes a single 3 × 3 convolutional layer with a stride of 2 for downsampling between hierarchical levels, reducing the spatial resolution by half while doubling the number of channels.The resulting feature maps are thus of sizes $\frac{h}{4}\times\frac{w}{4}\times{c}$, $\frac{h}{8}\times\frac{w}{8}\times{2c}$, $\frac{h}{16}\times\frac{w}{16}\times{4c}$ and $\frac{h}{32}\times\frac{w}{32}\times{8c}$.

DiNAT adopts a straightforward stacking of DiNA layers, following a similar structural pattern as other commonly used Transformers. For simplicity, we keep notations limited to single dimensional NA and DiNA. Given input $X\in\mathbb{R}^{n\times d}$, whose rows are d-dimensional token vectors, and query and key linear projections of $X$, $Q$ and $K$, and relative positional biases between any two tokens $i$ and $j$, $B(i,j)$, $\delta$-dilated neighborhood attention weights for the $i$-th token with neighborhood size $k$, $\mathbf{A}_{i}^{(k,\delta)}$, is defined as the matrix multiplication of the $i$-th token’s query projection, and its k nearest neighboring tokens’ key projections with dilation value $\delta$:

$$\mathbf{A}_{i}^{(k,\delta)}=\left[\begin{array}[]{c}Q_{i}K_{\rho_{1}^{\delta}(i)}^{T}+B_{\left(i,\rho_{1}^{\delta}(i)\right)}\\ Q_{i}K_{\rho_{2}^{\delta}(i)}^{T}+B_{\left(i,\rho_{2}^{\delta}(i)\right)}\\ \vdots\\ Q_{i}K_{\rho_{k}^{\delta}(i)}^{T}+B_{\left(i,\rho_{k}^{\delta}(i)\right)}\end{array}\right],$$

(1)

where $\rho_{j}^{\delta}(i)$ denotes $i$’s $j$-th nearest neighbor with dilation value $\delta$. The $\delta$-dilated neighboring values for the $i$-th token is similarly defined with neighborhood size $k$, $\mathbf{V}_{i}^{(k,\delta)}$:

$$\mathbf{V}_{i}^{(k,\delta)}=\left[\begin{array}[]{llll}V_{\rho_{1}^{\delta}(i)}^{T}&V_{\rho_{2}^{\delta}(i)}^{T}&\cdots&V_{\rho_{k}^{\delta}(i)}^{T}\end{array}\right]^{T},$$

(2)

where $V$ is a linear projection of $X$. DiNA output for the $i$-th token neighborhood size $k$ with dilation value $\delta$ is then defined as:

$$\operatorname{DiNA}_{k}^{\delta}(i)=\operatorname{softmax}\left(\frac{\mathbf{A}_{i}^{(k,\delta)}}{\sqrt{d_{k}}}\right)\mathbf{V}_{i}^{(k,\delta)},$$

(3)

where $\sqrt{d}$ is the scaling parameter, and $d$ is the embedding dimension. This operation is repeated for every pixel in the feature map.

Summary of DiNAT configuartions and dilation values will be provided in the supplementary material.

Decoders play a critical role in various vision tasks. Taking inspiration from multilevel feature fusion techniques employed in vision tasks, we propose a novel decoder, SCAF-MLA, to effectively utilize numerous channels in the feature maps output from transformer-based encoder. SCAF-MLA enables the supervision on multiple levels, and learns rich hierarchical features, thus enhances the performance of edge detection. Besides, SCAF-MLA Decoder is more computationally efficient, without the commonly employed PPM[47] and bottom-up path[19, 32], while experimental results demonstrate that our designed decoder achieves more superior performance.

Inspired by UAFM[31] in multi-level features fusing, we propose the Spatial and Channel Attention Fusion Module (SCAFM) as the main component of the SCAF-MLA. SCAFM is designed to extract both spatial and channel features, concurrently preserving the distinctive attributes of the current level while capturing higher-level features. The architecture of SCAFM is depicted in Figure 3. SCAFM consists of a spatial attention module (SAM) and a channel attention module (CAM) to compute inter-spatial and inter-channel weights, denoted as $\alpha_{\text{Sp}}$ and $\alpha_{\text{Ch}}$, respectively. Specifically, the upper-level feature is denoted as ${F}_{high}$ and the current-level feature as ${F}_{low}$. To begin with, bilinear interpolation is employed to upsample ${F}_{high}$ to the same size as ${F}_{low}$. Subsequently, convolutional operations is utilized to increase the channels of ${F}_{low}$ to match those of ${F}_{high}$, denoted as ${F}_{conv}$. For $\alpha_{\text{Sp}}$, we start by performing mean and max operations along the channel dimension on ${F}_{up}$ and ${F}_{conv}$, resulting in the generation of four features, each with a dimension of $\mathbb{R}^{1\times H\times W}$. Subsequently, these four features are concatenated and processed through convolutional and sigmoid operations, yielding $\alpha_{\text{Sp}}\in\mathbb{R}^{1\times H\times W}$. This process can be represented by thefollowing equations:

$\displaystyle\begin{split}F_{up}&=Up(F_{high}),\\ F_{conv}&=Conv(F_{low}),\\ \alpha_{\text{Sp}}&=Sigmoid(Conv(Cat(Mean(F_{up}),\\ &Max(F_{up}),Mean(F_{conv}),\\ &Max(F_{conv})))),\end{split}$

(4)

Regarding $\alpha_{\text{Ch}}$, average pooling and max pooling operations are applied on ${F}_{up}$ and ${F}_{conv}$, generating four features with dimensions $\mathbb{R}^{C\times 1\times 1}$. Then, these features are concatenated and subjected to convolutional and sigmoid operations, generating $\alpha_{\text{Ch}}\in\mathbb{R}^{C\times 1\times 1}$ , described as:

$\displaystyle\begin{split}\alpha_{\text{Ch}}&=Sigmoid(Conv(Cat(AvgPool(F_{up}),\\ &MaxPool(F_{up}),AvgPool(F_{conv}),\\ &MaxPool(F_{conv})))),\end{split}$

(5)

The input features are then fused with the generated weights $\alpha_{\text{Sp}}$ and $\alpha_{\text{Ch}}$ through multiplication and addition operations, resulting in features $F_{Sp}$ and $F_{Ch}$. Subsequently, these features are concatenated and convolved, generating $F_{SC}\in\mathbb{R}^{C\times H\times W}$. Then, we perform convolutions on $F_{high}$ followed by upsampling, generating features $F_{H}$ with dimensions $\mathbb{R}^{C\times H\times W}$. Afterwards, $F_{H}$, $F_{SC}$, and $F_{low}$ are concatenated and convolved to obtain the fused feature $F_{out}\in\mathbb{R}^{C\times H\times W}$. The aforementioned process can be described as:

$\displaystyle\begin{split}F_{up}=&Up(Conv(F_{high})),\\ F_{Sp}=&F_{up}\cdot\alpha_{\text{Sp}}+F_{conv}\cdot(1-\alpha_{\text{Sp}}),\\ F_{Ch}=&F_{up}\cdot\alpha_{\text{Ch}}+F_{conv}\cdot(1-\alpha_{\text{Ch}}),\\ F_{SC}=&Conv(Cat(F_{Sp},F_{Ch})),\\ F_{out}=&Conv(Cat(F_{SC},F_{H},F_{low})),\\ \end{split}$

(6)

We employ the loss function proposed in [42] for the 4 side edge maps and 1 primary edge map. Given an edge map $E$ and its corresponding ground truth $Y$, the loss function is computed as follows:

$\displaystyle\begin{split}\ell(E,Y)=&-\sum_{i,j}(Y_{i,j}\alpha\log\left(E_{i,j}\right)\\ +&\left(1-Y_{i,j}\right)(1-\alpha)\log\left(1-E_{i,j}\right)),\end{split}$

(7)

where $E_{i,j}$ and $Y_{i,j}$ are the $(i,j)^{th}$ element of matrix $E$ and $Y$, respectively. $\alpha=\frac{|Y^{-}|}{|Y^{-}|+|Y^{+}|}$ represents the percentage of negative pixel samples, with $|Y^{+}|$ and $|Y^{-}|$ denoting the number of positive and negative sample pixels, respectively. Since BSDS500 dataset is annotated by multiple annotators, we first normalize the multiple annotations into edge probability maps within the range of [0, 1]. Then, if the probability of a pixel is greater than a threshold value $\eta$, it is labeled as a positive sample; otherwise, it is labeled as a negative sample.

After the concatenation operation of the four feature maps output from our decoder, two 3×3 convolutions and one 1×1 convolution are applied to reduce the dimension of the concatenated feature maps. Similarly, a 3×3 convolution and a 1×1 convolution are applied to reduce the dimension of the four side feature maps. Subsequently, a sigmoid operation is performed on each of these reduced-dimensional feature maps to generate primary edge map and four side edge maps, denoted as $\mathcal{E}$, $\mathcal{S}_{1}$, $\mathcal{S}_{2}$, $\mathcal{S}_{3}$, and $\mathcal{S}_{4}$. We calculate the loss for both primary edge map and side edge maps to introduce additional supervision. To sum up, the overall loss function is as follows:

$$\mathcal{L}=\mathcal{L}^{\mathcal{E}}+\lambda\mathcal{L}^{\mathcal{S}}=\ell(\mathcal{E},Y)+\lambda\sum_{k=1}^{4}\ell\left(\mathcal{S}^{k},Y\right),$$

(8)

$\mathcal{L}^{\mathcal{E}}$ and $\mathcal{L}^{\mathcal{S}}$ represent the losses for the primary edge map $\mathcal{E}$ and the side edge maps $\mathcal{S}_{1},\mathcal{S}_{2},\mathcal{S}_{3},\mathcal{S}_{4}$, respectively. Meanwhile, $\lambda$ denotes the weight that balances $\mathcal{L}^{\mathcal{E}}$ and $\mathcal{L}^{\mathcal{S}}$. Based on [32] and our experimental observations, we set $\lambda$ to 0.4.

Two mainstream datasets are used to evaluate our proposed EdgeNAT, namely, BSDS500 and NYUDv2.

Our EdgeNAT is implemented with PyTorch and is based on mmsegmentation[7] and NATTEN[18]. We use the pre-trained weights of DiNAT[17] to initialize EdgeNAT’s transformer blocks. To generate binary edge maps, for BSDS500, we set the threshold $\eta$ to 0.3 to select positive samples. For NYUDv2, there is only one annotation per picture, so there is no need to set the threshold $\eta$.

We use the AdamW optimizer and train for 40k iterations using a cosine decay learning rate scheduler, where the first 15k iterations warm up the learning rate in a linear manner, and the remaining ones are decayed according to the scheduler. The initial learning rate is 0 and a preset learning rate is set to 6e-5. For BSDS500, we set its batch size to 8, and for NYUDv2, we set its batch size to 4.

All experiments were conducted on RTX 4090 GPU. The training of the L model of EdgeNAT (472.38MB) takes 6 hours, far more efficient than Transformer-based model EDTER (468.84MB)[32], which takes 26.4 hours. The inference runs at 20.87 FPS on RTX 4090, nearly ten times the speed of EDTER on V100 (2.2 FPS). During training, since our model is a one-stage edge detection model, for 320×320 images, the GPU memory requirement is about 20GB, 2/3 of EDTER(29GB).

Optimal Dataset Scale (ODS) and Optimal Image Scale (OIS) are two metrics for all datasets. Before evaluation, we perform non-maximum suppression on the predicted edge maps. For the maximum allowed tolerance distance between the detected edge and ground truth, we set it to 0.0075 for BSDS500 and to 0.011 for NYUDv2 as in previous works.

Ablation experiments are performed on the BSDS500 data set to verify the effectiveness of our proposed decoder. Specifically, we first compare the effect of pre-fusion (reduce the channels of feature map to C) and final-fusion (reduce the channels of feature map to 1); then the effect of bottom-up path is also verified. From the quantitative results shown in Table 1, it is clear that regardless of pre-fusion or final-fusion, Bottom-up Path has negative effects on edge detection performance, indicating it is not suitable for DiNAT. For edge detection models with relatively large number of feature map channels, pre-fusion without PPM will be a better choice.

Next, the effectiveness of PPM[47] and our proposed SCAFM are verified and compared. The quantitative results shown in Table 2 demonstrates that SCAFM works best without PPM, achieving best ODS and OIS score, 84.3% and 85.9% respectively. In summary, we will use the SCAF-MLA decoder without Bottom-up Path and PPM for the next experiments.

EdgeNAT-L has a relatively large amount of parameters (472.38MB). In order to adapt to different application scenarios, we conduct scalability experiments on different model sizes. The configuration settings of the encoder of the L, S0, S1, S2, and S3 variants of our EdgeNAT are the same as those of the Large, Mini, Tiny, Small, and Base versions of the DiNAT[17]. Extensive experiments are conducted to study the scalability and throughput of EdgeNAT variants. The result is shown in Figure 4. The models are all trained using the BSDS500 training and validation sets with or without PASCAL VOC, and evaluated with the BSDS500 test set. As expected, when the size of our model decreases, the ODS and OIS will decrease accordingly, and the throughput increases.

It is worth noting that the processing speed of the S0 model is much higher than that of other models. This should be contributed to the fact that its third level has only 6 layers, while the others models have 18 layers. The ODS and OIS of the L model are much higher than other models, due to the fact that the encoder is pre-trained on ImageNet-22K, while encoder of other models are pre-trained on ImageNet-1K. The results of multi-scale input experiment of S0, S1, S2, and S3 models as well as their visualization results will be provided in the supplementary material.