Optimal transport unlocks end-to-end learning for single-molecule localization

TL;DR

This paper introduces an optimal-transport loss and an iterative neural network architecture for single-molecule localization microscopy, enabling end-to-end training and improved performance at high emitter densities.

Contribution

It reformulates SMLM training as a set-matching problem using optimal transport, removing the need for non-differentiable NMS layers and enhancing denser emitter localization.

Findings

Outperforms state-of-the-art methods at moderate emitter densities.

Achieves superior results at high emitter densities.

Enables end-to-end training without NMS during inference.

Abstract

Single-molecule localization microscopy (SMLM) allows reconstructing biology-relevant structures beyond the diffraction limit by detecting and localizing individual fluorophores -- fluorescent molecules stained onto the observed specimen -- over time to reconstruct super-resolved images. Currently, efficient SMLM requires non-overlapping emitting fluorophores, leading to long acquisition times that hinders live-cell imaging. Recent deep-learning approaches can handle denser emissions, but they rely on variants of non-maximum suppression (NMS) layers, which are unfortunately non-differentiable and may discard true positives with their local fusion strategy. In this presentation, we reformulate the SMLM training objective as a set-matching problem, deriving an optimal-transport loss that eliminates the need for NMS during inference and enables end-to-end training. Additionally, we propose an iterative neural network that integrates knowledge of the microscope's optical system inside our model. Experiments on synthetic benchmarks and real biological data show that both our new loss function and architecture surpass the state of the art at moderate and high emitter densities. Code is available at https://github.com/RSLLES/SHOT.