D2SA: Dual-Stage Distribution and Slice Adaptation for Efficient Test-Time Adaptation in MRI Reconstruction

We sincerely thank the reviewer for the detailed summary and constructive comments. We appreciate your recognition of the practical value of our framework, the clarity of our experimental design, and the relevance of addressing domain shifts in medical imaging. Below, we respond to your main concerns regarding the framework, component contributions, hyperparameter sensitivity, and computational overhead, and we provide additional analyses to address the concerns.

Q1:The Method’s Complexity and the Use of Known Components

Conceptual contributions in adapting INR for test-time domain shift: To the best of our knowledge, INRs have not been previously used in test-time adaptation for image reconstruction. We leverage the compact functional expressiveness of INR to learn a new patient-specific representation at test time, not to generate images, but to guide the adaptation of source models in a structure-aware way. This is conceptually distinct from both conventional pretraining-based INR usage (e.g., Functa) and standard slice-wise TTA.

Unified modular architecture with pluggable design: Each component (MR-INR, affine modulation, SST refinement, and AD module) is designed to be lightweight, interchangeable, and orthogonal. This enhances not only performance, but also flexibility and practical adoption. For example:

If compute is limited, users can skip Stage 2 and rely on Stage 1 alone.

If specific slices are flagged by clinicians, only Stage 2 is applied.

A holistic test-time design perspective: Our framework proposes a structured pipeline that mimics how real clinical workflows operate—first broadly adapting to the patient, then selectively refining critical regions. This hierarchical control is not present in prior works and reflects a pragmatic yet theoretically motivated design.

While we do build on existing components, our architectural integration and adaptation strategy are novel, practical, and aligned with real-world test-time usage constraints—thus providing both scientific insight and engineering utility.

Q2:Necessity of Stage 1: Is most benefit from Stage 2?

We provide a detailed ablation comparing three configurations on both UNet and VarNet under anatomy shift, including Stage 2 only (FINE + SST + AD); Stage 1 + Stage 2 without MR-INR (FINE + SST + AD but with pretrained encoder from FINE) and Full D2SA (FINE + MR-INR + SST + AD).

Compared to Stage 2 only (FINE+SST+AD), which uses only FINE features, our Full D2SA introduces the Stage 1 (MR-INR) and consistently boosts PSNR by 0.18~0.4 dB More importantly, inference time is reduced by >30% (UNet) and nearly 3× (VarNet) when MR-INR provides better initialization for SST. This shows that Stage 1 is not redundant: it not only improves image quality but also accelerates slice-level optimization in Stage 2, especially combining MR-INR in the stage 1.

Q3: Hyperparameter sensitivity

To further support the robustness of our method in hyperparameters settings, we show the dedicated robustness results on the key hyperparameters in both Stage 1 (MR-INR) and Stage 2 (SST+AD), including: Latent code size (64 / 128 / 256)；INR network depth (3 / 4 / 5)；Latent code initialization std (0.1 / 0.01 / 0.001)

As shown above, the performance across PSNR / SSIM / LPIPS remains stable, with <0.15 dB variation in PSNR and negligible impact on LPIPS.

We also show the results about different combinations of loss weights for Stage 1: λ_Self / λ_INR / λ_Reg; Stage 2: λ_Self / λ_INR.

We used empirically chosen default settings (Stage 1: λ = (1, 1, 1e-4); Stage 2: λ = (1, 1)). The results remain stable across a range of values, indicating good robustness of our framework. However, we observe that setting λ_INR too low (e.g., 0.1) leads to performance degradation, particularly in Stage 1 and LPIPS in stage 2. This is expected, as MR-INR plays a central role in modeling patient-level distribution shifts. A sufficiently weighted INR loss encourages better representation learning of domain-specific patterns, which is crucial for boosting TTA effectiveness.

Q4:Computational Overhead vs. Gain:

While our method does introduce additional modules (INR and AD), we would like to clarify that we design the D2SA with practical efficiency and controllability in mind:

Frozen base networks reduce trainable parameters: Unlike methods like FINE that update all weights, our approach freezes the base CNN (U-Net/VarNet), and only updates light-weight modules (MR-INR, affine modulation, AD), leading to significantly fewer trainable parameters. This is evident in Table 2, where D2SA consistently requires fewer parameters than full-model tuning.

Stage-wise adaptation reduces convergence time: Although D2SA has more components, it achieves faster and more stable convergence. As shown in Figure 4 (main) and Figure 7 (supplement), inference time is reduced at least 60% compared to FINE+SST or DIP-TTT under the same early stopping setup.

Efficient initialization for refinement: Importantly, Stage 1 acts as a structured, patient-aware initialization for Stage 2, providing prior distribution alignment and faster convergence for slice-level refinement. This makes Stage 2 more effective and efficient, especially for structurally important slices.

In summary, D2SA offers a favorable trade-off: while FLOPs do increase, our design leads to lower parameter counts, faster convergence, and more robust reconstructions, making it more practical for clinical deployment than existing monolithic TTA approaches.

Thank you again for your thoughtful review. We would greatly appreciate it if you could share any further thoughts regarding our rebuttal. Your feedback would be invaluable in helping us improve the final version.

We thank the reviewer for the thoughtful suggestions. Your recognition of our dual-stage framework’s effectiveness and efficiency is greatly appreciated. Below, we address your comments regarding the clarity of architectural differences, the distinctions from diffusion-based test-time adaptation methods, and hyperparameter sensitivity.

Q1: Architectural Differences

Beyond architectural differences, our two-stage D2SA framework is designed with realistic clinical deployment in mind, where direct zero-shot transfer or pure slice-wise refinement often fails to provide reliable reconstructions—especially under large distribution shifts. This unreliability can directly impact early-stage screening and the localization of small or subtle lesions.

In clinical scenarios, radiologists often need initial high-quality reconstructions across all slices for efficient review, followed by selective refinement on suspicious or diagnostically relevant slices.

To this end:

Stage 1 (MR-INR) performs a fast, patient-level adaptation, leveraging inter-slice structure to learn global distributional shifts (mean, variance). It provides a reliable reconstruction baseline within ~5–7 minutes, without requiring slice-wise optimization.

Stage 2 (SST + AD module) then allows targeted refinement on demand, enhancing structural fidelity on specific slices without re-training the full network. The AD module preserves fine-grained anatomical details (e.g., tissue boundaries), which are often over-smoothed by conventional TTA methods.

Each component is deliberately positioned:

MR-INR enables compact patient-specific modeling, aligned with the notion of shared anatomical priors.

Affine adaptation (α, β) efficiently compensates for patient-specific mean/variance shifts with provable benefits (Section 4.3).

The AD module, integrated into SST, provides structure-aware filtering under frozen CNNs, striking a balance between efficiency and precision.

Q2: The distinctions from diffusion-based test-time adaptation methods

Unlike diffusion-based TTA approaches that require heavy pretraining and offer limited flexibility, our method is pluggable, interpretable, and clinically actionable. This hierarchical strategy is not merely a module combination but a functionally structured design tailored for practical, slice-selective, low-resource test-time adaptation in real-world medical workflows.

Q3: Sensitivity about hyperparameter on latent code and AD

To clarify robustness of our INR-based adaptation, we show additional sensitivity experiments on three critical hyperparameters: latent code size, INR network depth, and latent code initialization scale. As shown below, our method exhibits stable performance across a wide range of values:

Our default configuration (highlighted in bold: 128-d latent code, depth 4, init std 0.01) achieves a strong trade-off between stability and performance. Performance is consistent across hyperparameter ranges, with differences ≤ 0.02 PSNR or ≤ 0.01 LPIPS. This suggests our method is not overly sensitive and avoids instability even with pluggable configurations. The default setup ensures robust adaptation without excessive tuning overhead.

Moreover, to show the robustness of our method to the anisotropic diffusion (AD) step size, we conducted ablation experiments under modality shift (UNet with FINE backbone). As shown below, varying the AD step size from 1.0 to 0.01 leads to minimal changes in performance:

These results demonstrate that our method is robust to the choice of the diffusion step size, with less than 0.4 dB PSNR fluctuation and negligible variation in perceptual quality. The default choice of step=1.0 achieves the best trade-off between sharpness and stability.

We truly appreciate your initial feedback. If you have any further comments on our rebuttal, we’d be grateful to hear them during the discussion phase. Thank you again for your time and consideration.

We sincerely thank the reviewer for the detailed and constructive feedback. Your comments have greatly helped us identify key areas for the clarification of our method and results. For the spelling/grammar error, we will correct them in final version. Below, we address your main concerns.

Q1: Clarity Concerns

In the Introduction (Lines 47–70) and the beginning of Section 3 (Method), we summarize the motivation and high-level design of D2SA, highlighting its two-stage structure:

Stage 1: Patient-wise Distribution Adaptation – We use an MR-specific Implicit Neural Representation (MR-INR) to model shared representations across slices. Each slice’s coordinates are extended with a patient-specific latent code (z_p) and random Fourier features to stabilize convergence and express frequency shifts. These inputs are passed through a SIREN network to learn mean/variance modulation parameters (α,β), allowing distribution alignment for all slices in a patient.

Stage 2: Slice-wise Structural Refinement – For each slice, we freeze the encoder and patient latent code and fine-tune only the INR decoder and a lightweight, learnable Anisotropic Diffusion (AD) module. The AD module extracts directional gradients using difference convolutions (see Line 171–180) and generates adaptive edge-aware refinement weights.

Details of the AD module in Section 4.2&4.3: Regarding Section 4.2, due to space constraints, we initially omitted the explicit formulation of the five directional difference convolutions used to estimate gradient-aware anisotropic diffusion coefficients. These convolutions extract directional derivatives (horizontal, vertical, diagonal, etc.) with end-to-end learning machanism. For example, taking the Horizontal Difference Convolution (HDC) as a case: the horizontal gradient is first computed via differences between pixel pairs along the x-axis. After training, the learned kernel weights are re-arranged and applied directly as convolution filters on the input feature map. These filters exhibit properties similar to handcrafted operators (e.g., their horizontal weights sum to zero), enabling structural-aware modulation. We will briefly summarize this process in the main paper. For Section 4.3, the theoretical derivation of the LSS-based optimization is given in Appendix Section 5; we will highlight this in the main text.

Clarifying Lines 124–128: We rephrase it as follows for better clarity: “Following patient-level batch training, each slice is associated with a latent identity code. In the MR-INR branch, this latent code is concatenated with Fourier-encoded spatial coordinates before being passed to a SIREN network. This combination enhances convergence and representation power by encoding slice-specific anatomical priors and capturing periodic signals in MRI signal acquisition.”

Figure 4 and UNet result：To avoid table overcrowding and emphasize key trends, we visualized UNet results as charts. These highlight that: (i) MR-INR provides the largest PSNR gain under modality shift (top-left); (ii) our full method outperforms baselines with lower LPIPS and faster inference (bottom-left); (iii) MR-INR and SST consistently improve SSIM across all domain shifts (right). Full UNet tables are in Appendix Table 1, and VarNet charts in Appendix Fig. 1. We will move UNet table to main page and keep all charts in Appendix in the final version.

Q2: Ablation studies on VarNet

We extend the full-stage ablations to VarNet under the anatomy shift setting, including additional baselines including SST-only+AD and FINE+SST+AD. The VarNet part will be added into final version.

The extended ablation study highlights the functional contributions and trade-offs of each D2SA component across VarNet backbones. We observe that Stage 2 (SST+AD) with/without stage1 training offer reasonable performance, but at the cost of significantly higher inference time (~50–56 mins per patient). Incorporating Stage 1 (MR-INR) leads to stronger initialization for slice-wise refinement, reducing convergence time substantially (up to 3× faster inference) and yielding consistent performance gains. Notably, our final configuration (MR-INR (🧊 latent) + SST (🧊 CNN + 🔥 AD) achieves the best PSNR while maintaining low parameter count and the fastest inference time (e.g., 17.1 min/patient for VarNet). These results validate our core claim: D2SA’s staged modularity is not only conceptually sound, but also practically effective, enabling scalable, fast, and adaptive inference with minimal tuning. Each module remains plug-and-play, ensuring generalizability and ease of adoption in various architectures and deployment scenarios.

Q3:Choice of SIREN over Other INRs

To clarify our choice of SIREN as the INR backbone in MR-INR, we conducted an additional ablation comparing different INR backbones under both patient-wise (FINE+MR-INR) and single-slice (FINE+MR-INR+SST) training. Results are summarized below:

Our design combining random Fourier positional encoding with SIREN consistently yields the best or near-best performance, especially in perceptual quality (LPIPS). While Finer offers frequency adaptability, its more complex architecture increases overhead without clear performance gains; WIRE, though efficient, suffers from unstable convergence and limited representation in medical shift settings. Compared to INCODE, our SIREN-based MR-INR avoids the need for pretraining, supporting our lightweight, plug-and-play adaptation objective.

We will include these results in the final paper.

Q4:Result Significance and Efficiency

We thank the reviewer for the helpful suggestion. We would like to clarify statistical significance tests are already included in Supplementary Figure 8, where we compare our method (SSDU+MR-INR+SST) with DIP-TTT and +SST using Wilcoxon signed-rank tests. Under UNet (acceleration shift) and VarNet (sampling shift), we observe significant improvements in all metrics (p value < 0.05).

We also would like to highlight efficiency gains, even in cases with close average scores. For example:

In VarNet on anatomy shift (Table 1), our method improves SSIM from 0.878 → 0.882 and PSNR from 27.67 → 27.68, while reducing inference time from 52.5 min to 17.1 min, a 67% reduction.

In U-Net on sampling shift (Table 3), our method achieves LPIPS 0.311 vs. 0.314 (DIP-TTT) with 13.2 min vs. 27.1 min, yielding a>50% faster adaptation while maintaining superior quality.