PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Spectral Allocation: A Fractional Diffusion Framework for Adaptive Volumetric Segmentation

Yi-Hui Shen, Tie-Qiang Li

arXiv:2609.27217v1cs.CVcs.AI

TL;DR

The paper asks how much spectral mixing each stage of a lightweight 3D segmentation network needs. It derives and learns identity-containing FHEAT gates, then evaluates them in FHEAT-Seg. Under scarce labels, the method selects a single surviving diffusion stage and achieves strong segmentation results across three benchmarks.

  • Problem

    The paper asks where global spectral mixing should be spent in a lightweight 3D segmentation backbone.

    Source: §VI · p. 15

  • Method

    FHEAT-Seg equips stages with learnable identity-containing fractional diffusion gates and pairs the FHEAT operator with a KAN3D channel mixer.

    Source: §I · p. 2, C. FHEAT Block and KAN3D Mixing · p. 5

  • Results

    FHEAT-Seg achieves Dice of 90.47%, 78.79%, and 81.90%, surpassing five semi-supervised methods and Light-UNETR across three benchmarks; its 0.90G FLOPs are 79% fewer than Light-UNETR.

    Source: §I · p. 2, §VI · p. 15

  • Takeaways & Limitations

    Under scarce labels, training drives diffusion strength to zero in all but one of eight stages, consistently selecting the decoder block that produces the semi-supervised attention map.

    Source: §I · p. 2

    The fractional order is bounded to (0.3,0.9), so saturation at 0.9 does not establish a physically preferred order, and an unbounded parameterization remains untested.

    Source: §V · p. 15

Abstract

from arXiv · show

We address adaptive computation in 3D medical image segmentation: instead of designing another backbone, we ask how much spectral mixing each network stage needs and let optimization answer. We derive FHEAT, a two-parameter operator family, from the discrete cosine transform (DCT) solution of a fractional heat equation. A fractional order alpha and a diffusion strength D govern the operator, and at D=0 it is exactly the identity. Reparametrized by the semigroup time tau = D*alpha, same-resolution instances compose exactly, so any distribution of diffusion across same-resolution stages amounts to a single Sobolev-type regularizer of learned strength. This identity limit lets the optimizer of each layer, not the designer, decide whether global mixing is needed and how sharp it should be. We instantiate FHEAT in a lightweight U-shaped architecture (Light-UNETR) paired with a Kolmogorov-Arnold mixer (KAN3D) with adaptive rational activations, yielding FHEAT-Seg. At 5% to 20% label rates on three public benchmarks, training produces gradient-driven spectral sparsification: seven of the eight stage-level operators drive D to zero, and the survivor saturates at the sharpest low-pass (alpha ~ 0.9) in the decoder layer feeding the semi-supervised attention map. The retired layers become exact identity shortcuts at inference, cutting FLOPs from 4.29G to 0.90G (a 79% drop) at 0.975M parameters. Under a standard semi-supervised protocol, FHEAT-Seg reaches Dice scores of 90.47% (left atrium), 78.79% (Pancreas-CT), and 81.90% (BraTS 2019), ahead of five semi-supervised methods and the Light-UNETR baseline. The large variant also surpasses Light-UNETR-L under full supervision (Dice 93.09%, 85.11%, and 87.19%) with 2.851M parameters and 55.75G FLOPs. These results suggest that the allocation of spectral computation is a learnable property of optimization dynamics, not a manual design commitment.

I. INTRODUCTION

The paper makes per-stage global spectral mixing learnable through an operator with an exact identity limit, then evaluates its allocation in a lightweight segmentation model. Under scarce labels, training concentrates diffusion in one decoder stage while the resulting model outperforms comparison methods across three benchmarks at lower inference cost.

I. INTRODUCTION

  • The design targets fixed per-stage global-mixing budgets in lightweight models and the difficulty of separating backbone gains from semi-supervised protocol gains.The paper keeps the CSE protocol unchanged to isolate the backbone contribution.Source: §I · pp. 1–2
  • FHEAT embeds a fractional heat diffusion operator in a residual gated module, with learnable α controlling high-frequency attenuation and D scaling smoothing from identity to low-pass aggregation.The operator is derived from the fractional heat equation using the DCT-II basis under Neumann boundary conditions.Source: §I · pp. 1–2
  • Because spectral diffusion is a linear diagonal transform, FHEAT-Seg pairs it with KAN3D, whose learnable rational activations adapt channel mixing without widening the network.The mixer recombines channels rather than adding further spatial aggregation.Source: §I · p. 2
  • Across five tested settings, training drives D to zero in seven of eight stage operators, leaving the final decoder block that produces the SSL attention map as the survivor.The surviving operator reaches α ≈0.9; a controlled experiment allowing learnable D only at this block matches or slightly exceeds the fully learnable model.Source: §I · p. 2
  • The retired stages become exact identity shortcuts, yielding an inference graph with 0.975M parameters and 0.90G FLOPs, 79% fewer FLOPs than Light-UNETR.The identity mapping is reachable at zero diffusion strength.Source: §I · p. 2
  • FHEAT-Seg achieves Dice scores of 90.47% on LA, 78.79% on Pancreas-CT, and 81.90% on BraTS 2019, surpassing five semi-supervised methods and Light-UNETR under shared CSE evaluation.The large variant also outperforms Light-UNETR-L under full supervision on all three benchmarks at lower parameter and FLOP counts.Source: §I · p. 2
  • Same-resolution, same-order gates compose exactly by adding diffusion times, but the paper reports that this redundancy alone does not explain the observed concentration of diffusion.A fractional Dirichlet pairing of incoming features and gradients supplies a first-order criterion for whether a stage is useful.Source: §I · p. 2

II. RELATED WORK

Prior work improves 3D segmentation efficiency through compact blocks and lightweight attention, while semi-supervised methods exploit unlabeled scans. The paper positions learned, identity-containing spectral gates as a distinct way to allocate global mixing by stage, and holds the CSE protocol fixed to compare backbones.

A. Lightweight 3D Medical Image Segmentations

  • Lightweight 3D segmenters reduce block costs, but retain hand-crafted global-mixing allocation; this work instead learns each stage’s mixing strength so training can retain spectral computation only where needed.Light-UNETR is described as combining local grouped convolutions with low- and high-frequency dual-branch attention.Source: A. Lightweight 3D Medical Image… · p. 3

B. Semi-Supervised 3D Medical Image Segmentation

  • The study keeps CSE unchanged and compares five competing SSL methods under a shared-backbone setting to separate protocol effects from backbone differences.CSE combines attention-guided replacement with spatial masking consistency.Source: B. Semi-Supervised 3D Medical Image… · p. 3

C. Frequency-Domain and Learnable-Activation Operators

  • Unlike fully parameterized frequency filters and post-hoc gates, FHEAT uses a two-parameter fractional heat kernel with an exact identity limit, allowing a stage to opt out of mixing and compile to a shortcut.The paper distinguishes this from SlimFormer-3D’s post-training activation-similarity gates and AMBER-AFNO’s learned spectral filters at every stage.Source: C. Frequency-Domain and… · p. 3
  • FHEAT also differs from architecture search, pruning, stochastic depth, routing, and classical fractional diffusion by making per-stage mixing a continuous learned outcome with an exactly deployable no-mixing option.The paper learns the fractional order rather than hand-tuning it and embeds the operator in a trainable segmentation backbone.Source: C. Frequency-Domain and… · p. 3

III. METHOD

FHEAT-Seg replaces stage-specific fixed mixing recipes with learnable spectral gates in a compact U-shaped network. Its gates derive from fractional heat diffusion in the DCT-II basis and can reduce to identity mappings.

A. Overview

  • FHEAT-Seg instantiates learnable spectral gates in a compact U-shaped encoder-decoder with overlapping patch embedding, skip connections, and a fixed CSE testbed for semi-supervised evaluation.The architecture follows Light-UNETR design principles to provide an empirical baseline.Source: A. Overview · p. 4

B. Fractional-Order Spectral Diffusion

  • The Fractional Heat Conduction Operator (FHCO) uses the DCT-II basis under Neumann boundary conditions to apply a closed-form fractional-heat diffusion kernel to spectral coefficients.The cosine basis follows from the boundary conditions on the bounded voxel grid.Source: B. Fractional-Order Spectral Diffusion · pp. 4–5
  • Each stage shares two learned scalars: fractional order α shapes frequency attenuation, while diffusion strength D controls the amount of smoothing.The reparameterized diffusion time τ = D*α is the quantity stored and initialized by the implementation.Source: B. Fractional-Order Spectral Diffusion · p. 5
  • At D = 0, the spectral transform is exactly identity and can be bypassed; for α in (0,2], it is well-posed smoothing with stage-wise transform cost linear in channel count.The separable matrix-product implementation costs O(CDHW(D+H+W)); an FFT-based fast DCT can reduce this to O(CDHW log(DHW)).Source: B. Fractional-Order Spectral Diffusion · p. 5

C. FHEAT Block and KAN3D Mixing

FHEAT and KAN3D divide spatial and channel-mixing roles: FHEAT supplies frequency-selective global context, while KAN3D performs adaptive nonlinear channel recombination.

C. FHEAT Block and KAN3D Mixing

  • FHEAT is a linear, frequency-diagonal transform, so each block pairs it with a channel mixer to provide nonlinear feature recombination.The mixer uses learnable rational activations rather than a fixed activation.Source: C. FHEAT Block and KAN3D Mixing · p. 5
  • KAN3D uses learnable rational activations between 1 × 1 × 1 convolutions, adapting stage nonlinearities to feature statistics without widening the network.The rational functions are ratios of learnable polynomials initialized to approximate GELU.Source: C. FHEAT Block and KAN3D Mixing · p. 5
  • FHEAT provides global frequency-selective context, while KAN3D supplies adaptive pointwise recombination at negligible spatial cost.The final decoder block also exposes an attention map consumed by the semi-supervised protocol.Source: C. FHEAT Block and KAN3D Mixing · p. 5

D. Integration with the CSE Protocol

The method uses the CSE semi-supervised protocol unchanged as a controlled testbed, keeping the backbone as the varied factor. Its attention map is a read-only signal for region replacement.

D. Integration with the CSE Protocol

  • CSE combines Attention-Guided Replacement, which mixes labeled and unlabeled image regions with ground-truth and pseudo labels, and Spatial Masking Consistency between strongly masked and weakly augmented views.The protocol is held fixed rather than treated as part of the proposed method.Source: D. Integration with the CSE Protocol · p. 6
  • The AGR attention map comes from Block_Final and is stop-gradient, so its mixing loss cannot reshape the backbone features that generate it.Segmentation losses still train the backbone through segmentation outputs.Source: D. Integration with the CSE Protocol · p. 6

E. Inference Simplification and Efficiency Accounting

Training-selected zero-strength gates become exact identity shortcuts, simplifying inference. The optimization analysis distinguishes this exact bypass from the conditions that select which stages remain active.

E. Inference Simplification and Efficiency Accounting

  • 1.20G to 0.90G FLOPs for FHEAT-Seg and 57.77G to 55.75G for FHEAT-Seg-L follow from replacing zero-strength stages’ FHCO paths with identity shortcuts.Surrounding convolutions, projections, and gating remain; the bypass differs from identity by at most 1.5 × 10^-6 in fp32.Source: E. Inference Simplification and… · p. 6
  • At zero diffusion strength, the FHCO is exactly identity, making the stage’s inference replacement exact apart from floating-point rounding.The FHEAT block then collapses to its affine bypass around the transform.Source: F. Optimization Geometry of the… · p. 6

F. Optimization Geometry of the Identity Limit

  • The semigroup identity collapses only same-grid, same-order gates; different U-Net stages have distinct grids and learned orders, so their diffusion times cannot be combined across stages.Within one stage, blocks share the same parameters and can be collapsed exactly.Source: F. Optimization Geometry of the… · p. 6
  • Weight decay does not itself sparsify interchangeable gates: eight redundant gates remain active and near-equal, so observed sparsity requires stages to differ.The controlled symmetric case reports D = 0.248 ± 0.011 across gates.Source: F. Optimization Geometry of the… · p. 7, Figure 3 · p. 7
  • A positive fractional Dirichlet pairing means increasing a gate from zero strictly decreases loss; a nonpositive pairing makes zero a one-sided local minimum.To first order, the number of surviving stages is the number with positive pairing.Source: F. Optimization Geometry of the… · p. 7
  • With the ReLU parameterization and decoupled weight decay, a negative pairing can push a gate below zero and make its zero diffusion strength absorbing.Retirement is permanent during training absent an external perturbation; the authors note this locking is specific to this parameterization.Source: F. Optimization Geometry of the… · p. 7
  • Increasing diffusion strength acts, to leading order, like strengthening a stage-wise fractional Sobolev penalty on its features.At zero strength, that penalty is absent.Source: F. Optimization Geometry of the… · p. 7

A. Experimental Setup

The experiments test three public 3D segmentation benchmarks under scarce-label settings, using a shared semi-supervised protocol and Light-UNETR backbone for controlled comparisons.

Datasets

  • Three benchmarks cover LA (80 training, 20 test scans), Pancreas-CT (62 training, 20 test scans), and BraTS 2019 (250 training, 25 validation, 60 test scans).Label counts range from 4 and 8 on LA, 6 and 12 on Pancreas-CT, and 25 on BraTS 2019.Source: A. Experimental Setup · pp. 7–8

Metrics

Training protocol

  • Semi-supervised models use the CSE protocol for 16,000 iterations, and test scores are computed on test sets separate from validation data.Training uses two labeled and two unlabeled volumes per batch, AdamW, and cosine learning-rate decay.Source: A. Experimental Setup · p. 8

Statistical testing

  • Unless noted otherwise, improvements over the strongest baseline are tested with paired Wilcoxon signed-rank tests across test cases.Source: A. Experimental Setup · p. 8

Baselines

  • Comparisons use five semi-supervised methods trained with the same Light-UNETR backbone and official splits, isolating the backbone contribution.Reported fully supervised comparisons follow the same protocol.Source: A. Experimental Setup · p. 8

B. Learned Spectral Allocation

Across scarce-label semi-supervised settings, training retains one low-pass stage and retires the other seven as identities. The section examines the allocation's dynamics, cross-architecture behavior, and relation to attention and segmentation performance.

B. Learned Spectral Allocation

  • Across five semi-supervised settings and 15 seeded runs, exactly one of eight operators retains D > 0: the final decoder block, with α ≈ 0.90, while the other seven become identities.The same block produces the attention map consumed by AGR; under fully supervised training, sparsification persists but the survivor is not consistently localized there.Source: B. Learned Spectral Allocation · p. 8, Table I

THE UNIQUE STAGE WITH D > 0; ALL SEVEN OTHER STAGES

  • In the LA 5% training trace, seven stages reach D = 0 by the first recorded checkpoint, while the surviving stage's D grows and α approaches 0.9.The learned response is a strong low-pass at the survivor and identity elsewhere; projected dynamics make retirement absorbing, with AdamW decay preserving the negative sign.Source: THE UNIQUE STAGE WITH D > 0; ALL SEVEN… · pp. 8–9
  • A SegFormer3D port reproduces the pattern: seven stages reach D = 0 and dec4 is the sole survivor, with test Dice 0.9002 versus 0.8817 for the plain control.The survivor's location differs from FHEAT-Seg, while the unique-survivor pattern recurs across architectures.Source: THE UNIQUE STAGE WITH D > 0; ALL SEVEN… · p. 8, Figure 4 · p. 9
  • FHEAT-Seg bottleneck residuals place 33.3% of energy at high normalized radii (> 0.6), versus 0.2% for Light-UNETR, alongside stronger target-region attention concentration.Attention energy within the target region is 44.8% versus 3.4% on LA, 18.4% versus 11.6% on Pancreas-CT, and 22.5% versus 1.3% on BraTS 2019; the authors describe this as a byproduct, not a cause.Source: THE UNIQUE STAGE WITH D > 0; ALL SEVEN… · p. 9, Figure 6 · p. 11
  • FHEAT-Seg ranks first on Dice and Jaccard across the reported semi-supervised settings, with the best 95HD and ASD except Pancreas-CT-12 95HD (8.35 versus 8.69).At the lower label setting, it records 90.47% Dice on LA and 78.79% on Pancreas-CT, surpassing the CSE-trained Light-UNETR backbone and five SSL methods.Source: THE UNIQUE STAGE WITH D > 0; ALL SEVEN… · p. 9, Table II · p. 9, Table III · p. 11

D. Fully Supervised Comparison and Efficiency

Under full supervision, FHEAT-Seg and its large variant outperform their Light-UNETR counterparts across the reported benchmarks while using fewer parameters and FLOPs.

Fully supervised comparison

  • FHEAT-Seg improves over Light-UNETR Dice on LA (92.05 vs. 91.58), Pancreas-CT (83.30 vs. 80.45), and BraTS 2019 (84.28 vs. 79.93), with 27% fewer parameters and 79% fewer FLOPs.The comparison reports 0.975M / 0.90G for FHEAT-Seg versus 1.34M / 4.29G for Light-UNETR.Source: D. Fully Supervised Comparison and… · p. 11, Table IV · p. 12
  • FHEAT-Seg-L exceeds Light-UNETR-L Dice on LA (93.09 vs. 92.22), Pancreas-CT (85.11 vs. 84.30), and BraTS 2019 (87.19 vs. 84.32), at 2.851M / 55.75G versus 3.93M / 59.45G.The two FHEAT-Seg variants together rank first in twelve of Table IV's fourteen metric and efficiency columns.Source: D. Fully Supervised Comparison and… · p. 11, Table IV · p. 12

E. Ablation Studies

The ablations isolate measurable gains from KAN3D and show that one learnable spectral site can suffice; qualitative results show preserved thin anatomy. The discussion interprets the surviving operator as a bounded, stage-specific low-pass choice and states the evidence's scope.

E. Ablation Studies

  • 90.60% Dice with learnable diffusion only at Block_Final exceeds the fully learnable model's 90.47%, while freezing all diffusion strengths gives 90.42%.Replacing an equal-capacity GELU mixer with KAN3D raises Dice from 90.15% to 90.47% (+0.32 Dice, +0.53 Jac) at nearly identical parameter counts.Source: E. Ablation Studies · p. 13

F. Qualitative Results

  • FHEAT-Seg preserves papillary muscles on LA and the pancreatic tail on Pancreas-CT, structures competing methods blur or miss.The comparisons show three cases per dataset under LA 5% and Pancreas-CT 10% label settings, against five SSL methods, Light-UNETR, and ground truth.Source: F. Qualitative Results · p. 13

V. DISCUSSION

  • Across semi-supervised models, the surviving operator reaches the sharpest permitted low-pass near alpha = 0.9, while retired stages' alpha values of 0.60–0.75 are unused.The discussion identifies this saturation as a bounded optimum, not evidence that 0.9 is intrinsically special; an unbounded parameterization remains untested.Source: §V · p. 14
  • The surviving site is Block_Final, a high-resolution decoder stage before the segmentation head that also produces the attention map used for region selection.A Block_Final-only control is sufficient, while the all-frozen control indicates retired stages add negligible mutual information once surrounding convolutional paths are present.Source: §V · p. 14
  • The authors report two orders of magnitude more high-frequency detail from FHEAT blocks than the reference backbone's low-pass-only attention, alongside preserved thin structures and improved boundary distances.The cited boundary metrics are 95HD and ASD.Source: §V · pp. 14–15
  • The scope is bounded by alpha in (0.3, 0.9), a ReLU-and-decoupled-weight-decay retirement mechanism, and no measured pairing value on a saved medical checkpoint.The semigroup-collapse result applies only to gates sharing resolution and fractional order; it does not make the eight stage-level gates interchangeable.Source: §V · p. 15

VI. CONCLUSION

The conclusion presents an operator theory that predicts where spectral mixing is useful and how gradient descent can retire unnecessary stages. FHEAT-Seg applies this allocation under scarce labels, yielding a sparse, lower-cost model with strong benchmark performance.

VI. CONCLUSION

  • The theory uses a fractional Dirichlet pairing to predict first-order stage retention, while ReLU parameterization and weight decay convert an unfavorable sign into an exact deployable zero.Table VI agrees with the predicted ranking, and controlled simulation confirms the sign-predicts-survival mechanism on synthetic ground truth.Source: §VI · p. 15
  • 0.975M parameters and 0.90G FLOPs (79% fewer FLOPs than Light-UNETR) accompany the best compared semi-supervised segmentation on LA, Pancreas-CT, and BraTS 2019.Under full supervision, FHEAT-Seg also outperforms Light-UNETR-L at lower cost.Source: §VI · p. 15
  • Training drives diffusion strength to zero in seven of eight stages, consistently leaving the decoder block that produces the semi-supervised attention map.Reactivating retired stages degrades validation Dice, and the absorbing zero remains stable under learning-rate spikes.Source: §VI · p. 15
  • On a 3D SegFormer backbone at LA 5% labels, FHEAT gates achieve Dice 0.9002 versus 0.8817 for the plain control, with seven of eight stages reaching exactly D = 0.This reproduces spectral sparsification beyond the lightweight U-shaped backbone.Source: §VI · p. 15

Related papers

Loading 2609.27217v1…