Self-Supervised Bi-Channel Transformer Networks for Computer-Aided Diagnosis
Ronglin Gong, Xiangmin Han, Jun Wang, Shihui Ying, Jun Shi
Abstract
Self-supervised learning (SSL) can alleviate the issue of small sample size, which has shown its effectiveness for the computer-aided diagnosis (CAD) models. However, since the conventional SSL methods share the identical backbone in both the pretext and downstream tasks, the pretext network generally cannot be well trained in the pre-training stage, if the pretext task is totally different from the downstream one. In this work, we propose a novel task-driven SSL method, namely Self-Supervised Bi-channel Transformer Networks (SSBTN), to improve the diagnostic accuracy of a CAD model by enhancing SSL flexibility. In SSBTN, we innovatively integrate two different networks for the pretext and downstream tasks, respectively, into a unified framework. Consequently, the pretext task can be flexibly designed based on the data characteristics, and the corresponding designed pretext network thus learns more effective feature representation to be transferred to the downstream network. Furthermore, a transformer-based transfer module is developed to efficiently enhance knowledge transfer by conducting feature alignment between two different networks. The proposed SSBTN is evaluated on two publicly available datasets, namely the full-field digital mammography INbreast dataset and the wireless video capsule CrohnIPI dataset. The experimental results indicate that the proposed SSBTN outperforms all the compared algorithms.