PaperlayerView exact revision on arXiv

Source-linked AI summary

MARR: Decoupling Policy, Execution, and Calibration for All-in-One Medical Image Restoration

Haobin Chen, Ao Chang, Heqin Zhu, Rundong Wang, Ting Liu, Shaohua Kevin Zhou

arXiv:2609.06645v1cs.CV

TL;DR

All-in-one medical restoration must address distinct PET, CT, and MRI characteristics without relying on separate modality-specific models. MARR assigns policy routing, modality-private residual refinement, and output calibration distinct roles; it improves modality-balanced metrics over thirteen methods re-trained under the same protocol.

  • Problem

    PET, CT, and MRI differ in degradation, anatomical contrast, and output-space bias, creating a challenge for unified restoration that preserves shared priors without separate modality-specific models.

    Source: §I · p. 1

  • Method

    MARR separates degradation-aware prompt policy selection, prompt-gated modality-private feature refinement, and zero-initialized modality-specific output residual calibration.

    Source: §I · pp. 1–2, E. Zero-Initialized Output Calibration · p. 4

  • Results

    MARR improves modality-balanced PSNR, SSIM, and RMSE over thirteen methods re-trained under the same protocol; average PSNR is 34.43 dB.

    Source: D. Main Results · p. 5, §V · p. 8

  • Takeaways & Limitations

    MARR provides a compact way to balance common restoration priors with modality-specific residual allocation in one model.

    Source: §V · p. 8

    The evaluation uses full-reference restoration metrics that do not directly measure downstream clinical utility.

    Source: H. Limitations · p. 7

Abstract

from arXiv · show

All-in-one medical image restoration seeks to recover heterogeneous clinical images with a single model, but PET, CT, and MRI differ substantially in degradation statistics, anatomical contrast, and output-space bias. A fully shared network can entangle modality-specific residual errors, whereas separate modality-specific networks sacrifice the practical advantages of unified deployment. We therefore recast all-in-one restoration as a question of where limited adaptation should be placed: policy selection, feature execution, or output calibration. We propose MARR, a compact restoration framework that constrains multi-modality adaptation into degradation-aware policy routing, modality-private residual execution, and image-domain residual correction without requiring degradation labels or separate modality-specific models. The policy branch forms a routing prompt from input statistics, latent content, and modality identity, and uses it only as a control signal. Prompt-gated modality-private adapters then perform lightweight residual refinement at intermediate decoder stages, while zero-initialized modality-specific output heads calibrate the final image-domain residual without perturbing the initial shared prediction. On an all-in-one PET, CT, and MRI restoration benchmark, MARR outperforms thirteen methods re-trained under the same protocol, achieving PSNR values of 37.34 dB, 33.85 dB, and 32.09 dB on PET, CT, and MRI, respectively, and the best modality-average PSNR of 34.43 dB. The code is publicly available at https://github.com/CHB-learner/MARR.

I. INTRODUCTION

All-in-one medical restoration must accommodate distinct PET, CT, and MRI restoration needs while retaining the deployment advantages of a shared model. MARR assigns policy selection, modality-specific feature refinement, and output correction separate roles.

  • PET, CT, and MRI have distinct restoration failure modes, while separate modality-specific models increase validation and maintenance costs and complicate deployment.PET must suppress low-count noise without erasing weak uptake; CT must remove structured artifacts while preserving calibrated intensity; MRI must recover anatomical edges and tissue contrast.Source: §I · p. 1
  • Prior all-in-one methods use task routing, adaptive attention, or degradation-aware prompts, but may assign one prompt selection, feature modulation, and final correction together.The cited methods include –.Source: §I · p. 1
  • MARR estimates a degradation-aware policy prompt, applies prompt-gated modality-private decoder adapters, and uses modality-specific output heads for image-space residual calibration.The framework keeps common restoration priors shared while assigning explicit roles to decision, execution, and correction.Source: §I · pp. 1–2
  • MARR uses image statistics, latent content, and modality identity to guide restoration while preserving shared anatomical priors.Source: §I · p. 2
  • MARR reports consistent gains over thirteen representative methods on PET synthesis, CT denoising, and MRI super-resolution under a unified evaluation protocol.The comparison covers all three restoration tasks.Source: §I · p. 2

II. RELATED WORK

The related work spans generic restoration backbones, modality-specific and all-in-one medical restoration, and lightweight prompting or adaptation. MARR distinguishes its approach by assigning these mechanisms constrained roles within one model.

A. Generic Image Restoration

  • Generic restoration work includes transformer, activation-free, state-space, and prompt-learning approaches, but most cited methods do not explicitly separate modality routing, feature specialization, and output calibration.The passage cites SwinIR, Restormer, NAFNet, PromptIR, and other models –.Source: A. Generic Image Restoration · p. 2

B. All-in-One Medical Image Restoration

  • Medical restoration includes modality-specific PET, CT, and MRI methods and all-in-one approaches such as AMIR, RAT, TAT, and DaPT.The cited work also explores stronger priors for all-in-one restoration and synthesis –.Source: B. All-in-One Medical Image Restoration · p. 2

C. Prompting and Lightweight Adaptation

  • MARR constrains prompts to control, private adapters to feature-space residuals, and output heads to image-domain bias correction, making the shared-private mechanism auditable.Related lightweight adaptation work includes prompts and other mechanisms –[31].Source: C. Prompting and Lightweight Adaptation · p. 2

III. METHOD

MARR formulates unified PET, CT, and MRI restoration as a shared-model problem with modality input. Its design separates policy, private feature refinement, and final image correction in response to differing modality residuals.

A. Problem Formulation and Overview

  • MARR uses one restoration model for degraded images and their modality, with a policy prompt controlling behavior, private adapters refining decoder features, and output heads correcting residual bias.The three paths are decision, execution, and correction.Source: A. Problem Formulation and Overview · p. 2, Figure 1 · p. 3

B. Design Rationale

  • PET, CT, and MRI differ in noise, intensity bias, artifacts, and anatomical texture, so MARR keeps common priors shared while adding lightweight private capacity for modality-specific residuals.The design constrains each path to a limited error type to reduce conflation of task selection, feature specialization, and image-domain correction.Source: B. Design Rationale · p. 3

C. Degradation-Aware Prompt Policy

MARR builds a routing prompt from image degradation statistics, modality identity, and latent content, then uses it to control residual adaptation rather than inject image features. Its lightweight policy regularizer encourages prompt diversity and discourages premature routing collapse.

C. Degradation-Aware Prompt Policy

D. Prompt-Gated Modality-Private Residual Adaptation

MARR uses prompt-gated, modality-private residual adapters at intermediate decoder stages to refine features while retaining a shared restoration representation. Their placement balances access to contextual and spatial detail against the risks of overly early or output-only adaptation.

D. Prompt-Gated Modality-Private Residual Adaptation

  • At selected decoder levels, a prompt-dependent gate controls a lightweight modality-specific residual adapter applied to the decoder feature.Each branch uses depthwise convolution, GELU, and pointwise convolution; its final pointwise convolution is zero-initialized and its residual scale is small.Source: D. Prompt-Gated Modality-Private… · p. 4
  • The shared trunk learns modality-agnostic priors while private adapters absorb modality-specific residual corrections through a conservative, lightweight residual path.Source: D. Prompt-Gated Modality-Private… · p. 4
  • Intermediate decoder stages provide bottleneck context while retaining spatial resolution for local refinement, unlike private modules placed too early or only at the output.Source: D. Prompt-Gated Modality-Private… · p. 4

E. Zero-Initialized Output Calibration

MARR adds modality-specific output calibration to address residual bias that can remain after feature-level adaptation. Zero initialization leaves the initial shared prediction unchanged while the heads learn image-domain corrections.

E. Zero-Initialized Output Calibration

F. Training Objective

MARR training uses a modality-balanced reconstruction objective with policy regularization, and selects models using a validation score that accounts for every modality. This validation choice addresses the risk that MRI-heavy sampling can obscure PET or CT degradation.

F. Training Objective

  • MARR is trained with a modality-balanced reconstruction objective and a policy regularizer.Source: F. Training Objective · p. 4
  • The policy regularization weight is small, keeping the reconstruction objective dominant while stabilizing prompt routing.Source: F. Training Objective · p. 5
  • Model selection uses a modality-aware validation score because the MRI-dominated test set can make a sample-weighted criterion hide PET or CT degradation.Source: F. Training Objective · p. 5

IV. EXPERIMENTS

The evaluation uses paired low-quality and high-quality PET, CT, and MRI slices, with PSNR, SSIM, and RMSE. Thirteen methods are retrained and evaluated under a shared protocol.

A. Dataset and Evaluation Metrics

  • The benchmark pairs low-quality and high-quality slices across PET, CT, and MRI, and reports PSNR, SSIM, and RMSE.Higher PSNR and SSIM are better; lower RMSE is better.Source: A. Dataset and Evaluation Metrics · p. 5
  • Because test-set sizes differ by modality, Table I’s arithmetic average reflects modality-balanced PET, CT, and MRI performance.The average is not dominated by the largest test subset.Source: A. Dataset and Evaluation Metrics · p. 5

B. Comparison Methods and Implementation

  • Thirteen representative methods are retrained and evaluated on the same PET/CT/MRI split, making Table I a reproduced-protocol comparison.The table excludes separately reported or released numbers that may use different training schedules or checkpoint-selection rules.Source: B. Comparison Methods and Implementation · p. 5
  • Each method’s reported checkpoint is selected by sample-weighted average PSNR, and all metrics use the same full-test restored outputs and targets.Latest and best-validation checkpoints are evaluated when available.Source: B. Comparison Methods and Implementation · p. 5

C. Implementation Details

MARR achieves the strongest reported results across PET, CT, and MRI in the reproduced comparison. Diagnostics show modality-structured prompts, stronger CT/MRI feature corrections, and the largest output calibration shifts for PET.

C. Implementation Details

  • MARR uses uniform modality sampling, batches of five per modality, and 128 × 128 random crops for training.Training uses AdamW with learning rate 2 × 10−4 and a maximum 250k-iteration schedule; the final tested checkpoint was selected at 26,000 iterations.Source: C. Implementation Details · p. 5
  • Comparison methods are trained independently, and reporting each method’s stronger checkpoint reduces the risk of penalizing baselines for checkpoint timing.All use the same data paths and split files.Source: C. Implementation Details · p. 5

D. Main Results

  • MARR leads all three reproduced modality comparisons, raising PET PSNR from 37.25 to 37.34 dB, CT from 33.79 to 33.85 dB, MRI from 32.05 to 32.09 dB, and modality-average PSNR from 34.36 to 34.43 dB.Average SSIM improves from 0.9357 to 0.9361, and average RMSE decreases from 12.4940 to 12.4811.Source: D. Main Results · p. 5

E. Mechanism Diagnostics

  • Prompt embeddings form modality-aware clusters, with a silhouette score of 0.83; adapter residuals peak at decoder level 2 for CT and MRI, while output calibration shifts are largest for PET.The diagnostics are descriptive rather than causal.Source: E. Mechanism Diagnostics · p. 5

F. Component Analysis

The component analysis finds that feature adaptation and output calibration address complementary modality-specific residuals, while the discussion presents their compact combination as a practical all-in-one trade-off.

F. Component Analysis

  • Combining private adaptation and output calibration gives the best MRI score and modality-average PSNR, while extended variants do not surpass compact MARR and can lower the balanced average.Output calibration alone is strong for PET and CT, whereas private adaptation improves MRI over calibration alone; the margin over the strongest single-role variants is small.Source: F. Component Analysis · p. 6, Table II · p. 6

G. Discussion

  • Diagnostics associate stronger level-2 feature corrections with CT and MRI, but larger output-calibration shifts with PET; no single mechanism dominates across modalities.The diagnostics are consistent with calibration contributing more to PET/CT and feature adaptation being more relevant to MRI.Source: G. Discussion · p. 6, Figure 4 · p. 7
  • Visual comparisons place LQ input, representative methods, MARR, and HQ target in each row, with matched zoom-in regions.The layout supports direct visual comparison across the three modalities.Source: Figure 3 · p. 7
  • Frequency and anchor variants can improve MRI while lowering PET or CT enough to reduce the balanced average, making compact MARR a practical deployment compromise.Table I’s average is the arithmetic mean over PET, CT, and MRI.Source: Table I · p. 6, G. Discussion · p. 7

H. Limitations

MARR improves modality-balanced restoration results, while the authors identify limits in the evaluation scope and supporting evidence. The conclusion frames its gains as arising from complementary adaptation roles.

H. Limitations

  • The evaluation uses full-reference restoration metrics that do not directly measure downstream clinical utility.The authors call for evaluating effects on segmentation, quantification, and diagnosis.Source: H. Limitations · p. 7
  • MARR uses modality labels during training and testing; fully automatic modality inference could improve robustness to incomplete metadata.Source: H. Limitations · p. 7
  • The decomposition is supported by ablations and diagnostics, not a formal domain-shift bound; it uses statistic cues rather than a learned uncertainty estimator.The authors suggest uncertainty-aware calibration may improve reliability in difficult cases.Source: H. Limitations · p. 7

V. CONCLUSION

  • MARR organizes unified PET, CT, and MRI restoration into degradation-aware policy selection, modality-private residual adaptation, and zero-initialized output calibration.Source: §V · pp. 7–8
  • MARR improves modality-balanced PSNR, SSIM, and RMSE over thirteen methods re-trained under the same protocol.Source: §V · p. 8
  • The authors suggest the gains reflect complementary adaptation roles rather than one uniformly dominant module, and identify additional modalities and uncertainty-aware calibration as future work.Source: §V · p. 8
Loading 2609.06645v1…