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

2025-12-11

AI Summary

Single-molecule localization microscopy (SMLM) suffers from degraded localization accuracy and excessive computational cost under high-density labeling. Existing deep learning approaches rely on non-differentiable non-maximum suppression (NMS), which often erroneously suppresses true emission events. To address this, we propose the first end-to-end differentiable learning framework for SMLM: (i) we formulate single-molecule localization as a set-matching problem; (ii) we design an optimal-transport-based differentiable loss function that fully eliminates NMS; and (iii) we integrate optical system priors into an iterative physics-informed neural network (PINN). Evaluated on both synthetic and real biological datasets, our method significantly outperforms state-of-the-art approaches—particularly in medium-to-high density regimes—achieving superior localization precision and recall. The framework enables faster, higher-fidelity super-resolution imaging of live cells.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Security and Privacy: Security and privacy of machine learning and AI applications

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.

Problem

• Enables end-to-end training for single-molecule localization microscopy
• Eliminates non-differentiable non-maximum suppression layers in deep learning
• Improves super-resolution imaging at moderate and high emitter densities

Innovation

• Optimal transport loss replaces non-differentiable NMS layers
• Iterative neural network integrates microscope optical system knowledge
• End-to-end training enables handling dense fluorophore emissions