Let Video Teaches You More: Video-to-Image Knowledge Distillation using DEtection TRansformer for Medical Video Lesion Detection

The pipeline for V2I-DETR is illustrated in Fig. 2, which is developed upon a teacher-student architecture. Both the teacher and student are implemented using the same DEtection TRansformer model. During training, the teacher model takes multiple images as input, including current frame $I^{c}\in\mathbb{R}^{C\times H\times W}$ and several reference frames $I^{r}\in\mathbb{R}^{C\times(T-1)\times H\times W}$. Through the Multi-scale Spatiotemporal Interaction (MSI) module, the teacher model adeptly captures both spatial and temporal contexts from these diverse frames. Subsequently, these learned contexts are effectively transferred to the student model. This transfer is facilitated by techniques like Target-guided Feature Distillation (TFD) and Cross-view Query Distillation (CQD), which contribute to enhancing the student model’s understanding and performance. During inference, the teacher model is disregarded, and only the student model is utilized for making predictions, which takes a single image (current frame $I^{c}$) as input, where $C,T,H,W$ are the channel, frame number, image height, and width, respectively. Further specifics of V2I-DETR will be discussed in the following sections.

As indicated by [12], video inherits a high degree of contextual redundancy and dependency. Specifically, low-level features typically rely on local rather than global contexts to learn object relations, but high-level features require global features to build long-range dependencies across distant frames. Motivated by this, we propose the MSI module to enable low-level/high-level feature interaction between different scales of multiple frames. As shown in Fig. 3, the MSI comprises two interaction layers, each of which consists of a temporal interact block $\mathcal{A}$ and a Feed-Forward Network (FFN). Given the multi-scale teacher features $F_{t}^{l}\in\mathbb{R}^{d\times T\times h\times w},~{}l=0,1,2,3$, where $h=\frac{H}{2^{l}},w=\frac{W}{2^{l}}$ are the height and width of feature at different level $l$. The enhanced feature $Y_{t}^{l}$ at each level can be obtained:

where PE is the patch embedding to integrate 3D position information into tokens, and $\mathcal{A}^{l}$ is the interact block of level $l$. In the shallow layers ($l=0,1$), we extract the spatial and temporal affinity between tokens in small volumetric regions by local interactor using the 3D convolutional layer to encourage the network to capture contextual information between adjacent frames. In the deep layers ($l=2,3$), we focus on capturing long-term token dependency by global interactor. We compute the feature similarity of all tokens by self-attention mechanisms [13]. Hence, our model exhibits insensitivity to the changes in the size and location of moving objects.

The detection performance is greatly influenced by the feature representations of the backbone, primarily due to their abundance of semantic information associated with objects. Therefore, it is necessary to mimic the spatiotemporal features of the teacher model to boost the performance of the student model. An intuitive manner to transfer the teacher’s knowledge to the student is reducing the Euclidean distance between feature $F^{3}_{s}$ and $Y^{3}_{t}$ [14]:

where $N$ is the total pixel number, $(\text{F}^{3}_{s})$ and $(F^{3}_{t})$ denote the last-stage feature of the current frame from the student and teacher backbone, respectively. However, foreground objects are relatively smaller compared to backgrounds. If features from the teacher model and the student model are directly aligned, foreground features might be overshadowed and a subsequent drop in performance.

Hence, we employ the ground truth bounding boxes as indicators to guide the model to focus more on the foreground portions of each image. Furthermore, instead of adopting a rigid hard selection method like [15], which exclusively emphasizes foreground pixels within the ground truth box for distillation points, we opt for a more adaptable approach by employing soft selection since the background information also holds valuable insights and contributes to improved model generalization. Specifically, given a bounding box $B$ of an object, with the size of $H\times W$ and centered at $\tilde{x},\tilde{y}$, we splat all ground truth points onto a two-dimensional Gaussian heatmap $H=exp(-\frac{(x-\tilde{x})^{2}+(y-\tilde{y})^{2}}{2\sigma_{p}^{2}})$, where $(x,y)\in B$ and the $\sigma_{p}$ is object size-adaptive factor along two directions. The final soft selection mask is obtained by normalization $M=\frac{H}{||H||_{2}}$. In particular, if several Gaussian masks overlap over a single pixel $(x,y)$, we take the element-wise maximum. Fig. 4a illustrates some images and the corresponding Gaussian masks. The soft selection mask is only effective within the activated Gaussian mask, thus only useful information is transferred to the student model. With the delicately designed Gaussian mask, the backbone features are distilled via minimizing the following loss:

Based on previous research [16], well-optimized queries in the teacher model consistently achieve a stable bipartite assignment and improve the training stability of the student model. Thus, previous image-to-image knowledge distillation methods for DETR directly apply consistency regularization between object queries. However, in the multi-frame collaboration paradigm, the information in the current frame and the reference frame is not completely consistent due to the camera moving.

To overcome this issue, we propose cross-view query distillation that enables the DETR-based framework to learn semantically invariant characteristics of object queries from two different views. The details of the cross-view query distillation process are shown in Fig. 2. Specifically, for object query $\text{Q}_{t}\in\mathbb{R}^{N\times d}$ of each frame in the teacher model, we randomly select $n$ queries $q_{t}\in\mathbb{R}^{n\times d}=\{\text{Q}_{t}^{i},~{}i\in\Phi(N)\}$ from $\text{Q}_{t}$ as the cross-view query embeddings, where $n=\lfloor N/T\rfloor$ and $\Phi(\cdot)$ is a random selection function that returns a sequence of list index of the teacher object query. We propose the random selection strategy to regulate the number of cross-view queries to be roughly equal to the number of student queries, thus student queries will not be overwhelmed by larger cross-view queries. Then the cross-view query embedding is attached to the original student object queries in another view to serve as the input of the decoder:

To sum up, the total loss for training the V2I-DETR is a weighted combination of bounding box supervision and knowledge distillation:

where $\mathcal{L}_{det}$ is the sum of f1 loss and giou loss as described in [6], $\lambda_{bk}$, and $\lambda_{qe}$ are the weights for each individual distillation loss. To guarantee the loss terms are roughly of the same magnitude, we fine-tune weights with various combinations as shown in Fig. 5(a), and empirically fix those to $\lambda_{bk}=1$, and $\lambda_{qe}=5$ in the experiments.

The results of colonoscopy lesion detection are displayed in Tab. I. Primarily, the video-based teacher model, utilizing multiple frames as inputs, exhibits a significant performance advantage over the original Deformable DETR by $9.2\%$ and $7.6\%$ f1-score on the easy and hard test set. However, The cost of performance improvement is very low inference frames per second (FPS). In contrast, our V2I-DETR model, employing knowledge distillation, consistently enhances the performance of the student model. It yields improvements of f1-score by $7.2\%\sim 9.0\%$ on both test sets. Specifically, V2I-DETR based on SAM DETR achieves the state-of-the-art (SOTA) results, attaining an f1-score of $85.5\%$ (easy) and $84.3\%$ (hard). Notably, V2I-DETR will not introduce any computational overhead into the inference stage and achieve $\times 7$ faster FPS. Even compared with video-based models, our V2I-DETR based on SAM DETR surpasses the SOTA method YONA by $4.3\%$ (easy) and $3.6\%$ (hard) f1-score. Tab. II demonstrates the comparison on the BUV benchmark. Specifically, among the comparison methods, the video-based CVA-Net achieves the highest performance of 36.1% AP, 65.1% $\text{AP}_{50}$, 38.5% $\text{AP}_{75}$. Compared to CVA-Net, our V2I-DETR based on Deformable DETR improves the AP by 0.7% and $\text{AP}_{75}$ by 1.2%. While our V2I-DETR based on SAM DETR further boosts the performance by 1.3% AP and 1.7% $\text{AP}_{75}$.

Fig. 4b visually compares detection results for polyp and breast lesion detection results produced by our network and four compared methods. Lack of temporal consistency information, image-based model Deformable DETR fails to detect the polyp and breast lesions, as shown in the first row of Fig. 4b, STFT and YONA obtain a better lesion detection result but also fail to detect lesions or produce many false positive results ($3^{rd},4^{th}$ in polyp and $2^{nd}$ in the breast). In contrast, our method can correctly detect the lesions of all video frames. Confirming that our method is effective in precisely detecting lesion regions from moving-object medical videos.

Effectiveness of component We investigate the contribution of each proposed component. The results are shown in Tab. III. Compared with the original Deformable DETR (# 1), our proposed components enjoy consistent performance improvement. Specifically, by introducing the TFD (# 2) for feature distillation, it outperforms the baseline by $3.8\%$ on the f1-score. Further integrating the MSI (# 3) for spatiotemporal feature interaction in the teacher model brings an improvement of $6.4\%$. Finally, by combining both feature and query distillation, our V2I-DETR framework w/o (# 4) or w/ (# 5) MSI can significantly boost the performance by $7.3\%$ and $8.6\%$ on f1-score. These results show that our proposed components are complementary to each other and therefore prove the effectiveness of each component in V2I-DETR.

Analysis on distillation feature Tab. IV summarized the results of distillation from different stages of DETR. Firstly, we apply direct knowledge transfer with MSE loss to the baseline model (# 2). This design slightly improves the performance by $1.7\%$ f1-score, showing that mimicking the teacher backbone learning during training is helpful to the student model. However, distilling the feature embedding from the DETR encoder (# 6) causes an obvious performance drop by $4.0\%$ f1-score. we further apply the self-attention mechanism introduced in [23] to restrict the region of interest (# 7), but the performance is still worse than the baseline. Lastly, we transfer the teacher decoders’ object query to the student, which shows a positive contribution to the performance with $1.2\%$ f1-score promotion. Therefore, we mainly focus on the design of FD and QD in our model.

Analysis on the region of FD Based on the above observation, we dig deeply with different features distillation (FD) strategies as shown in the first part of Tab. IV. We first follow the idea in [23] that computes the attention mask between object queries and backbone features to select the valuable region for distillation (# 4). However, the improvement is minimal compared with vanilla distillation (# 2) and requires significant additional computation. Next, we guided the distillation using masks generated from the ground truth boxes, which improved precision by 2.6% in F1-score. Finally, we converted the hard ground truth masks to soft masks using a Gaussian distribution. This further improved performance by 3.8%, demonstrating that our method effectively balances the information around the target boundaries.

Analysis on the number of CQD Meanwhile, we also conduct several ablations on different cross-view numbers for Cross-view Query Distillation (CQD), as shown in The third part of Tab. IV. First, we manually set the number equal to the number of the sum of all teacher queries (# 9). However, because this number was much larger than the number of student queries, it misled the student queries and significantly reduced performance. Reducing the ratio could mitigate the undesirable result (# 10). Lastly, we use random sampling to keep the cross-view query number and student one to be approximately same (# 11), resulting in an improvement of $4.2\%$ on f1-score.

Analysis on the number of training reference We investigated the impact of different support frame numbers on V2I DETR in Fig. 5(b). We keep the inference frame number as $1$ and change the number of training frame numbers. Training with 3 frames achieves the best balance between accuracy and computational complexity (5 frames reach the memory cap).