Scalable GPU-Accelerated Euler Characteristic Curves: Optimization and Differentiable Learning for PyTorch

AI summary

Overview

Research area: Topological deep learning and topological data analysis, specifically GPU-accelerated and differentiable computation of Euler Characteristic Curves (ECCs) for dense 2D/3D imaging data, integrated with PyTorch.

Technical level: Intermediate. Readers need some familiarity with CUDA memory hierarchies, PyTorch custom autograd functions, and basic topological concepts (sublevel sets, cubical complexes, Euler characteristic). The mathematical content is limited to sigmoid relaxation and analytic gradients.

Scope: The paper presents optimized full-sweep CUDA kernels for non-differentiable ECC computation on dense grids, plus a differentiable PyTorch layer that learns thresholds and a single direction for ECC via sigmoid relaxation, benchmarked on synthetic 2D and 3D grids against a prior chunk-based GPU implementation.

What This Paper Is About

The Euler Characteristic Curve summarizes the topology of an image or volume as a threshold sweeps over intensity values, and the Euler Characteristic Transform extends this over multiple directions to give a theoretically sufficient shape descriptor. Computing ECCs quickly on dense 2D and 3D grids, and making them differentiable so they can be trained end-to-end inside neural networks, are both unresolved problems. This paper addresses both: fast GPU kernels that process an entire dense grid in a single sweep, and a PyTorch layer that learns ECC thresholds and a direction through gradient descent.

Key Contributions

1. Efficient GPU kernels for ECC computation on modern hardware. Kernels designed for Ampere GPUs that respect data locality and the memory hierarchy, using 128-byte coalesced global memory access and per-block shared-memory histograms, achieving 16–2000× speedups over the prior chunk-based GPU implementation of Wang et al. (2023) on synthetic grids.

2. A differentiable single-direction ECC layer in PyTorch. A custom CUDA forward kernel implementing a DECT-style sigmoid relaxation, paired with a custom CUDA backward kernel using analytically derived gradients with respect to image values, thresholds, and the learnable direction.

3. End-to-end learning of geometric invariants for dense cubical complexes. Unlike DECT (Röell and Rieck, 2024), which targets graphs and point clouds, and unlike prior GPU-based ECC, which is non-differentiable, this work enables gradient-based discovery of task-relevant topological structure on 2D/3D images.

4. Released code (linked in the paper) to support adoption.

Main Findings

• 2D speedups peak in the mid-range grid sizes. Against the chunking baseline on 2D grids, speedups were 466× at 128² (7.7 ms baseline vs. 0.0165 ms new), 1379× at 256² (24 ms vs. 0.0174 ms), 1900× at 512² (39.68 ms vs. 0.021 ms), 2031× at 1024² (65 ms vs. 0.032 ms), 836× at 2048² (92 ms vs. 0.11 ms), 377× at 4096² (154 ms vs. 0.408 ms), and 166× at 8192² (267 ms vs. 1.6 ms).

• 3D speedups decrease as volume size grows. On 3D grids: 270× at 128³ (27 ms vs. 0.1 ms), 128× at 256³ (91 ms vs. 0.71 ms), 25× at 512³ (140 ms vs. 5.61 ms), and 16.5× at 1024³ (800 ms vs. 48.47 ms).

• The headline speedup range is 16–2000×. The paper reports consistent 16×–2000× speedups over the prior state-of-the-art chunking implementation and "stable performance across image sizes."

• Chunk-based overheads are the identified cause of the baseline's cost. The paper attributes prior latency to repeated kernel launches, inter-chunk synchronization, redundant boundary recomputation, and global atomic updates scaling as O(N_chunks × N_voxels).

• Full-sweep design eliminates per-chunk synchronization. Each voxel is assigned directly to a thread, each block keeps a private shared-memory histogram of H bins with __syncthreads() barriers, and atomics are issued to global memory only after all voxels in a block are processed, amortizing kernel launch cost over the entire dataset.

• Gradients exist for both thresholds and image values under moderate λ. The sigmoid relaxation recovers the hard indicator ECC as λ → ∞, and the analytic gradient with respect to the direction u is −α Σ c(p) σ′_λ(z_p(τ,u)) p, with ‖u‖₂ = 1 enforced (e.g., via reparametrization).

• Hardware and evaluation setting. All benchmarks used a single NVIDIA RTX A6000 GPU with 48GB of VRAM on the Ampere architecture. Benchmarks are on synthetic grids only; the paper does not report results on real imaging datasets, downstream task accuracy, or ablations over λ.

Methodology in Plain English

The Euler characteristic of a binary mask can be written as a sum over grid points, where each point contributes a coefficient c(p) ∈ {−1, 0, 1} determined by its local 3×3 (2D) or 3×3×3 (3D) binary neighborhood. That makes ECC computation a histogram-style accumulation problem over the whole image.

Earlier GPU work split grids into chunks small enough to fit in limited shared memory, processed each chunk, and merged partial results. That approach made sense for older hardware but pays for it with repeated launches, synchronization, boundary recomputation, and heavy global atomic traffic. This paper instead loads the whole dense grid in one sweep on modern hardware: one thread per voxel, coalesced 128-byte memory transactions, a private shared-memory histogram per thread block accumulated behind __syncthreads() barriers, and a single round of global atomics per block at the end. Boundary handling happens at the warp level so no redundant passes are needed.

For differentiability, the hard threshold — the indicator of X(p) ≤ τ — is replaced with a sigmoid of (τ − X(p)), scaled by a steepness parameter λ. This soft ECC is smooth, so gradients flow to the thresholds and to the image values. To support a learnable direction, the argument is extended to τ − X(p) − α⟨u, p⟩, with the direction constrained to unit norm. The forward pass is a custom CUDA kernel; the backward pass is a second custom CUDA kernel implementing the analytic derivatives given in the appendix.

Why This Matters

Impact on research. The paper makes the case that algorithm–hardware co-design — framed through Hooker's (2020) "hardware lottery" idea — is what turns theoretically motivated topological invariants into practical neural-network components. It fills a gap between DECT, which handles graphs and point clouds, and prior GPU ECC code, which is fast but not differentiable, by delivering both capabilities for dense cubical complexes. It also claims to enable real-time learning algorithms and the inclusion of topological inductive biases in large vision models.

Real-world applications (as discussed in the paper):

• Cosmic Microwave Background radiation analysis (referenced via Novikov et al., 1999), where differentiable ECC/ECT on dense grids may be useful because simplicity and efficiency are paramount.
• Medical imaging, specifically the Virtual Imaging Clinical Trials for Regulatory Evaluation (VICTRE) setting referenced via Badano et al. (2018), where geometric priors can complement learned features while remaining computationally tractable.
• Scientific imaging more broadly, where global geometric structure is informative but topological features have been too costly to adopt.
• Deep learning pipelines on 2D/3D images, where the released PyTorch layer allows topological descriptors to be trained as geometry-aware features rather than used as fixed preprocessing.

Industry relevance. A drop-in PyTorch module with custom CUDA forward and backward kernels, benchmarked on a widely available A6000-class GPU, lowers the integration cost for teams wanting topological features in production vision pipelines. The latency reductions — sub-millisecond ECC at 1024² in 2D — make real-time feature extraction plausible in a way chunk-based pipelines were not.

Future Directions

• Support for more modern GPU types beyond the Ampere architecture used in the benchmarks.
• Batched execution, which the abstract lists alongside multi-GPU extensions as needed to broaden adoption.
• Multi-GPU data and model parallelism pipelines for scaling to larger workloads.
• A robust PyTorch API release and exploration of application domains where fast, learnable ECC/ECT features benefit medical and scientific imaging.

Open questions the paper leaves unaddressed include how the learnable layer performs on real imaging tasks rather than synthetic grids, how sensitive results are to the relaxation steepness λ, and how the single learnable direction compares against the full multi-direction ECT once batching is available.

Target Audience

Graduate students and researchers in topological data analysis, topological deep learning, and geometric deep learning who want differentiable topological descriptors for dense imaging data. Also relevant to practitioners building PyTorch pipelines for medical or scientific imaging who need fast GPU feature extraction, and to systems-oriented researchers interested in the algorithm–hardware co-design argument for replacing chunk-based kernels with full-sweep designs on modern accelerators. Readers without CUDA or PyTorch custom-operator background can follow the motivation and benchmark results but may find the kernel design and gradient derivations hard going.

Authors’ abstract

Topological features capture global geometric structure in imaging data, but practical adoption in deep learning requires both computational efficiency and differentiability. We present optimized GPU kernels for the Euler Characteristic Curve (ECC) computation achieving 16-2000Ö speedups over prior GPU implementations on synthetic grids, and introduce a differentiable PyTorch layer enabling end-to-end learning. Our CUDA kernels, optimized for Ampere GPUs use 128B-coalesced access and hierarchical shared-memory accumulation. Our PyTorch layer learns thresholds in a single direction via a Differentiable Euler Characteristic Transform-style sigmoid relaxation. We discuss downstream relevance, including applications highlighted by prior ECC work, and outline batching/multi-GPU extensions to broaden adoption.