AdaLoRA-QAT: Adaptive Low-Rank and Quantization-Aware Segmentation

1. Problem Statement

The deployment of large foundation models (such as the Segment Anything Model, SAM) for medical image segmentation, specifically Chest X-ray (CXR) lung segmentation, faces significant barriers in clinical settings:

• Computational Constraints: Large models require substantial memory and inference power, making them unsuitable for resource-constrained clinical hardware.
• Generalization Challenges: Existing deep learning models (e.g., nnU-Net, DeepLabV3+) struggle with anatomical variability, pathological distortions, and imaging artifacts.
• Precision vs. Efficiency Trade-off: Reducing model size via quantization often leads to a degradation in structural fidelity and segmentation accuracy, which is unacceptable for diagnostic reliability.

The goal is to develop a framework that achieves high diagnostic accuracy (matching full-precision baselines) while drastically reducing trainable parameters and model size to enable deployment on edge devices.

2. Methodology: AdaLoRA-QAT

The authors propose AdaLoRA-QAT, a two-stage fine-tuning framework that integrates Adaptive Low-Rank Adaptation (AdaLoRA) with Quantization-Aware Training (QAT).

Core Architecture

The method adapts the pre-trained Segment Anything Model (SAM) for binary lung mask prediction using bounding box prompts.

Stage 1: Full-Precision Adaptive Low-Rank Training

• Mechanism: Instead of standard LoRA (which uses a fixed rank), the framework employs AdaLoRA. It models weight updates as a low-rank residual $\Delta W = P \text{diag}(\Lambda \odot m) Q$.
• Adaptive Allocation: It dynamically allocates rank capacity across transformer layers based on task sensitivity.
  • Pruning: Redundant components are pruned based on an importance metric $I_i = |\lambda_i| \cdot |\partial L/\partial \lambda_i|$, retaining only the top-$r_{target}$ components.
  • Orthogonal Regularization: Ensures stable subspace learning and prevents "rank collapse."
• Optimization: Parameters ($P, Q, \Lambda, m$) are optimized in FP32 using a hybrid loss function ($L_{BCE} + L_{Dice} + \lambda_{ortho}L_{ortho}$).

Stage 2: Quantization-Aware Fine-Tuning (QAT)

• Mixed-Precision Strategy: To enable INT8 deployment without losing critical structural information, a selective quantization approach is used:
  • Quantized to INT8: Encoder feed-forward layers, the decoder, and the prompt encoder (using symmetric per-tensor quantization).
  • Retained in FP32: Attention QKV projections and the AdaLoRA parameters ($P, Q, \Lambda$). This is crucial to prevent rank collapse in SVD-parameterized layers and maintain orthogonality.
• Freezing Mechanism: The rank masks ($m$) learned in Stage 1 are frozen. Only the singular values ($\Lambda$) are fine-tuned to compensate for quantization-induced shifts.
• Objective: The model adapts to quantization noise while preserving the learned subspace topology.

3. Key Contributions

1. Unified Framework: Introduction of AdaLoRA-QAT, coupling adaptive low-rank encoder tuning with full-model quantization-aware fine-tuning.
2. Novel Mixed-Precision Design: A strategy that selectively quantizes most components to INT8 while keeping attention mechanisms and low-rank adaptation parameters in FP32. This prevents the degradation of structural fidelity often seen in aggressive quantization.
3. Efficiency and Performance: The method achieves state-of-the-art efficiency, reducing trainable parameters by 16.6× and compressing the model size by 2.24× without statistically significant loss in accuracy.

4. Experimental Results

The framework was evaluated on a large-scale dataset comprising 64,590 CXRs from diverse sources (JSRT, QaTa-COV19, COVID-19 Radiography, etc.).

• Segmentation Accuracy:
  • Achieved a Dice Score of 95.59%, which is statistically indistinguishable from the full-precision SAM Decoder baseline (95.55%) and outperforms traditional models like nnU-Net (94.85%).
  • Statistical Validation: A Wilcoxon signed-rank test confirmed no significant difference ($p > 0.05$) between the proposed method and the baseline across all metrics (Dice, IOU, NSD).
• Parameter Efficiency:
  • Trainable Parameters: Reduced from 3.8M (SAM Decoder FT) to 5.4M (Ours) relative to the base model, but compared to full fine-tuning, it represents a 16.6× reduction in trainable parameters.
  • Model Compression: Achieved 2.24× compression compared to the baseline.
• Structural Fidelity:
  • SSIM (Structural Similarity Index) heatmaps showed stronger agreement along lung boundaries and vascular regions compared to the baseline.
  • Quantization error analysis revealed a zero-mean Gaussian distribution ($\mu \approx 2.7 \times 10^{-5}$), indicating unbiased rounding and preserved numerical fidelity.

5. Significance and Conclusion

AdaLoRA-QAT addresses the critical bottleneck of deploying foundation models in low-resource clinical environments.

• Clinical Relevance: By proving that INT8 quantization does not degrade diagnostic accuracy, the method enables the use of powerful AI tools on standard hospital hardware.
• Scalability: The framework serves as a proof-of-concept for task-agnostic medical image segmentation, extensible to other modalities and organs.
• Future Impact: The work paves the way for scalable, AI-assisted diagnosis in settings with limited computational resources, balancing the triad of accuracy, efficiency, and structural trustworthiness.

The code and pretrained models are publicly available, facilitating further research and clinical integration.