Cross-Timestep: 3D Diffusion Model with Trans-temporal Memory LSTM and Adaptive Priori Decoding Strategy for Medical Segmentation
Authors: Shangqian Wu
Siyuan Shen
Yahan Li
Zhijian Huang
Ziyu Fan
Yuanpeng Zhang
YI WANG
Lei Deng
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Models; Medical Image Segmentation; LSTM
Primary area: applications to computer vision, audio, language, and other modalities

Abstract: Diffusion models have recently demonstrated significant robustness in medical image segmentation, effectively accommodating variations across different imaging styles. However, their applications remain limited due to: (i) current successes being primarily confined to 2D segmentation tasks—we observe that diffusion models tend to collapse at the early stage when applied to 3D medical tasks; and (ii) the inherently isolated iteration along timesteps during training and inference. To tackle these limitations, we propose a novel framework named Cross-Timestep, which incorporates two key innovations: an Adaptive Priori Decoding Strategy (APDS) and a trans-temporal memory LSTM (tLSTM) mechanism. (i) The APDS provides prior guidance during the diffusion process by employing a Priori Decoder(PD) that focuses solely on the conditional branch, successfully stabilizing the reverse diffusion process. (ii) The tLSTM integrates convolution and linear layers into the LSTM gating structure, and enhances the memory cell mechanism to retain temporal state, explicitly preserving and propagating continuous temporal states across timesteps. Experimental results demonstrate that Cross-Timestep performs favorably on heterogeneous 3D medical datasets. Three experiments further analyze the collapse phenomenon in 3D medical diffusion models and validate that APDS effectively prevents initial-stage collapse without excessively constraining the model, while tLSTM facilitates the performance and scalability of diffusion models.

Published as a conference paper at ICLR 2026
CROSS-TIMESTEP: 3D DIFFUSION MODEL WITH TRANS-TEMPORAL MEMORY LSTM AND ADAPTIVE PRIORI DECODING STRATEGY FOR MEDICAL SEG- MENTATION
Shangqian Wu', Siyuan Shen', Yahan Li', Zhijian Huang', Ziyu Fan', Yuanpeng Zhang', Yi Wang?., Lei Deng
1 School of Computer Science and Engineering, Central South University, ChangSha, China PPeingg University People's Hospital
(wushanggian, leiden Leideng) (ac.g))csssueee eus..edu.co "Corresponding author
ABSTRACT
Diffusion models have recently demonstrated significant robustness in medical image segmentation, effectively accommodating variations across different imag- ing styles. However, their applications remain limited due to: (i) current suc- cesses being primarily confined to 2D segmentation tasks-we observe that diffu- sion models tend to collapse the early stage when applied 3D medical tasks; and (ii) the inherently isolated iteration along timesteps during training and infer- ence. To tackle these limitations, we propose a novel framework named Cross- Timestep, which incorporates two key innovations: an Adaptive Priori Decoding Strategy (APDS) and trans-temporal memory LSTM (ILSTM) mechanism. (i) The APDS provides prior guidance during the diffusion process by employing Priori De)oder(PD) that focuses solely on the conditional branch, successfully stabilizing the reverse diffusson process. (ii) The tLSTM integrates convolution and linear layers into the LSTM gating structure, and enhances the memory cell mechanism to retain temporal state, explicitly preserving and propagating con- tinuous temporal states across timesteps. Experimental results demonstrate that Cross-Timestep performs favorably on heterogeneous 3D medical datasets. Three experiments further analyze the collapse phenomenon in 3D medical diffusion models and validate that APDS effectively prevents initial-stage collapse without excessively constraining the model, while tLSTM facilitates the performance and scalability of diffusion models.
INTRODUCTION
Medical image segmentation-thur the precise delineation of anatomical structures -is essential for di- agnosis and treatment planning (Zhang al. 2025b Luo et al. 202 2022). Building robust models usually requires aggregating images across hospitals and scanners, which introduces large appearance variability and degrades the performance conventional networks (Zaman al. 2024 Luo et al. 2021b). Denoising Diffusion Probabilistic Models (DDPMs) are an attractive alternative because their coarse-to- -ine denoising process encourages recovery of global structure before fine details, suggesting greater robustness to style variation (Salsi et al. 2025 Zhang et al. 2025a Ding
al. 2024, Wu et al. 2025) Shuai et al. 2024).
Despite this promise, applying diffusion models to 3D volumetric segmentation remains challeng- ing. We identify a failure mode we call "initial-stage collapse": when reverse sampling is started from very high-noise timesteps (near pure Gaussian noise), models adapted from 2D often produce incoherent outputs and fail to recover the target structure (see Fig. []. additional details and visual- izations are provided in Appendix A. The difficulty stems from the enlarged manifold of 3D volumes and the vanishingly weak structural cues present at extreme noise levels, while standard diffusion samplers lack mechanisms to accumulate evidence across timesteps.

---

Published as a conference paper at ICLR 2026
To address these issues, propose Cross-Timestep. 3D diffusion framework that (i) uses an Adaptive Priori Decoding Strategy (APDS) to supply winehheehtted structural priors from the con- ditional image, and (ii) equips the denoiser with an explicit cross-timestep memory via trans- temporal memory LSTM (LLSTM). APDS provides a coarse, image-driven scaffold that is strongest early denoising steps and decays the learned denoiser gains confidence; tLSTM persistently car- ries low-frequency structure, residual statistics, and uncertainty-awer timesteps refine rather than re-discover structure. We implement two concrete ILSTM instantiations (Conv-tLSTY, and Linear-tGRU) and two extensions (SC-tLSTM and FFT-LLSMM to demonstrate the approach's flexibility. Experiments on multi-center 3D datasets show Cross-Timestt Timestep prevents initial-stage stage collapse and improves segmentation robustness under severe domain shift. In summary, the contributions of this study are as follows:
saliency acToss steps so later
We introduce Cross-Timestee Timestep, pioneering the use of diffusion models for 3D medical seg- mentation tasks. This framework establishes a foundational strategy for diffusion-bart approaches, offering a possible universal solution for future research and achieving robust performance heterogeneous datasets.
We propose the APDS to specifically address 'initial-stage collapse', providing stable struc- tural guidance during high-noise stages without interfering with later refinement stages.
We develop the tLSTM module, introducing a recurrent controller that carries cross- timestep state, turning local, myopic denoising into temporally coherent trajectory.
Implement tLSTM as modular controller (Conv-tLSTM, Linear-tGRU) and demonstrate extensions (SC-LLSTM, FFT-LSTMM) that demonstrate substantial scalability and potential for further enhancements of diffusion models via these units.
2 RELATED WORK
Diffusion models have been adapted for discrimina- tive tasks including image segmentation, with many 2D methods integrating conditional guidance, uncer- tainty modeling, or transformer-hase based context mod- ules to improve boundaries and global consistency (Zaman et al. 2026 & Chung 2024 Wu et al.) 2024 Tang et al. 2024). Work 3D volumetric
Sturt 1000
Z1000
Xster
X
Diftusion samping
SStrtt 600
diftusion has focused mainly on generative synthesis and on engineering memory/efficiency solutions to handle large 3D tensors (Nie et al., 2025 Liu et al. 2025 Wang al., 2024a Chen et al. 2024), while some recent methods (e.g.. Diff-UNet) fuse multi-
Xsm
samping
Start1 = 1000
step outputs to reduce variance (Xing al., 2025
X600
These prior approaches either (i) concentrate on gen- eration and computational feasibility in 3D, or (ii) apply post-hoe fusion of independent per-step pre- dictions. In contrast, Cross-Timestep targets two complementary problems simultaneously: reliable initialization from high-noise timesteps (via APDS) and explicit temporal features accumulation during sampling (via LLSTM). By integrating structured pri- with stateful denoiser, our framework addresses the root causes of initial-stage collapse rather than relying on post-hoc output fusion.
AP23
Figure 1: "Initial-stage collapse". For 3D medical data, the diffusion model will crash when sampling starts from the high-noise stage (equivalent random noise), but it can correctly sample from the middle and low time steps. Introducing APDS enables cor- rect sampling starting from random noise.
METHODOLOGY
In this section, introduce proposed Cross-Timeste method for stable 3D medical image segmentation using diffusion models (Fig. 2 We aim to make the reverse diffusion process both reliably initiable under extreme noise (via APDS) and temporally coherent across steps (via LLSMM.
2

---

Published as a conference paper at ICLR 2026
Cross-Timestep Diffusion Model
connect
SC
3C
SC-ALSTH
sc
SC-LLSTM
SC-LLSM
SC
meme
SC-LLSTM
SC-LLSTM
Uhanppng
UPaamolin
UPamsshee
PD
Skip connect
(d) Reverse Addition
(b) SC
1-ta.
CN
ConV-tLSTM
COOA
AN5000
(c) FFT-LLSTM
Liner-tGRU
ULSTM
X
Eoner STM
LIwn-GGUU
Prise Deccddrr
Tingret
Figure 2: (a) The framework the Cross- cross-Timestep, we propose APDS and ILSTM to construct stable diffusion architecture.
(b) Time-weighted control RA, using the prior mask obtained from the PD to guide the main branch. (c) Detailed design Conv-tLS tLSTM, using convolution to improve the gating mechanism of LSTM and enhance its memory cell to remember the temporal state. (d) Detailed design Linear-tGRUU, using Linear combined with GRU to reduce resource requirements, compared with Conv-IISTM (e) Structure of SC-1LSTM, improving the traditional SC attention adapt to the diffusion model. (f) Structure FFT-LLSMM, transforming the time domain into the frequency domain for denoising.
APDS provides structural priors from the clean conditional image and injects them with a decaying time weight. tLSTM explicitly carries cross cross-timesten timestep memory evidence accumulates instead of being re-dixcovered.
3.1 ADAPTIVE PRIORI DECODING STRATEGY
As illustrated in Fig. b. the APDS mechanism operates exclusively on the conditional branch of the network. It takes the conditional image X., which is embedded with the current timestep information The deepest features from this are passed through a series of bottleneck blocks be- fore entering a Prior Decoder (PD). The PD fuses these bottleneck features with multi-scale skip connections from the conditional encoder's path. This process yields preliminary segmentation mask, F prioo.. which serves robust approximation of the target structure z According to the Reverse Addition(RA))) (Fig. 2)) concept, pprior is reversed guide generating the refined feature map, Frefined:
the main branch F Fmain: thereby
Frefined Fmain
(FFrrorr)
Although there was no explicitly shown embedded time t during this process, X. already incorpo- rates the time embedding information, as it allows the APDS to be aware of the current diffusion

---

Published as a conference paper at ICLR 2026
stage, enabling it provide more global, structural guidance during high-noise timesteps and more detailed, refined guidance as the noise level decreases.
Further still, to prevent this strong prior from overly interfering with the main model's predictions, especially during the later, low-noise stages of diffusion, we introduce a time-weighted mechanism. The RA module is integrated at each upsampling stage of the main branch, and fuses them using dime-dependent weight, w/. The fused feature map, Ffuseds computed as:
Ffused = (1 We) O Prefined + we C Fprior
The weight is designed to be high for large values of /. (initial high-noise stages) and diminish as approaches This ensures that the guidance from the prior is strongest when the main branch is most unstable and gracefully recedes the model's own predictive capabilities become more reliable, for details, please refer to Appendix B. The above design constitutes APDS, which can be seamlessly integrated into any 3D diffusion model to form stable infrastructure.
3.2 CONV-TLSTM
The Conv-LLSTM module, illustrated in Fig. 2. extends the classical LSTM to 3D volumetric data by replacing matrix multiplications with 3D convolutions all gating operations. Unlike conven- tional LSTMs that only store abstract vector states, Conv-tLSTM maintains both the hidden state h2 and the memory cell C₁ as full 3D tensors, thereby preserving spatial correlations across voxels.
To align with our cross-timestep memory design, we modify the hidden state update with time- aware modulation. Specifically, the previous hidden state h4-1 is first modulated by timestep embedding F,, producing My which allows the recurrent unit to remain aware of the denoising stage.
All three gates input gate i, forget gate ft, and output gate o are computed using 3D convolu- tions:
4
Conv(Wz: X, + Wai h4 + b.), fe aWa, X. + Whr 14 + by),
(3)
= o zo X, + Who * h4 + bo),
(4)
where X, is the noisy input timestep t, and denotes 3D convolution.
Based on these gates, the candidate memory Čt is computed as: and the cell state is updated as:
Č7 = tanh(Wae * Xt + Who hig + be),
(5)
(6)
Finally, the hidden state is updated as:
hu = o @ lanh
(7)
Intuitively, the cell state O2 serves as the cross-timestep memory carrier, explicitly retaining: Low- frequency structural sketches (Č, contributions), Residual noise statistics (filtered through fa), Uncertainty-aware saliency cues (modulated by and oz). Thus, Conv-ILSTM transforms the re- verse diffusion process into memory-guided trajectory where each denoising step refines, rather than re-discovers, structural evidence.
3.3 LINEAR-TGRU
While Conv-tLSTM excels at capturing spatial dependencies, its reliance 3D convolutions incurs significant computational cost. To achieve a more lightweight yet effective recurrent design, we propose Linear-tGRU (Fig. 2d), which builds upon the Gated Recurrent Unit (GRU) architecture. Compared to LSTM, GRU merges the cell and hidden states into single hidden representation, and fuses the forget and input gates into an update gate, reducing redundancy.
Similar to Conv-LLSTM, Linear-tGRU introduces timestep-awarc modulation by embedding the cur- rent step f into the hidden state h2-1, yielding h4. Two gates are then computed:
F+ oWW., X. Whot,, br), 24 oW + Whehi, bz).
(8)
4

---

Published as a conference paper at ICLR 2026
The reset gate r2 controls how much past information is forgotten, while the update gate 2 balances between retaining history and incorporating new evidence. The update of the candidate hidden state is provided in Appendix C.
By design, LiRuar-IGRU primarily retains global descriptors of the diffusion trajectory. Specifically, 2, ensures persistent memory of coarse structures and residual statistics across timesteps, while re selectively refreshes saliency regions when strong new evidence emerges. Compared with Conv- ILSTM, it sacrifices fine-grained spatial encoding for computational efficiency, making suitable for long-horizon denoising.
3.4 SC-TLSTM
Building upon our foundational Conv-LLSTM and Linear-tGRU units, the Spatial-Channel Trans- temporal Memory LSTM (SC-LLSTM) module(Tig. (b) adapts the conventional spatial-channel attention mechanism to be stateful and temporally aware. This module processes an input feature map through two recurrent branches tLSTM module creates a dynamic, stateful version of spatial-channel attention, using our recurrent units to learn what and where focus throughout the entire reconstruction process.
spatial attention branch and channel attention branch. The SC-
3.4.1 SPATIAL ATTENTION BRANCH
we first generate compact feature descriptors from the input feature map F E RCCxxHxW This is achieved by applying average pooling independently along the X, Y, and 7. axes, thereby creating three distinct spatial summaries, Pzyz. These summaries concatenated and fed into a recurrent block, primarily composed "Conv-LLSTM'. which maintains memory of spatial patterns across diffusion timesteps. The output a 3D spatial attention map, M,. computed as:
Payz Concat Pool, Poolz
M, (LS)
(10)
3.4.2 CHANNEL ATTENTION BRANCH
Concurrently, to model the temporal evolution of cnter-channel relationships, the module aggregates spatial information using both average-pooling and max-pooling operations across the spatial di- mensions of the input feature map F. The resulting channel descriptors Pehannet are then processed by another recurrent block, which predominantly uses 'Linear-tGRU" This branch tracks the evo- lution of channel importance over time, producing a channel attention map, Mo:
Pehannel Concall (AugPo(F) MarPool(F)),
(11)
M. tGRU (GRUUPPaanna (Pehannne)).
(12)
3.4.3 FEATURE REFINEMENT
The two attention maps are applied sequentially to the input feature map F to adaptively refine it, the channel attention is applied first, followed by the spatial attention, yielding the final refined output, Fout:
F' = M. @ F.
(13)
Foot M.OF.
(14)
3.5 FFT-TLSTM
The core principle of FFT-LLSMM shown in Fig. 2 is to leverage the frequency domain, where structural information and noise components often more separable. The mechanism involves sequence transformations and stateful modulations. First, both the noisy input image X, and the conditional image X. are transformed from the spatial domain to the frequency domain using a 3D Fast Fourier Transform (FFT), F; and Fe:
FFT(X)), FF=((((X)),
(15)
The noisy and conditional spectra are then combined and filtered. This fused representation is pro- cessed by stateful recurrent block. Similarly, the memory of time steps enables it to better grasp the scale of frequencies that belong noise in the frequency domain space. The output is subsequently
5

---

Published as a conference paper at ICLR 2026
modulated by the conditional spectrum Fe. which acts as a gate to amplify relevant structural fre- quencies, obtained F. This modulated spectrum is then transformed back to the spatial domain via an inverse FFT (iFFT), obtain the noisy image after frequency-doman denoising. Finally, the sult is combined with the original noisy input through residual connection, yielding final output, Xou The entire process can be formulated as follows:
f = LLLLLM LLSTM + F.)) Fo:
(16)
Xout = (FFT((F) + Xx.
(17)
3.6 DIFFUSION PROCESS
The diffusion process, whose overall network structure is depicted in Fig. 2h. consists of two primary stages: a fixed forward diffusion process and learned reverse denoising process. forward pro- cess progressively corrupts the ground truth segmentation mask Tpp over sequence of T timesteps by adding Gaussian noise. This is fixed Markov chain defined as:
(( 1+- = (m= v1 - Be
(18)
where {.3. 1 is a predefined variance schedule. As f + T, the distribution of æy approaches a standard isotropic Gaussian distribution, TT (O, I).
The reverse process where our neural network learns to denoise these corrupted inputs. is learned Markov chain that begins with pure Gaussian noise IT and iteratively refines to generate the final segmentation mask Zo. To achieve image-conditional segmentation, the reverse process is guided by the clean medical image, X-. Our network, denoted as Me. is trained to predict the noise component from the noisy mask 2; at given timestep I, conditioned on X The training objective follows the simplified loss function common in DDPMs:
Lsimple Ex. [l|e Molse, 1)
(19)
where xy = V0, + vT- aie and E (O.
At inference time, the segmentation process starts by sampling a tensor of pure Gaussian noise, TT- The model then iteratively applies the learned denoising function for = ....., , to produce a progressively cleaner sample 2 from Ft. This iterative refinement, guided at each step by the innovations within our architecture, can be formally expressed as:
DSCAADs (ESc.FTTEE f), Xe,
(20)
1
au
-
ax Ep 022.
(21)
V
vI
Here, the sampling step utilizes the predicted noise co. which is generated by our complete network. The term ESCIFT represents the encoder, which leverages 'FFT-LLSMM' for noise resilience and 'SC-LLSTM" for stateful feature extraction. The term DSC,APDS represents the decoder, which uses 'SC-LLSTM blocks for reconstruction and is critically stabilized by the 'APDS' mechanism using the conditional image X.. This process is repeated until the final, high-fidelity segmentation mask To is produced. At each step we update cross pooss-tmmesep state S, and condition the denoiser on S,, enabling evidence accumulation across steps:
S, = SS H(F,, Xe,t))
(22)
EXPERIMENTS
4.1 EXPERIMENTAL SETUP
4.1.1 DATASETS
To rigorously evaluate the robustness and generalization of our proposed model under domain shift, we utilized two multi-center nasopharyngeal carcinoma (NPC) segmentation datasets: LNCTVSeg
(Luo et al.] 2024 2025) and OAseg (Wang et al. 2024b). LNCTV comprises 440 computed tomography (CT) images from 262 patients across four medical institutions, focusing on lymph node clinical target volume (CTV) delineation. OAseg includes Tl T--weigtted MRI scans for Gross Tumor Volume (GTV) segmentation collected from three distinct institutions: Center A (50 cases), Center B (50 cases), and Center C (60 cases). t-SNE visualization (Appendix D) confirmed significant stylistic variability across these datasets, providing an ideal testbed for assessing model performance under substantial domain shifts.

---

Published as a conference paper at ICLR 2026
4.1.2 BASELINE METHODS
We compared our method against several state-of-the-art (SOTA) segmentation models, including
TransBTS (Wang al. 2021). SwinUNETR (Hatamizadeh al.). UNETR (Hatamizadeh et al.
2022
3DUXNET (Lee et al. 2022), inFormer (Zhou et al. 2023). and Perspective+ (Hu et al,
2024
Detailed training parameters are described in Appendix E.
4.2
ADDRESSING 'NITIAL-STAGE COLLAPSE'
To empirically demonstrate the phenomenon of 'initial-stage collapse", we assessed the model's segmentation capability by initiating the reverse diffusion process from varying initial noise levels (timesteps ranging from low-noise (t=100) to high-noise (t=1000)). The average Dice scores were subsequently calculated as a function of these starting timesteps.
1.0
DIT ILSTM
DI 6.5TM
DET A/OS
De APDS
UUSTM APDS
DT 8,5TM APO3
ANg Dioe
Smestage
Smesteps
25.0
600
200
.000
200
600
200
1000
a) LNCTVSeg
b) OASeg
Figure 3: The average dice value obtained by performing reverse diffusion starting from noisy im- ages at different time steps.
The results, depicted in Fig. for both the LNCTVSeg and OAseg datasets, clearly illustrate this collapse. The baseline diffusion model ('Diff") and the variant enhanced solely with our temporal consistency module ('Diff ILSTM') exhibited robust segmentation accuracy at low-to-moderatt noise levels (t <600). However, their performance dramatically deteriorated as noise increased, collapsing entirely at timesteps beyond 700, where the input images were equivalent to random noise. resulting in incoherent predictions.
Conversely, models incorporating APDS ('Diff + APDS' and the full 'Diff + ILSTM + APDS' variant) exhibited remarkable stability across all noise levels. Although their performance slightly declined with increased noise, these models consistently maintained effective segmentation capa- bilities, even in high-noise regimes where conventional models failed. This result conclusively indicates that APDS provides essential structural guidance, preventing 'initial-stage collapse' and ensuring the reverse diffusion process remains viable solution trajectory from the outset. Con- sequently, APDS emerges as universally applicable framework for enhancing stability in future 3D diffusion- based medical segmentation models.
4.3 MITIGATING APDS INTERFEREEEEE
A potential concern regarding the introduction of the Adaptive Priori Decoding Strategy (APDS) its strong structural prior, which might inadvertently cause the model overly depend on this guidance. Such dependence could degrade the model into a simpler refinement mechanism resem- bling a conventional U-Net. To investigate potential over-interference by APDS, we conducted an experiment to compare the performance between the main diffusion output ("Diff Out') and the APDS geDeereneetted prior ('APDS Out') throughout the reverse diffusion process.
Specifically, the reverse diffusion process commenced from timestep t=1000, and Dice scores were recorded for both 'Diff Out' and 'APDS Out' at each subsequent timestep. The results, depicted in Fig. indicate that at initial high-noise stages, the 'Diff' Out' initially struggles, relying heavily on the stable structural prior provided by APDS. However, as the timesteps decrease, the performance of Diff Out' rapidly improves, eventually surpassing 'APDS Out' and achieving significantly higher accuracy. This trend confirms that APDS effectively functions as intended supportive structural scaffold during unstable early stages ttahest constraining the model's learning capacity. Con-

---

Published as a conference paper at ICLR 2026
0.85
0.75
Aug Dice
0.70
0.70
0.00
Caf Out APDS Cul
DIf Out
APDS Out
0:55
1200
820
(44)
(P0
200
1000
890
600
600
250
timestops
timesteps
a) LNCTVSeg
b) OASeg
Figure 4: During a complete reverse diffusion process, the changes in the predicted results by the main denoising network ('Diff Out') and the Prior Decoder within APDS ('APDS Out').
sequently, approach maintains the progressive refinement characteristic inherent diffusion models, successfully mitigating concerns about APDS over-interference.
Table 1: Ablation study on tLSTM components.
Modules
LNCTVSeg
OASeg
LSTM Conv-LSTM
Lincar-GRU
t-cell
Dice f
IoU 1
HD95
Diee
IoU
HD95
66.6
81.1
68.2
5.48
70.6
54.6
8.89
65.6
7.12
52.8
10.11
82.5
70.2
3.81
71.7
56.0
7.49
74.2
2.44
72.8
65.4
6.24
4.4
MEMORY
The tLSTM module aims to ensure temporal consistency within the reverse diffusion process by preserving stateful memory across timesteps. To qualitatively validate this capability, we visualized the model's 'ss feature heatmap across different stages of the denoising trajectory for representative sample, as illustrated in Fig. 5 (details provided in Appendix F). Initially, at the high-noise stage (I=999), the model's attention diffuse, emphasizing broad structural reconstruction. As denoising progresses, the tLSTM leverages its temporal memory, refining attention patterns and increasingly focusing on the target structure. By timestep 1=799, the model starts developing more defined target-region focus, which progressively sharpens at timesteps t=399 and t=100. This visualization clearly demonstrates the evolving stateful attention driven by tLSTM's temporal memory.
1 869
1 799
399
10100
Figure 5: An example of the heat map during the reverse diffusion.
To evaluate the efficacy of the tLSTM module, we conducted comprehensive ablation study quan- tifying its performance contributions (T. Beginning with a baseline model devoid of temporal modules, the introduction
standard LSTM yielded notable performance improvements. Replac- ing standard LSTM with our specialized 'Conv-LSTM' which integrates 3D convolutional opera- tions with LSTM gates, and 'Linear-GRU', which combines linear layers with the GRU structure, specifically designed for 3D imaging, further enhanced performance. The optimal performance was attained by integrating our proposed 't-cell' innovation. These findings underscore the importance of explicit temporal awareness in achieving maximum stability and accuracy.
8

---

Published as a conference paper at ICLR 2026
Moreover, this analysis validates the versatility of 'Conv-LLSTM' and "Linear-tGRU" units, demon- strating their fundamental role as adaptable attention modules within diffusion models. Their recom- bination potential offers extensive opportunities for future architectural innovations and advance- ments.
4.5 COMPARISON WITH STATE-OF-THE-ARR METHODS
We conducted a comprehensive comparison be- tween proposed model and several state- of-the-art (SOTA) segmentation methods on the LNCTVSeg and OAseg datasets. As detailed in Table our model consistently outperformed existing approaches across all evaluated metries. Specifically, our method achieved the highest Dice and loU scores, demonstrating superior ac- curacy in volumetric overlap with ground truth segmentations. Additionally, our model exhib- ited the lowest HD95 scores, reflecting enhanced precision in segmentation boundary delineation compared to other leading methods. Visual com- parisons provided Appendix further illus- trate these advantages. These results underscore the robustness and effectiveness of our proposed architecture, particularly when addressing signif- icant stylistic variability across multi-center medical imaging datasets.
Table 2: Comparison with tate-of-theeart meth- ods.
Method
LNCTVSeg
OASeg
Dicet IoU* HD95, Dicet IoU1 HD95
TransBTS SwinUNETR
78.6
68.7
5.14
66.3
55.4
12.7
80.7
71.1
4.14
65.9
58.3
9.50
UNETR
78.1
69.1
6.87
67.1
59.2
9.04
unFormer
80.3
71.5
4.31
68.4
62.4
7.76
3DUXNET
81.6
72.9
3.52
68.9
61.3
7.93
Perspective+
82.4
73.6
3.27
69.6
62.8
7.09
Diff-UNet
81.7
72.2
3.91
71.5
64.2
6.88
Ours
83.7
74.2
2.44
72.8
65.4
4.6 QUANTITATIVE COMPARISON OF COMPUTATIONAL COST
To more rigorously quantify the computational efficiency of the proposed Cross-Timessee Timestep frame- work, we conducted a comparative analy- sis against several representative 3D segmen- tation architectures namely UX-Net, Swin- UNETR, nnFormer, and the diffusion-based Diff- UNet-on the LNCTVSeg dataset. Four comple- mentary indicators were assessed to capture over- all computational cost: training time, inference latency, GFLOPS, and GPU memory consump- tion. All models evaluated under consistent experimental conditions, and the complete proto- col is documented in Appendix H.
Table 3: Ablation study APDS, SC, and FFT modules.
Modules
LNCTVSeg
OASeg
APDS SC FFT
Dicef
ToUf HD95
Dicet IoUf HD95
77.3
63.0
8.51
67.5
50.9
11.53
82.1
69.6
4.32
71.3
55.4
7.83
81.8
69.2
4.85
70,9
54.9
8.14
v
83.7
74.2
2.44
72.8
65.4
6.24
As summarized in Table the proposed method exhibits a balanced efficiency profile in comparison with both transformer- based and diffusion-based baselines. Notably, Cross-Timeste requires sub- stantially less training time than large transformer architectures such as SwinUNETR and nnFormer, while maintaining moderate memory footprint and competitive runtime. Moreover, relative the diffusion- based Siff-UNNt-whitne strates comparable resource usage while offering markedly improved segmentation accuracy. These results collectively indicate that Cross- Cross-Timested Timestep achieves favorable trade-off between computa- tional efficiency and predictive performance, strengthening its practicality for real-world clinical deployment.
is optimized for lightweight computation method demon-
4.7 ABLATION STUDY
To quantify the contribution of individual modules within our architecture, performed an ablation study, as summarized in Table The baseline model included the core diffusion framework and the essential Adaptive Priori Decoding Strategy (APDS) to address 'initial-stage collapse'. Introduc- ing either the SC-LLSMM (SC) or FFT-LLSTM (FFT) modules individually significantly improved performance both datasets, demonstrating the efficacy of spatio-channel temporal memory and

---

Published as a conference paper at ICLR 2026
Table 4: Detailed computational cost including training/inference time, GPU memory usage, and GFLOPS
Method
Training Time (h)
Inference Time (s)
GFLOPS GPU Mem (MiB)
UX-Net
33.6
0.11
631.9
14722
SwinUNETR
45.3
328.8
16364
nnFormer
80.4
0.03
235.3
11626
Diff-UNet
29.2
0.12
981.1
10834
Ours
33.4
0.17
453.5
15152
frequency-domate noise resilience enhancements. These findings highlight the versatility and poten- tial of Conv-LLSTM and Lincar-tGRU modules as foundational components that can be effectively leveraged to further extend and refine diffusion-bart based segmentation models.
CONCLUSION
In this study, we identified and addressed critical challenge- aoaalengee"'nittil-sttttee 'initial-stage collapse' -aaapssee-hhtttt signif- icantly constrains the applicability of diffusion models in 3D medical segmentation tasks. We troduced Cross- Timestep, a versatile 3D diffusion framework integrating APDS and tLSTM. These components collectively provide structural stability during early diffussinn phases and maintain tem- poral consistency across timesteps. Experiments conducted on multi-center datasets demonstrated that Cross- Timestep effectively resolves 'initial-stage collapse' and achieves state-of-the-art seg- mentation performance. Our framework provides robust foundation for advancing 3D diffusion- based medical segmentation methodologies and illustrates considerable potential for integration with other sophisticated architectures.
REPRODUCIBILITY Our code and data
are
made
publicly
available at
ACKNOWLEDGMENTS. This research was supported by the National Natural Science Founda- tion of China (Grant Nos. U23A20321 and 62272490); the Natural Science Foundation of Hunan Province of China (Grant No. 2025JJ20062); the National Natural Science Foundation of China (Grant No. 82471964),
REFERENCES
Tao Chen, Chenhui Wang, Zhihao Chen, Yiming Lei, and Hongming Shan. Hidiff: Hybrid diffusion framework for medical image segmentation. IEEE Transactions on Medical Imaging, 43(10); 3570-3583, 2024.
Weiping Ding, Sheng Geng, Haipeng Wang, Jiashuang Huang, and Tianyi Zhou. Fdiff-fusion Denoising diffusion fusion network based on fuzzy learning for 3d medical image segmentation. Information Fusion, 112:102540, 2024.
A Hatamizadeh, V Nath, Y Tang, D Yang, HR Roth, and D Xu. Swin transformers for semantic segmentation of brain tumors in mri images. Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries, pp. 272-284.
Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Land- man, Holger R Roth, and Daguang Xu. Unetr: Transformers for medical image segmentation. In Proceedings of the IEEE/CVF winter conference applications of computer vision, pp. 574- 584, 2022.
Jintong Hu, Siyan Chen, Zhiyi Pan, Sen Zeng and Wenming Yang. Perspective+ unet: Enhancing segmentation with bi bi-path fusion and efficient non-local attention for superior receptive fields. In International conference on medical image computing and computer-assisted intervention, pp- 499-509. Springer, 2024.
10

---

Published as a conference paper at ICLR 2026
Wen Ji and Albert CS Chung. Diffusion-bart based domain adaptation for medical image segmentation using stochastic step aligmment. In International Conference on Medical Image Computing and Computer-Assiseed Intervention, pp. 188-198. Springer, 2024.
Ho Hin Lee, Shunxing Bao, Yuankai Huo, and Bennett A Landman. 3d ux-net: large kernel volumetric convnet modernizing hierarchical transformer for medical image segmentation, 2022.
Tiange Liu, Jinze Li, Drew A Torigian, Yubing Tong, Qibing Xiong, Kaige Zhang, and Jayaram K Udupa. Diffusion semantic segmentation model: A generative model for medical image segmen- tation based on joint distribution. Medical physics, 52(7)=e17225 2025.
Xiang Luo, Guotai Wang, Weixiong Liao, Jiacheng Chen, Shichuan Zhang, Hong Wang, Tao Zhou, Tao Song, Ying Chen, Yuan Zhang, Huai Liu, Wen-Fei Li, Ying Sun, Shaoting Zhang, Heye Zhang, and Jian Zhou. Semi-supervised medical image segmentation via uncertainty rectified pyramid consistency. Medical Image Analysis, 80: 102517, 2022.
Xiang Luo, Weixiong Liao, Yan Zhao, Luting Lin, Teng Liu, Yuan Tian, Bi-Xiu Huang, Wei Li, Ying Chen, Yuan Zhang, Jia-Wei He, Xiao-Jing Wang. Wen-Fei Li, and Ying Sun. multicenter dataset for lymph node clinical target volume delineation of nasopharyngeal carcinoma. Scientific Data, 11(1):1085, 2024.
Xiang Luo, Hong Wang. Jing Xu, Jiacheng Chen, Tao Zhou, Shichuan Zhang, Yuan Zhang, Ying Chen, Wen-Fei Li, Huai Liu, and Ying Sun. Generalizable magnetic resonance imaging-based nasopharyngeal carcinoma delineation: Bridging gaps across multiple centers and raters with active learning. International Journal of Radiation Oncology* Biology* Physics, 121(5):1384- 1393, 2025.
Xiangde Luo, Jieneng Chen, Tao Song, and Guotai Wang. Semi- supervised medical image seg- mentation through dual-task consistency. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 8801-8809, 2021a.
Xiangde Luo, Guotai Wang, Tao Song, Jingyang Zhang, Michael Aertsen, Jan Deprest, Sebastien Ourselin, Tom Vercauteren, and Shaoting Zhang. Mideepseg: Minimally interactive segmentation of unseen objects from medical images using deep learning. Medical image analysis, 12:102102, 2021b.
Zuxin Nie, Jiahong Yang, Chengxuan Li, Yaqin Wang, and Jun Tang. Diffbts: A lightweight diffu- sion model for 3d multimodal brain tumor segmentation. Sensors, 25(10):2985, 2025.
Mattia Salsi, Yunying Wang, Chen Hu, Yueyue Hu, Hanchi Ren, Jingjing Deng, and Xianghua Xie. Guided latent diffusion for universal medical image segmentation. In International Conference on Al-Generated Contens (AIGC 2024), volume 13649, pp. 48-55. SPIE, 2025.
Zhihao Shuai, Yinan Chen, Shunqiang Mao, Yihan Zho, and Xiaohong Zhang. Diffseg: a segmen- tation model for skin lesions based on diffusion difference. arXiv preprint arXiv:2404.16474, 2024.
Qian Tang, Qikui Zhu, Yuxuan Xiong, Yongchao Xu, and Bo Du Edge-and-mask integration- driven diffusion models for medical image segmentation. IEEE Signal Processing Letters, 31: 2665-2669, 2024.
Haoshen Wang. Zhentao Liu, Kaicong Sun, Xiaodong Wang, Dinggang Shen, and Zhiming Cui. 3d meddiffusion: A 3d medical diffusion model for controllable and high-quality medical image generation, 2024a.
Hongqiu Wang, Jian Chen, Shichen Zhang, Yuan He, Jinfeng Xu, Mengwan Wu, Jinlan He, Wenjun Liao, and Xiangde Luo. rufl-reference source-free active domain adaptation for nasopharyngeal carcinoma tumor segmentation across multiple hospitals. IEEE Transactions on Medical imaging. 43(12))4078-400) 2024b.
Wenxuan Wang, Chen Chen, Meng Ding, Hong Yu, Sen Zha, and Jiangyum Li. Transbts: Mul- timodal brain tumor segmentation using transformer. In International conference on medical image computing and computer-asiited intervention, pp. 109-119. Springer, 2021.
11

---

Published as a conference paper at ICLR 2026
Junde Wu, Wei Ji, Huazhu Fu, Min Xu, Yueming Jin, and Yanwu Xu. Medsegcifftv2. Diffusion- based medical image segmentation with transformer. In Proceedings of the AAAI conference on artificial intelligence, volume 38, pp. 6030-6038 2024.
Rong Wu, Ziqi Chen, Liming Zhong, Heng Li, and Hai Shu. Unleashing diffusion and state space models for medical image segmentation, 2025.
Zhaohu Xing, Liang Wan, Huazhu Fu, Guang Yang, Yijun Yang, Lequan Yu, Baiying Lei, and Lei Zhu. Diff-unet: A diffusion embedded network for robust 3d medical image segmentation.
Medical Image Analysis, 105:103654, 2025. ISSN 1361-8415. doi:
media.2025.0 03654, URL https: www. sci encedi rect com wttps/ccienncarrre pi
Fahim Ahmed Zaman, Mathews Jacob, Amanda Chang. Kan Liu, Milan Sonka, and Xiaodong Wu. Surf-cdm: Score-based surface cold-diffusion model for medical image segmentation. In 2024 IEEE International Symposium on Biomedical Imaging (ISBI), pp- 1-5. IEFF, 2024.
Fahim Ahmed Zaman, Mathews Jacob, Amanda Chang, Kan Liu, Milan Sonka, and Xiaodong Wu. Latent diffusion for medical image segmentation: End-to-end learning for fast sampling and accuracy. Biomedical signal processing and control, 114: 10:300,,,, 2026.
Le Zhang, Fuping Wu, Kevin Bronik, and Bartlomiej W Papiez. Diffuseg: domain diffusion for medical image segmentation. IEEE Journal of Biomedical and Health Informatics, 29(5): 3619-3631, 2025a.
Li Zhang, Basu Jindal, Ahmed Alaa, Robert Weinreb, David Wilson, Eran Segal, James Zou, and Pengtao Xie. Generative ai enables medical image segmentation in ultra low-data regimes. Nature Communications, 16(1))6486, 2025b.
Hong- Yu Zhou, Jiansen Guo, Yinghao Zhang, Xiaoguang Han, Lequan Yu, Liansheng Wang, and Yizhou Yu. anformer: Volumetric medical image segmentation via 3d transformer. IEEE Iransactions on image processing, 32:4036-4045, 2023.
SUPPLEMENTARY MATERIAL
ETHICS STATEMENT
This work adheres to the ICLR Code of Ethics. In this study, no human subjects or animal ex- perimentation was involved. All datasets used were sourced in compliance with relevant usage guidelines, ensuring no violation of privacy. We have taken care to avoid any biases discrimi- natory outcomes in our research process. No personally identifiable information was used, and no experiments were conducted that could raise privacy or security concerns. We are committed to maintaining transparency and integrity throughout the research process.
REPRODUCIBILITY STATEMENT
We have made every effort to ensure that the results presented in this paper are reproducible. The ex- perimental setup, including training steps, model configurations, and hardware details, is described in detail in the paper.
Additionally. all datasets used in this paper secieene/statss/rros-imeeee sistency and reproducibility of the evaluation results.
publicly available ensuring
resources(httss//anonyym
the con-
We believe these measures will enable other researchers to reproduce our work and further advance the field.
12

---

Published as a conference paper at ICLR 2026
Diffusion sampling
Diff
APDS
1000
800
500
Timestep
Figure 6: Visualization of the reverse diffusion process with and without APDS.
LLM USAGE
Large Language Models (LLMs) were used to aid in the writing and polishing of the manuscript. Specifically, we used an LLM to assist in refining the language, improving readability, and ensuring clarity in various sections of the paper. The model helped with tasks such sentence rephrasing, grammar checking, and enhancing the overall flow of the text.
important to note that the LLM experimental design. All research concepts, ideas, and analyses were developed and conducted by the authors. The contributions of the LLM were solely focused on improving the linguistic quality of the paper, with no involvement in the scientific content or data analysis.
involved in the ideation, research methodology, or
The authors take full responsibility for the content of the manuscript, including any generated or polished by the LLM. We have ensured that the LLM- L_M-generated text adheres to ethical guidelines and does not contribute plagiarism or scientific misconduct.
APPENDIX A
To further illustrate the stabilizing effect of APDS, Fig. 6 presents visualization of the reverse diffusion process with and without APDS (APDS and Diff). We show representative timesteps (t = 1000, 800, 500) and the final prediction. Without APDS, the model produces unstable, structureless patterns at high-noise timesteps, reflecting the initial-stage collapse phenomenon. contrast, APDS injects explicit prior guidance early timesteps, allowing coarse anatomical structures gradually emerge and resulting in a stable denoising trajectory throughout the diffusion process.
APPENDIX B
The Adaptive Prior-guidance Denoising Schema (APDS) introduces a dynamic weighting factor, was to modulate the influence of the external prior throughout the reverse diffusion process. The core design principle is to apply strong guidance during the initial, high-noise stages of generation (when t is large) and to progressively reduce this guidance as the model's own generative capacity becomes more reliable in later, low-noise stages (as - 0), as illustrated Fig. This ensures stability without sacrificing the fine details learned by the model.
To provide explicit control over the maximum guidance strength, we introduce a hyperparameter, [0, 1], which acts scaling factor threshold. This allows for fine-tuning the upper limit of the prior's 's influence. For instance, a value of a = 0.5 would cap the maximum weight at 50%. The weight w, is defined function of the normalized time Inormalired [0, 1], where mmormalized 1/T for a total of timesteps. The formulation is given by:
Wi a exp(-5.0 (1 nnomalized)) (1 exp(-10.0 (nommalized)
(23)
To understand the behavior of w;, we can analyze its constituent terms: Primary Decay Term (1 exp(- 10.0- (nnmalived)): 1
This component ensures that the guidance is nullified at the very end of
13

---

Published as a conference paper at ICLR 2026
1.0
Wt
0.8
Weight Value ah
0.6
0.4
0.2
0.0
1000
800
600
400
200
timestep
Figure 7: "Initial-stage collapse".
the diffusion process. As normalized - 0, the term exp(-10.0 (nomalizedd - 1, which causes the entire expression to approach O. This is critical for allowing the model to generate high-fidelity details without being constrained by the prior. For Envemalized 0, this term rapidly increases towards 1, effectively "activating" the guidance. Shaping Term (exp(-5.0 (1 nnommllize4):: This term sculpts the decay curve. At the beginning of the process, when Inormalized - this term approaches exp(0) = 1. Consequently, the overall weight we approaches its predefined maximum value of a. As 'nommalieed decreases, this term introduces smooth, exponential decay, ensuring graceful transition from a high-guidance regime to www-guiaance one, rather than an abrupt cutoff.
The combined effect of these terms weighting curve that starts O, rises peak value of approximately a in the early-to-mo mid stages, and then decays back to 0. The plot of versus the nor- malized time (from down 0) visually confirms this behavior. This carefully designed weighting function is integral to the stability and efficacy of the APDS framework, providing robust guidance where needed while preserving the generative autonomy of the underlying 3D diffusion model.
APPENDIX C
The Conv-tLSTM module is the core recurrent component of our architecture, designed to process 3D volumetric feature maps across the denoising trajectory. It extends the standard LSTM by in- tegrating 3D convolutional operations directly into its gating units, enabling it to effectively learn and propagate spatio-temporal information. This appendix provides a detailed breakdown of the mathematical operations within each gate and state update of the module.
In this module, all states (C,, fız) and inputs (X,) are 3D tensors. The *** operator denotes 3D convolution, and @ denotes the Hadamard (element-wise) product. The previous hidden state ht first modulated by learned time embedding E, produce time-aware hidden state h/, which is then used in the gate computations.
The module's behavior is governed by three primary gating units that regulate the flow of informa- tion.
Input Gate (i:): This gate determines which new information from the input X, will be stored in the cell state. It computes update mask by combining the current input and the time-aware previous hidden state. Its computation as follows:
oW zi X. + Who h4 b.)
(24)
14

---

Published as a conference paper at ICLR 2026
where X2 is the current input, W terms are 3D convolutional kernels, b, is the bias term, and is the sigmoid activation function.
Forget Gate (ft): This gate is responsible for deciding which information to discard from the previ- ous cell state, C2-1. This is critical for the model to forget irrelevant features from earlier stages of the long denoising trajectory. Its computation is:
fi o X1 + Way hi + by)
(25)
Output Gate (ot): This gate controls which information from the updated cell state will be used to compute the new hidden state. It acts as a filter on the cell's memory to produce the module's output for the current timestep. It is computed as:
= o X + Who K4 + bo)
(26)
Based on the outputs of these gates, the cell and hidden states are updated for the current timestep I.
First, a candidate cell state, Č₁ is computed from the current input and previous hidden state. This represents the new information that *could* be added the cell's memory.
C = tanh Wee nc X₁ + Who * hú + be)
(27)
Next, the new cell state, C, produced by combining the previous cell state, Cr-1, with the can- didate state. The forget gate determines what to discard from Ct-1, while the input gate determines what to add from C..
C. = fe C. + の Ć.
(28)
Finally, the new hidden state, h4, is determined by passing the updated cell state through a tanh function and then filtering it with the output gate.
h o2 C tan
(29)
Through this sequence of operations, the Conv-tLSTTM module selectively update its memory over time, preserving relevant spatial features while discarding irrelevant ones, making it a powerful tool for sequential generative processes.
The Linear-IGRU module serves as Domputationally efficient recurrent component within our framework, designed to balance performance with operational cost. It adapts the Gated Recurrent Unit (GRU) architecture, replacing computationally intensive 3D convolutions with standard linear transformations. This makes well-suited for processing feature vectors where full spatial convo- lution is not required. This appendix details the mathematical formulations of its gating units and state update process.
In this module, the ! operator denotes matrix multiplication, and ( represents the Hadamard (element-wise) product. As with the ConV-LLSTM, the previous hidden state f is modulated by time embedding E, to create a time-aware hidden state hy before being used in calculations.
The Linear-1GRU simplifies the information flow using two primary gating units:
Reset Gate (rt): This gate determines how much of the previous hidden state should be forgot- ten when computing the new candidate hidden state. This allows the model to effectively discard irrelevant past information that is not pertinent to the current input. Its computation is:
o oWWr. Xt + Whr 11 + b.)
(30)
where W terms are weight matrices, b, is the bias term, and a is the sigmoid function.
Update Gate (2ạ): This gate functions as a combination of the forget and input gates in an LSTM. decides how much of the previous hidden state to retain and how much of the new candidate state to incorporate, controlling the flow of information into the final hidden state. computed as:
2 (W 2= X₁ + Wha R/ + bz)
(31)
The final hidden state is updated in two-step process using the outputs from the gating units.
15

---

Published as a conference paper at ICLR 2026
First, candidate hidden state, heg is computed. The reset gate (rt) modulates the influence of the previous state (he), allowing the model to focus on generating new information from the current input X₁.
h4 = tanh(Wa., sh X + Whs (rx @ h4) + by.)
(32)
Finally, the new hidden state, h, is determined by linear interpolation between the previous hidden state h4-1 and the candidate hidden state h2. The update gate 2, controls this interpolation, balancing the information to be kept from the past with the new information to be added.
h = (1 - 2x) 3 h1-1 + 2x 3 hi
(33)
This streamlined architecture allows the Linear-tGRU to efficiently update its state, capturing tempo- ral dependencies without the overhead of maintaining separate cell state, thus offering powerful yet lightweight recurrent building block.
APPENDIX D
To visually substantiate the claim of significant domain shifts across the evaluation datasets, we employed the t-Distributed Stochastic Neighbor Embedding (t-SNE) technique. t-SNE is non- linear dimensionality reduction method particularly well-suited for visualizing the structure of high- dimensional data in low-dimensional space (e.g., 2D plot). For this analysis, we extracted deep features from representative set of images from each source institution and projected them into 2D space. The resulting plots vividly illustrate the data distribution and intriiinstttutionall variability.
Center
Center B
Center
Center
150
50
1100
-150
-150
Figure 8: I-SNE visualization of the LNCTVSeg dataset. Data points are colored according their source institution. The formation of separate clusters confirms the multi-center heterogeneity of the data.
Fig. X displays the t-SNE plot for the LNCTVSeg dataset. The data points, representing CT images from four different medical institutions, also form distinct groups. This visualization confirms that despite all being CT scans for NPC, there are notable stylistic and technical differences between the data from each source.
16

---

Published as a conference paper at ICLR 2026
CenterA
CenterB
CenterC
20
20
-40
-60
-80
-60
-40
-20
0
20
40
60
Figure 9: t-SNE visualization of the OAseg dataset. Data points are colored according to their center of origin (Center A, B, and C). The distinct clustering demonstrates significant domain shift across the three institutions.
As illustrated in Fig. the t-SNE visualization of the OAseg dataset reveals three semi-distinct clusters, each corresponding to one of the contributing medical centers (Center A, Center B, and Center C). The clear separation between these clusters indicates substantial differences in the un- derlying data distributions. These variations likely arise from differences in MRI scanners, imaging protocols, and patient populations, presenting a significant domain adaptation challenge.
Collectively, these visualizations provide compelling evidence of the domain heterogeneity present both within each multi-center dataset and between the two datasets (CT vs. MRI). This confirms that chosen datasets provide rigorous and realistic testbed for evaluating the robustness and domain generalization capabilities of our proposed model.
APPENDIX E
To ensure a fair and reproducible comparison, all experiments were conducted using a consistent training and evaluation framework. This appendix details the hardware, software, and hyperparam- eters used for training our proposed model and the baseline methods.
Hardware: All models were trained server equipped with four NVIDIA A6000 graphics cards, each with 48 GB of VRAM. Software: The implementation was based on the Py Torch deep learning framework (version 1.13). The training environment utilized Python 3.8, CUDA 11.7, and cuDNN 8.5 for optimized performance.
The following hyperparameters were used for training our model and for re-implementing the base- line models to ensure fair comparison: Optimizer: We used the AdamW optimizer for all training processes, as it is well-suited for both CNN and Transformer-bant The initial learning rate was set to 5 10-5 Learning Rate Scheduler: A cosine annealing scheduler was employed to adjust the learning rate over the course of training. A warm-up period 50 epochs was used at the beginning of training, during which the learning rate gradually increased to its initial value. Loss Function: A compound loss function, combining Dice Loss and Cross- -nttroyy
architectures. Learning Rate:
Loss, used supervise the segmentation. The final loss computed as L Len-
17

---

Published as a conference paper at ICLR 2026
This combination provides balance between region- and pixel-wise accuracy. Epochs and Batch Size: All models were trained for a total of 200 epochs to ensure full convergence. Due to the memory constraints of 3D data, a batch size of 4 was used (distributed as I samples per GPU). Input Patch Size: During training, we extracted random patches of size 96 x 96 x 96 from the full-resolution images to serve as inputs to the models.
To improve model robustness and prevent overfitting, we applied a series of on-the-fly fly data aug- mentation techniques the input patches during training. These augmentations include: Random rotation in the range of [-30%, 30 Random scaling with a factor between 0.7 and 1.4, Random elastic deformation, Random gamma correction with a gamma value between 0.7 and 1.5, Random horizontal and vertical flipping (mirroring).
For the baseline models (TransBTS, SwinUNETR, etc.), we utilized their officially available code- bases and default configurations where possible, and we carefully reproduced them within our framework using the same training strategy described above to maintain consistency across all com- parisons.
APPENDIX F
To qualitatively evaluate the (LSTM module's ability to maintain stateful memory across the reverse diffusion process, we visualized the evolution of the model's internal feature representations dif- ferent timesteps. The feature heatmaps shown in Fig. 4 were generated from the output feature maps of the final tLSTM module in our model's architecture. The process began by selecting a represen- tative sample from the test set and running the full reverse diffusion trajectory. We captured the output 3D feature tensor from 1=999 to =100 To create a single-channel heatmap from this multi- channel tensor, we computed the mean activation across the channel dimension, which summarizes the model's s spatial focus at that instant. Finally, the resulting 3D heatmap was upsampled to the original image resolution and overlaid on the corresponding axial slice of the input image using "jet" colormap, where red indicates high activation and blue indicates low activation.
The sequence of heatmaps in Fig. 10 provides clear visual narrative of the tLSTM's temporal dynamics. Initially, at the high-noise stage, the model's ' attention is diffuse, with the heatmap show- ing broad, low-intensity activations that reflect the initial uncertainty and focus reconstructing the general anatomy. As the denoising process continues, the tLSTTM begins to leverage its mem- ory, causing the heatmap concentrate as background activations diminish and the focus on the potential target region intensifies. This localization becomes significantly more pronounced, where the model's ' attention is clearly centered on the target structure. By the final refinement stage, the heatmap is sharp and precisely aligned with the target's boundaries, with background noise almost entirely suppressed. In summary, this progressive sharpening from a diffuse to a highly localized focus provides strong qualitative evidence that the tl tLSM module effectively maintains and utilizes stateful temporal memory, mechanism crucial for ensuring coherence and accuracy throughout the multi-step reverse diffusion process.
APPENDIX G
Fig. presents qualitative segmentation results on both the LNCTVSeg (CT) and OAseg (MRI) datasets. To provide a comprehensive visual comparison, we benchmark our method against several strong and widely adopted 3D medical segmentation baselines, including SwinUNETR, nnFormer, 3DUXNET, Perspective+, and the diffusion-basee model Diff-UNet.
Across both datasets, competing methods frequently show issues such as incomplete tumor cover- age, over-segmentttion into surrounding tissues, or blurred and irregular boundaries- particclaaall under strong domain shifts in the multi-center MRI data (OAseg). In contrast, our method consis- tently produces segmentations that exhibit more accurate localization, cleaner boundaries, and better volumetric alignment with the ground truth.
On the LNCTVSeg dataset, our model more faithfully captures the irregular lymph-node CTV shapes, while the baseline models often miss small subregions introduce spurious fragments. On the OAseg dataset, where significant inter-center variability exists, our model maintains robust performance and yields cleaner GTV predictions, demonstrating strong generalizability.
18

---

Published as a conference paper at ICLR 2026
1:899
599
499
399
1=299
Figure 10: Visualization of feature heatmaps at different denoising timesteps. The progressive fo- cusing of the heatmap demonstrates the tLSTM module's 's ability to leverage temporal memory to refine its attention on the target structure.
LNCTVSS
Image
SwinUNETR
unFormer
3DUXNET
Perspeetive+
DifF-UNa
Figure 11: Visual comparison on representative cases from the LNCTVSeg and OAseg datasets.
These visual comparisons validate the quantitative results reported in the main paper, confirming that approach not only achieves state-of otate-of-the-art the-art numerical performance but also delivers more reliable and clinically meaningful segmentation outputs.
APPENDIX H
To further evaluate the computational efficiency of the proposed framework, we benchmarked our model against several representative SOTA segmentation models on the LNCTVSeg dataset. All experiments were conducted under strictly controlled and consistent settings: the input voxel res- olution was uniformly resampled to 96x96x96, the batch size was fixed at 1, and all models were trained and tested using the same GPU environment. For each method, measured four key in- dicators trrining time, inference time, GFLOPS, and GPU memory consumption- to provide a comprehensive assessment of computational cost.
19

---

Published as a conference paper at ICLR 2026
As summarized in Table 4 the proposed model demonstrates balanced computational profile. Although it is not the fastest model in inference nor the most memory-effiiiet, achieves compet- itive performance while maintaining moderate resource demand. In particular, model requires substantially less training time than transformer-hant architectures. These results suggest that the proposed model attains effective trade-off between segmentation performance and computational efficiency.
APPENDIX
Diffusion models inherently involve a trade-off between computational cost and reconstruction qual- ity. Using fewer diffusion steps reduces runtime and memory consumption but may lead to de- graded performance, while more steps typically improve output quality at the cost of higher com- putation. Following commonly adopted settings (500-1000 steps), evaluated our framework the LNCTVSeg dataset using 300, 500, and 1000 diffusion steps. The results are summarized in T showing that insufficient steps (e.g., 300) lead notable performance drops, whereas 1000 steps yield the best overall accuracy.
Table 5: Performance under different diffusion steps on the LNCTVSeg dataset.
Steps
Dicet
IoUf
HD95
300
79.3
68.9
4.86
500
82.8
73.0
3.39
1000
83.7
74.2
2.44
20


Reviews:
[{"review_id": "KJeNWMIoDE", "reviewer": "Reviewer_mjAf", "summary": "This paper presents Cross-Timestep, a framework for 3D medical image segmentation using diffusion models. The method introduces two key innovations. First, an Adaptive Priori Decoding Strategy (APDS) generates a time-weighted prior mask to stabilize the model during the early, high-noise stages. Second, a trans-temporal memory unit (tLSTM) with variants (SC-tLSTM and FFT-tLSTM) is designed to overcome the isolated, step-by-step nature of the denoising process, allowing the model to accumulate evidence across timesteps. Experiments on two multi-center datasets confirm that APDS successfully mitigates “initial-stage collapse” and tLSTM improves the coherence of the denoising trajectory.", "questions": "1. Computational Cost: Could the authors provide a table comparing inference time and VRAM requirements for the proposed model and other baselines?\n2. t-cell Explanation: Could the authors clarify what the \"t-cell\" in Table 1 refers to? How is it architecturally distinct from the Conv-tLSTM and Linear-tGRU modules?\n3. Sensitivity of ω_t: The time-weight ω_t (Appendix A) has a rather complex formulation. How sensitive is the model's performance to this specific function? Did the authors experiment with simpler decay functions (e.g., linear, exponential decay)?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Identifying an Under-Explored Problem: The paper identifies and addresses the \"initial-stage collapse\" problem in 3D medical diffusion models. As demonstrated in Fig. 1, the authors show that standard 3D diffusion models fail when sampling from high-noise timesteps, but succeed when starting from mid-level noise. This finding renders the proposed Adaptive Priori Decoding Strategy (APDS) an intuitively appealing and well-motivated solution.\n2. Achieving SOTA Performance: The paper thoughtfully combines effective components (SC-tLSTM and FFT-tLSTM) to create a powerful, stateful denoiser. This architecture achieves state-of-the-art performance on two multi-center 3D datasets (LNCTVSeg (CT) and OAseg (MRI)), demonstrating the framework's robustness and practical utility.", "weaknesses": "1. Limited Architectural Novelty: The proposed framework appears to be a complex integration of existing and well-established components, rather than a new algorithmic contribution. \n    + On the SC-tLSTM: The core idea of tLSTM is to apply a recurrent model (LSTM or GRU) to maintain a state across the diffusion timesteps. However, the combination of tLSTM and diffusion models tends to be inelegant. It forces an external, stateful memory (tLSTM) onto the diffusion model's fundamentally Markovian process, which seems unjustified from a theoretical standpoint. As demonstrated by prior work [1], LSTM-based U-Net backbones have already been explored and proven effective for 3D medical segmentation. The SC-tLSTM module itself appears to be a straightforward combination of the standard spatial-channel attention mechanism [2] and the tLSTM component.\n    + On the FFT-tLSTM: The novelty of FFT-tLSTM is questionable. Leveraging the frequency domain (via FFT) within a diffusion model framework for medical image segmentation has already been published [3]. The proposed FFT-tLSTM thus appears to be a minor variation that simply inserts the tLSTM components into an established FFT pipeline.\n    + On the APDS: The APDS uses a segmentation decoder that operates solely on the conditional image to generate a coarse structural prior. This concept of using the conditional input to create explicit guidance is a well-established technique in conditional diffusion models and thus cannot be regarded as a significant innovation.\nThe authors should more clearly highlight their components' algorithmic novelty over prior research. They should also strengthen the explanation of how the integration of these modules provides synergistic benefits that validate their effectiveness.\n\n2. Insufficient Experimental Validation\n    * Lack of Comparative Qualitative Visualizations: The paper's quantitative claims are not supported by adequate qualitative evidence. While Appendix G (Fig. 10) shows the proposed model's output, it fails to provide essential side-by-side visual comparisons against the SOTA baselines. Without these direct comparisons, it is impossible to qualitatively assess the model's claimed advantages.\n    * Missing Computational Cost Analysis: The framework introduces significant architectural complexity (e.g., APDS, tLSTMs) without any analysis of the computational overhead. The authors should provide metrics on training/inference time and GPU memory usage. This information is crucial for evaluating the method's practical feasibility and understanding the trade-off between performance and efficiency.\n    * Inadequate Comparison with State-of-the-Art Diffusion Models: The authors cite Diff-UNet [4] in the related work section, which makes their omission from the quantitative comparison a significant weakness. I strongly recommend the authors include comparisons against these recent diffusion-based models [4-6]. Such a comparison is essential to properly benchmark the \"Cross-Timestep\" framework and convincingly validate its claimed segmentation superiority.\n\n3. Clarity and Writing Quality\n    * Discrepancy in Abstract: The abstract explicitly claims: \"Three real-world cases further analyze the collapse phenomenon...\". However, these three specific case studies are not clearly presented in the main paper or appendices. \n    * Minor writing issues: There are duplicate headings (\"2 Related Work\" and \"3 Related Work\") , and scattered typos/formatting issues in Section 5 (e.g., the use of quotes around ’Diff Out’, ’APDS Out’ , ’Conv-tLSTM’, and ’Linear-tGRU’ ).\n\n[1] Chen, Tianrun, et al. \"xlstm-unet can be an effective 2d & 3d medical image segmentation backbone with vision-lstm (vil) better than its mamba counterpart.\" BHI 2024.\n\n[2] Si, Yunzhong, et al. \"SCSA: Exploring the synergistic effects between spatial and channel attention.\" Neurocomputing 634 (2025): 129866.\n\n[3] Jiang, Yuxuan, et al. \"Diff-sfct: A diffusion model with spatial-frequency cross transformer for medical image segmentation.\" BIBM 2023.\n\n[4] Xing, Zhaohu, et al. \"Diff-UNet: A diffusion embedded network for robust 3D medical image segmentation.\" Medical Image Analysis (2025): 103654.\n\n[5] Chen, Tao, et al. \"HiDiff: Hybrid diffusion framework for medical image segmentation.\" IEEE TMI (2024).\n\n[6] Shuai, Zhihao, et al. \"Diffseg: a segmentation model for skin lesions based on diffusion difference.\" arXiv preprint arXiv:2404.16474 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uF2U1U4ML8", "reviewer": "Reviewer_dE8V", "summary": "This paper proposes Cross-Timestep, a novel diffusion model framework for 3D medical image segmentation. The work aims to resolve the \"initial-stage collapse\" issue frequently observed in 3D diffusion models during early denoising steps. To this end, the authors introduce two main components: a Trans-Temporal Memory LSTM (tLSTM) to accumulate structural information across diffusion time steps, and an Adaptive Priori Decoding Strategy (APDS) to stabilize the reverse process using a weighted prior. Although the method achieves competitive performance on some medical datasets, the overall framework is overly complex, lacks adequate demonstration of flexibility and generalizability, and fails to provide sufficient experimental analysis and key visualizations to justify its high design and computational cost.", "questions": "1. Could the authors provide a detailed report on GPU memory usage and training time for the basic 3D diffusion model versus the full Cross-Timestep model to quantify the specific computational overhead introduced by the tLSTM module?\n\n2. Considering the complexity of APDS and tLSTM, have the authors explored a more simplified and flexible design (e.g., using only APDS or a simpler cross-step attention mechanism)? How does the performance of these simpler variants compare to the full model?\n\n3. Can the authors provide a more convincing visual analysis that clearly demonstrates how the memory information accumulated by tLSTM and the prior guided by APDS specifically suppress or reverse the \"initial-stage collapse\" during the early denoising time steps?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper addresses a practical and significant pain point in applying 3D diffusion models to the medical domain, specifically the \"initial-stage collapse,\" which is a valuable research direction in volumetric medical data processing.\n\n2. The combination of tLSTM and APDS reflects a novel approach to addressing time-step dependency and denoising stability. Conceptually, tLSTM, acting as a memory mechanism across time steps, has some inherent reasoning.", "weaknesses": "1. The proposed Cross-Timestep framework is excessively complex and potentially redundant. Integrating the tLSTM module inside the 3D U-Net significantly increases the model's parameter count and computational complexity, while also introducing multiple new hyperparameters. This over-engineered design reduces the framework's generality and flexibility for deployment.\n\n2. Given the extremely high resource demands of the 3D U-Net itself, the additional memory and training time overhead introduced by the tLSTM module are not adequately and rigorously quantified or justified. The paper fails to clearly demonstrate that the performance gain is worth such a substantial increase in design complexity and computational cost.\n\n3. The experimental results rely heavily on quantitative metrics like the Dice Score in tables. Crucial visual evidence and generalization analysis are lacking, especially a direct visualization of how tLSTM and APDS internally mitigate the \"collapse\" across early time steps, which weakens the conviction of the core claim.\n\n4. Although an ablation study is provided, the discussion regarding the strict necessity of both tLSTM and APDS is not deep enough. The authors should explore whether simpler, lighter mechanisms could achieve similar stabilization effects to justify the complexity of the current design.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UoO6iqUmi9", "reviewer": "Reviewer_oCBU", "summary": "This paper introduces Cross-Timestep, a novel 3D diffusion framework designed to address the challenge of initial-stage collapse in 3D medical image segmentation. The authors propose two key components: 1) Adaptive Priori Decoding Strategy (APDS), which provides time-weighted structural guidance during high-noise stages of the diffusion process; 2) Trans-temporal memory LSTM (tLSTM), a recurrent module that maintains temporal state across denoising steps to ensure coherence and stability. The method is evaluated on two multi-center 3D medical imaging datasets (LNCTVSeg and OAseg) and demonstrates state-of-the-art performance in terms of Dice, IoU, and HD95 metrics. The paper also includes extensive ablation studies and visualizations to validate the contributions of each component.", "questions": "1. How does the performance of Cross-Timestep scale with the number of diffusion steps? Is there a trade-off between segmentation accuracy and inference speed?\n2. Could the tLSTM mechanism be adapted to other generative or discriminative tasks beyond segmentation, such as 3D image synthesis or reconstruction?\n3. The APDS relies on a prior decoder trained only on the conditional branch. How sensitive is the model to the quality of this prior? What happens if the prior is noisy or inaccurate?\n4. Have the authors considered using learned or adaptive schedules for $\\omega_t$ instead of the fixed exponential decay? Could this further improve performance?\n5. The paper mentions that APDS prevents over-interference in later stages. Is there a risk of under-guiding in early stages if the prior is too weak?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. **Significance:** The paper identifies and addresses a critical failure mode (initial-stage collapse) in 3D diffusion models, which has been largely overlooked in prior work.\n2. **Technical Rigor:** The proposed APDS and tLSTM modules are carefully designed and thoroughly evaluated. Multiple variants of tLSTM (Conv-tLSTM, Linear-tGRU, SC-tLSTM, FFT-tLSTM) are introduced, showcasing flexibility and scalability.\n3. **Comprehensive Evaluation:** Experiments are conducted on two challenging multi-center datasets with significant domain shifts. The method outperforms several strong baselines and is supported by both quantitative and qualitative results.\n4. **Reproducibility:** The paper includes detailed methodology, ablation studies, and references to publicly available datasets and code, enhancing reproducibility.", "weaknesses": "1. **Computational Cost:** While the authors propose Linear-tGRU to reduce computational demands, the overall framework (especially with 3D convolutions and recurrent units) is likely still resource-intensive. A more detailed analysis of training/inference time and memory usage would be helpful.\n2. The design of the time-weighting function $\\omega_t$ in APDS is heuristic. While it works well, a more systematic analysis or learning-based approach for $\\omega_t$ could further strengthen the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ejEeskrJwr", "reviewer": "Reviewer_1Gpu", "summary": "The paper proposes a 3D diffusion framework for medical image segmentation that addresses the instability of conventional diffusion models when applied to volumetric data. It introduces two key components: the Adaptive Priori Decoding Strategy (APDS), which provides structural guidance to prevent early-stage collapse during high-noise timesteps, and the trans-temporal memory LSTM (tLSTM), which preserves temporal information across diffusion steps for consistent refinement. Together, these innovations stabilize the reverse diffusion process and improve segmentation accuracy across heterogeneous datasets, achieving state-of-the-art results and demonstrating strong robustness under domain shifts.", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Enhanced temporal coherence: The trans-temporal memory LSTM (tLSTM) explicitly retains and propagates structural and contextual information across diffusion timesteps, turning each step into a coherent continuation rather than an independent reconstruction.\n\nImproved stability under high noise: By combining tLSTM with the Adaptive Priori Decoding Strategy (APDS), the model accumulates temporal evidence effectively, preventing early-stage collapse and ensuring reliable recovery from highly noisy initial states.", "weaknesses": "What is the reason for the initial-stage collapse? Why are 2D models utilized in a 3D scenario?\nSome methods, such as Diff-UNet, which use 3D diffusion models, can avoid the initial-stage collapse — is that correct?\nSometimes, the authors describe their methods in a confusing and complex way. For example, what does “explicit temporal evidence accumulation” mean?\nThere are no diffusion-based methods included in the comparison experiments — why?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)