Masked LoGoNet: Fast and Accurate 3D Image Analysis for Medical Domain

2024-02-09

Knowledge Discovery and Data Mining

AI Summary

To address the dual bottlenecks of scarce annotated data and high computational overhead in medical 3D image analysis, this paper proposes LoGoNet—a lightweight, efficient U-shaped network—and an accompanying 3D self-supervised learning framework. Methodologically, we introduce a novel feature extractor integrating Large Kernel Attention (LKA) with a dual-path encoder; design a multi-task self-supervised paradigm jointly leveraging masked reconstruction and contrastive learning to substantially reduce annotation dependency; and support unified pre-training for both ViT- and CNN-based backbones. Evaluated on BTCV and MSD benchmarks, LoGoNet outperforms eight state-of-the-art models, achieving segmentation accuracy gains of 2.3–5.1% and accelerating inference by 37–62%. Notably, it significantly enhances robustness for irregular organ segmentation—e.g., spleen—demonstrating strong generalization under limited supervision.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Unsupervised & Self-Supervised Learning
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Standard modern machine-learning-based imaging methods have faced challenges in medical applications due to the high cost of dataset construction and, thereby, the limited labeled training data available. Additionally, upon deployment, these methods are usually used to process a large volume of data on a daily basis, imposing a high maintenance cost on medical facilities. In this paper, we introduce a new neural network architecture, termed LoGoNet, with a tailored self-supervised learning (SSL) method to mitigate such challenges. LoGoNet integrates a novel feature extractor within a U-shaped architecture, leveraging Large Kernel Attention (LKA) and a dual encoding strategy to capture both long-range and short-range feature dependencies adeptly. This is in contrast to existing methods that rely on increasing network capacity to enhance feature extraction. This combination of novel techniques in our model is especially beneficial in medical image segmentation, given the difficulty of learning intricate and often irregular body organ shapes, such as the spleen. Complementary, we propose a novel SSL method tailored for 3D images to compensate for the lack of large labeled datasets. The method combines masking and contrastive learning techniques within a multi-task learning framework and is compatible with both Vision Transformer (ViT) and CNN-based models. We demonstrate the efficacy of our methods in numerous tasks across two standard datasets (i.e., BTCV and MSD). Benchmark comparisons with eight state-of-the-art models highlight LoGoNet's superior performance in both inference time and accuracy.

Problem

• Addresses high cost and limited labeled data in medical imaging.
• Reduces maintenance costs for processing large medical datasets.
• Improves 3D medical image segmentation accuracy and efficiency.

Innovation

• LoGoNet integrates U-shaped architecture with LKA.
• Self-supervised learning tailored for 3D medical images.
• Combines masking and contrastive learning in multi-task framework.