PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimization of Deep Learning Model for Denoising Dose Profiles in Proton Therapy

Junaid Jawaid, Matteo Spezialetti, Gabriele Libardi, Giovanni Stilo, Filippo Mignosi

arXiv:2610.08843v1q-bio.QMeess.IV

TL;DR

Monte Carlo dose calculation is accurate but computationally expensive, motivating denoising of low-history simulations. This study extends a differentiable gamma-index loss to native 3D and applies structured pruning to a multiscale network. The GPR loss improves test-set dose agreement over MSE, while pruning reduces model size but not inference latency.

  • Problem

    Monte Carlo provides high-fidelity proton dose calculation but is computationally costly, and MSE training does not directly optimize the clinical Gamma Passing Rate.

    Source: §1 · p. 2

  • Method

    The study trains a multiscale network with a differentiable GPR loss extended to native 3D, then applies L1-norm structured filter pruning to the GPR-trained model.

    Source: §1 · p. 3

  • Results

    The mean GPR rises from 92.10% to 99.47% at 2mm/2% and by 24.54 percentage points at 1mm/1% with the differentiable GPR loss; pruning removes up to 20% of parameters with 99.33% versus 99.47% at 2mm/2% after fine-tuning.

    Source: §5 · p. 16

  • Takeaways & Limitations

    Pruning produces a smaller stored model and a lighter backbone for follow-on work, but does not reduce inference latency.

    Source: §5 · p. 16

    The network was trained and evaluated only on fixed-axis beams; arbitrary gantry angles require volume rotations that introduce trilinear interpolation errors.

    Source: §4.5 · p. 16

Abstract

from arXiv · show

Proton therapy offers a dosimetric advantage over conventional radiotherapy through the Bragg peak, enabling precise dose delivery to the tumour while sparing surrounding tissue. Monte Carlo (MC) simulation provides the most accurate dose calculation but is computationally expensive for clinical workflows. A promising compromise is to simulate with a reduced number of primary histories and recover a high-quality dose distribution through a deep learning denoising step. In this work, we take a 3D multiscale U-Net from recent literature as a case study and investigate two improvements critical for clinical deployment. First, we replace the standard MSE training objective with a differentiable $\gamma$-index-based loss, extended to native 3D operation, and demonstrate that directly optimising the clinical evaluation metric yields a mean Gamma Passing Rate of 99.47\% at 2mm/2\% and 95.67\% at 1mm/1\% on the test set, compared to 92.10\% and 71.13\% for MSE. Second, we apply L1-norm structured filter pruning to reduce the model's memory footprint, achieving a 13.3\% parameter reduction with negligible performance loss under partial pruning, and 20\% reduction under full-block pruning after fine-tuning.

1 Introduction

Monte Carlo dose calculation is accurate but computationally costly, while existing deep-learning denoising work used a pixel-wise objective that does not directly optimize clinical Gamma Passing Rate. This study tests a native-3D differentiable GPR loss and structured pruning on a multiscale network, alongside mixed-precision training.

1 Introduction

  • Monte Carlo offers high-fidelity proton dose calculation but can be computationally unviable, motivating low-count simulation followed by deep-learning denoising.Prior work used a multiscale 3D network to denoise 5 × 10^3-history dose to quality equivalent to a 10^6-history simulation, but trained with MSE rather than the clinical GPR metric.Source: §1 · p. 2
  • The research questions assess whether GPR-based training improves dose agreement and worst-case behavior, and whether structured pruning reduces parameters and inference latency.Pruning is examined at partial and full scope across the multiscale elaboration blocks.Source: §1 · p. 2
  • The study extends Martinot et al.'s differentiable GPR surrogate to native 3D operation and applies it to proton dose denoising.It also uses bfloat16 automatic mixed precision to reduce GPU memory consumption.Source: §1 · p. 3
  • 13.3% fewer parameters with GPR drops of only 0.06% at 2mm/2% and 0.64% at 1mm/1%; full-block pruning removes 20% and retains 99.33% at 2mm/2% after fine-tuning.Full-block pruning still incurs a 1.22% loss at the stricter 1mm/1% criterion.Source: §1 · p. 3

2 Related Work

Prior work establishes deep learning for proton dose denoising, identifies a mismatch between common voxel-wise losses and clinical GPR, and motivates structured pruning as a compression strategy.

2.1 Deep Learning for Proton Dose Denoising and Prediction

  • Spezialetti et al.'s multiscale 3D network denoises low-count Monte Carlo dose, while other U-Net-based models have predicted or refined radiotherapy dose distributions.The network processes 256×32×32 volumes and its memory requirement grows cubically with input edge length.Source: §2.1 · p. 3

2.2 Loss Functions and the Gamma Index in Dose Prediction

  • Martinot et al. made a gamma-index surrogate differentiable using a parametric sigmoid, but demonstrated it on 2D arrays with bilinear interpolation.This study extends that formulation to native 3D volumes using trilinear interpolation and compares it with MSE.Source: §2.2 · p. 4

2.3 Model Compression via Structured Pruning

  • L1 filter-weight norms provide an inexpensive criterion for structured pruning, which removes whole filters to produce smaller dense networks without sparse-runtime support.The study validates pruning on held-out batches and fine-tunes the pruned network to recover performance loss.Source: §2.3 · p. 4

3 Materials and Methods

The study compares MSE and native-3D differentiable GPR training on a shared network, then prunes the better-performing GPR model to measure compression–accuracy trade-offs.

3 Materials and Methods

  • The same architecture was trained from scratch under identical conditions with MSE and differentiable γ-index losses, evaluated on a held-out test set using pymedphys GPR at 2mm/2% and 1mm/1%, then the better-performing model underwent structured filter pruning.Source: §3 · p. 4

3.1 Dataset

  • 39 CT studies were split at patient level, transformed into beam-aligned 32×32×256 subvolumes, and used to generate training, validation, and test samples.Beams were simulated along the negative z-axis; the dataset totals 624 training, 208 validation, and 266 test samples.Source: §3.1 · p. 5
  • TOPAS simulations generated noisy inputs from 5 × 10^3 primary histories and denoising targets from 10^6 histories, with both dose volumes normalized to 10^5 particles.Beam energies ranged from 87 to 162 MeV and spot sizes were sampled from 2, 3, or 5 mm.Source: §3.1 · p. 5

3.2 The Network

  • The network takes CT and low-count Monte Carlo dose as two channels and predicts a single-channel denoised dose, using a multiscale U-Net-derived architecture.The PyTorch reimplementation was trained with MSE or 3D differentiable GPR loss, both using bfloat automatic mixed precision; pruning followed on the better-performing model.Source: §3.2 · p. 6

3.3 Gamma Loss Function

The loss approximates clinical Gamma Passing Rate with a differentiable gamma-index objective computed natively in 3D. It searches nearby predicted-dose locations and smooths the pass/fail decision for gradient-based training.

3.3 Gamma Loss Function

  • The local gamma index is the minimum generalized distance combining dose difference and spatial separation between a reference dose and a predicted dose nearby.The search neighborhood is bounded by the Distance-to-Agreement threshold.Source: §3.3 · p. 7
  • A parametric sigmoid smooths the gamma pass/fail criterion, approaching the hard criterion as its sharpness parameter increases.The loss minimizes Soft-GPR over voxels above a dose cutoff, with the sharpness parameter typically annealed during training.Source: §3.3 · pp. 7–8
  • The native-3D loss resamples predicted and reference volumes from 2 mm to 1 mm by trilinear interpolation, then searches a discrete DTA-bounded neighborhood in parallel.The elementwise minimum yields gamma values, with gradients propagating through the minimizing offset; values are clamped at 10.Source: Three-dimensional search and… · p. 8
  • Training uses DTA = 2 mm, Δ = 2%, and a 20% lower-dose cutoff, while test-set GPR is computed independently with pymedphys.The test evaluation therefore does not depend on the loss implementation's resampling and search procedure.Source: Three-dimensional search and… · p. 8

3.4 Mixed Precision Training

Automatic mixed precision reduces memory use during 3D dose-volume training while retaining full precision for numerically sensitive operations.

3.5 Mixed Precision Training

  • bfloat16 halves low-precision activation and weight memory relative to FP32 and yields 1.5–2× throughput improvements on Tensor Core hardware.Convolutions use bfloat16, while numerically sensitive reductions and FP32 master weights remain full precision.Source: §3.4 · p. 9
  • bfloat16 preserves FP32’s 8-bit exponent, so gradients do not underflow and loss scaling is unnecessary for stable convergence.The reduced mantissa precision is the difference from FP32 noted in the passage.Source: §3.4 · p. 9
  • The pruning target of 8 channels aligns with cuDNN’s multiple-of-8 requirement for Tensor Core acceleration, unlike the original 12-channel blocks.The channel alignment applies to FP16/BF16 mode.Source: §3.4 · p. 9

3.5 Post-Processing: Dilated Dose Mask

Because the GPR loss does not supervise low-dose and zero-dose voxels, it can leave spurious out-of-field predictions; a dilated mask from the MC input suppresses these artifacts without changing reported GPR.

3.5 Post-Processing: Dilated Dose Mask

  • GPR-only training leaves low-dose and zero-dose predictions unsupervised, allowing spurious non-zero activations outside the irradiated volume.The GPR loss evaluates only voxels whose reference dose exceeds its cutoff, whereas MSE penalizes all voxels uniformly.Source: §3.5 · p. 9
  • The mask thresholds the MC input at 10% of its maximum, dilates the region in 3D by 5 voxels, then zeros predicted dose outside it.The spherical structuring element is used to accommodate spatial spread and avoid clipping the Bragg peak penumbra.Source: §3.5 · p. 9
  • The mask does not change GPR because it leaves every reference voxel above the metric’s 20% cutoff unmodified.The 10%-thresholded input mask is dilated by 5 voxels and comfortably encloses those evaluated voxels; mask computation adds negligible inference overhead.Source: §3.5 · pp. 9–10

3.6 Pruning

Structured L1-norm pruning removes low-ranked filters from the MultiScaleElaboration blocks to create a smaller dense network while preserving residual-path dimensions. The study compares partial and full-block pruning, followed by fine-tuning.

3.6 Pruning

  • Filters are ranked by L1-norm and removed with corresponding normalization parameters and downstream input channels, yielding a physically smaller dense network.The 12 internal elaboration channels are reduced to 8, the nearest multiple-of-8 target compatible with the stated Tensor Core alignment constraint.Source: §3.6 · p. 10
  • Pairwise pruning slices upstream convolution and instance-normalisation outputs with downstream inputs to preserve dimensional consistency in residual paths.A validation dry-run on 20 batches precedes commitment, with pruning accepted only if the GPR drop remains below 50%.Source: §3.6 · p. 10
  • Partial pruning targets conv3, conv4, and conv4_1, while full pruning covers all seven convolutions in each elaboration block across encoder, bottleneck, and decoder.Both configurations are fine-tuned for up to 20 epochs with Adam at 10^-6; the best checkpoint is selected by validation GPR, with early stopping after five unimproved epochs.Source: §3.6 · p. 10

3.7 Evaluation Metrics

Evaluation uses the Gamma Passing Rate, which jointly accounts for dose difference and distance-to-agreement. Results are reported at routine and stricter criteria, including variability and worst-case values.

3.7 Evaluation Metrics

  • GPR jointly scores dose difference and distance-to-agreement, addressing limitations of dose-only or geometry-only comparisons.The joint measure is especially relevant near steep gradients such as the distal Bragg peak, where small shifts can produce large local dose differences.Source: §3.7 · p. 11
  • GPR is reported at 2 mm/2% and 1 mm/1%, both with a 20% lower-dose cutoff; the former is routine clinical QA and the latter is stricter.The cutoff excludes near-zero out-of-field voxels, where agreement is not clinically meaningful.Source: §3.7 · p. 11
  • The evaluation reports mean, standard deviation, minimum, and maximum GPR because worst-case performance matters clinically; batch-one forward-pass time is secondary for RQ3.Source: §3.7 · p. 11

3.8 Training Setup

The training pipeline normalises CT, dose, and network-output volumes before fitting the denoising models. MSE and differentiable GPR models use separate loss configurations, with the GPR loss’s sharpness increased during training.

Training Setup

  • bfloat16 automatic mixed precision was used on a single 20 GB GPU, enabling training on full 256 × 32 × 32 volumes without memory overflow.Source: §3.8 · p. 11

Data Normalisation

  • CT volumes were scaled to approximately [0, 1], and low- and high-history dose volumes were independently min-max normalised to [0, 1].Training samples were also augmented with four axial rotations: 0°, 90°, 180°, and 270°.Source: Data Normalisation · p. 11

Loss Configuration

  • The MSE model used standard MSE loss, while the GPR model used a 2 mm distance-to-agreement, 2% global dose threshold, and 20% cutoff.For the GPR model, the sharpness parameter β increased from 0.1 during warmup to 5.0, and the best checkpoint was selected by maximum validation GPR.Source: MSE Model · pp. 11–12, GPR Model · p. 12

4 Results and Discussion

The section compares dose-denoising losses and pruning configurations on a held-out test set of 266 samples, then reports inference timing, an indicative end-to-end cost, and limitations. The reported results show strong dose agreement alongside trade-offs in compression and evaluation stringency.

4.1 Baseline and Loss Function Comparison

  • For the representative case, GPR scored 99.63% / 97.28% versus MSE’s 91.64% / 69.96% at 2mm/2% and 1mm/1%, respectively.The figure shows GPR agreeing more closely with the reference in the high-dose region, while also showing spurious zero-dose activations.Source: Figure 2 · p. 13
  • The undenoised 5 × 10^3-primary-particle dose had mean GPR of 26.75% at 2mm/2% and 13.54% at 1mm/1%.Source: §4.1 · pp. 12, 14

4.2 Pruning Results

  • 13.3% partial pruning cut parameters from 3,957,258 to 3,432,042, with GPR shifting 99.47%→99.41% at 2mm/2% and 95.67%→95.03% at 1mm/1% after retraining.Source: §4.2 · p. 14
  • 20.0% full-block pruning cut storage from 15.10 MB to 12.07 MB; fine-tuning retained 99.33% at 2mm/2% but yielded 94.45% at 1mm/1%.Source: §4.2 · pp. 14–15, §4.3 · p. 15

4.3 Inference Time

  • The unpruned model’s 91.53 ms mean inference time was 95.42 ms after partial pruning and 96.37 ms after full-block pruning, with changes within run-to-run variability.The proposed bandwidth and kernel-launch explanation was not confirmed by kernel-level profiling.Source: §4.3 · p. 15

4.4 Clinical Applicability and Real-Time Feasibility

  • The estimated end-to-end cost was about 0.74 s per beam geometry, versus 130.30 s for the 10^6-primary-particle reference, or roughly two orders of magnitude faster.Source: §4.4 · p. 16

4.5 Limitations

  • Training and evaluation used beams along a fixed axis; arbitrary gantry angles remain untested, and alignment rotations may introduce interpolation errors near bone–tissue interfaces.The GPR loss also leaves below-cutoff voxels unsupervised, requiring a mask generated from the low-history input dose.Source: §4.5 · p. 16

5 Conclusion and Future Work

The study evaluates a differentiable clinical-metric loss, structured pruning, and mixed-precision training for proton dose denoising. Its findings support learned refinement as a potential fast dose engine, while identifying integration and generalization work still needed.

5 Conclusion and Future Work

  • 99.47% versus 92.10% mean GPR at 2mm/2%, with a 24.54-percentage-point gain at 1mm/1%, followed replacing MSE with native-3D differentiable GPR loss.The study reports a tighter worst-case tail as well.Source: §5 · p. 16
  • Up to 20% of parameters can be removed at essentially no dosimetric cost, but pruning did not reduce inference latency.After fine-tuning, the pruned model reached 99.33% versus 99.47% at 2mm/2%; its value is a smaller stored model and a lighter backbone.Source: §5 · p. 16
  • Well under a second per geometry versus roughly two minutes for the 10^6-primary reference positions learned refinement as a fast dose engine, subject to integration, commissioning, and quality assurance.The low-count simulation retains the physical model while the network restores statistical quality.Source: §5 · p. 17
  • Future work targets out-of-field artifacts with composite GPR and weighted-MSE supervision, and arbitrary gantry angles with a rotation-invariant architecture.The current fixed-axis approach requires interpolation-inducing volume rotations for arbitrary gantry angles.Source: §5 · p. 17

Related papers

Loading 2610.08843v1…