EMA-SAM: Exponential Moving-average for SAM-based PTMC Segmentation

TL;DR

This work tackles real-time PTMC tumor segmentation in ultrasound during RFA, a setting plagued by low contrast, probe motion, and bubble-induced occlusions. It extends the Segment Anything Model-2 with a confidence-weighted exponential moving-average pointer that is never evicted from memory and is up-weighted during cross-attention, yielding a stable latent tumor prototype $\mathbf{m}_t$ that adapts when reliable evidence reappears via $\mathbf{p}_t$, $\alpha_t = \alpha_0(1 - c_t)$, $\mathbf{m}_t = \alpha_t \mathbf{m}_{t-1} + (1 - \alpha_t) \mathbf{p}_t$, and $\||\mathbf{m}_t||_2 = 1$. This EMA pointer is appended to the memory with a gain $\gamma$ so the decoder remains guided by a robust prior, effectively mitigating drift during occlusion without sacrificing responsiveness. Across a PTMC-RFA dataset (13 patients, 124 minutes), EMA-SAM improves maxDice from $0.82$ to $0.86$ and maxIoU from $0.72$ to $0.76$, while reducing false positives by $29\%$, and it achieves additional gains on VTUS and colonoscopy polyp benchmarks. Importantly, the approach adds less than $0.1\%$ FLOPs and sustains ~30 FPS on an A100 GPU, offering a practical, real-time solution that bridges foundation-model segmentation with the stringent demands of interventional ultrasound.

Abstract

Papillary thyroid microcarcinoma (PTMC) is increasingly managed with radio-frequency ablation (RFA), yet accurate lesion segmentation in ultrasound videos remains difficult due to low contrast, probe-induced motion, and heat-related artifacts. The recent Segment Anything Model 2 (SAM-2) generalizes well to static images, but its frame-independent design yields unstable predictions and temporal drift in interventional ultrasound. We introduce \textbf{EMA-SAM}, a lightweight extension of SAM-2 that incorporates a confidence-weighted exponential moving average pointer into the memory bank, providing a stable latent prototype of the tumour across frames. This design preserves temporal coherence through probe pressure and bubble occlusion while rapidly adapting once clear evidence reappears. On our curated PTMC-RFA dataset (124 minutes, 13 patients), EMA-SAM improves \emph{maxDice} from 0.82 (SAM-2) to 0.86 and \emph{maxIoU} from 0.72 to 0.76, while reducing false positives by 29\%. On external benchmarks, including VTUS and colonoscopy video polyp datasets, EMA-SAM achieves consistent gains of 2--5 Dice points over SAM-2. Importantly, the EMA pointer adds \textless0.1\% FLOPs, preserving real-time throughput of $\sim$30\,FPS on a single A100 GPU. These results establish EMA-SAM as a robust and efficient framework for stable tumour tracking, bridging the gap between foundation models and the stringent demands of interventional ultrasound. Codes are available here \hyperref[code {https://github.com/mdialameh/EMA-SAM}.

Figures (4)

• Figure 1: Workflow of the proposed EMA-SAM. The ultrasound frame $\mathbf{I}_{t}$ is encoded into tokens $\mathbf{E}_{t}$, combined with the FIFO memory bank $\mathcal{B}_{t}$, and processed by the memory–attention module $\mathcal{A}$. Our novelty is the confidence-weighted EMA pointer $\mathbf{m}_{t}$ (blue), which is never evicted and is up-weighted in cross-attention (gain $\gamma$), providing a stable latent prototype of the PTMC even when the tumour disappears due to probe pressure or bubble shadowing.
• Figure 2: Qualitative comparison of segmentation continuity across time. The first row shows original ultrasound video frames at three time steps, where the PTMC region, if any, is highlighted by a different color. The second row presents SAM2 predictions (overlayed by a light green color), which fail to recover the PTMC segmentation once it disappears (in middle frames, it wrongly predicts a non-PTMC region as a PTMC). The third row shows EMA-SAM results, where the model successfully retains spatial memory of the PTMC location and continues consistent segmentation despite transient invisibility. The dots between columns indicate that intermediate frames are omitted from the visualization.
• Figure 3: Visualization of tumour localisation during RFA using Grad-CAM++ chattopadhay2018grad. Consecutive ultrasound frames (left) are shown alongside predicted masks (middle) and Grad-CAM++ heatmap overlays (right). The heatmaps reveal the gradient-based regions that most strongly influenced the model’s segmentation decision. Red areas indicate high-confidence tumour localisation, while surrounding cooler colors (green/blue) reflect less discriminative regions. This visualization confirms that the network consistently attends to the PTMC lesion across frames despite acoustic shadows and probe-induced tissue deformation.
• Figure 4: Temporal IoU across an entire RFA sequence. Frame-wise IoU vs. ground truth for SAM-2 and EMA-SAM. The plot indicates spike events (large negative deviations from a rolling baseline), typically aligned with occlusions or rapid probe motions. EMA-SAM displays a flatter curve with fewer and shallower spikes, maintaining stable tumour outlines even during seconds-long disappearances.