Technical Summary: GRASP – Guided Residual Adapters with Sample-wise Partitioning

Problem Statement

Text-to-image flow matching transformers exhibit severe performance degradation in long-tail distributions, particularly in domains like medical imaging where rare pathologies dominate real-world scenarios. While models perform well on frequent "head" classes, they suffer from mode collapse on "tail" classes, generating images with poor fidelity and limited diversity. The authors identify the root cause as an optimization-level pathology: during fine-tuning, the gradient updates for head and tail classes are mismatched in both norm and direction. These gradients average destructively at the parameter level, causing the model to prioritize frequent classes and collapse the representation of rare ones. Existing interventions, such as post-hoc sampling strategies or training-time loss modifications, fail to address this specific gradient misalignment, often leaving the underlying distribution unchanged or modifying the generative objective in ways that do not directly target the optimization failure.

Methodology

The paper proposes GRASP (Guided Residual Adapters with Sample-wise Partitioning), a non-invasive architectural intervention designed to improve gradient alignment without altering the flow-matching objective or the sampling process.

1. Sample-wise Partitioning via Conditioning

Instead of learning a gating network (which risks being dominated by head-class gradients), GRASP employs a static, deterministic partition of the conditioning space.

• Theoretical Basis: In conditional flow matching, condition values index distinct sets of probability paths. The authors argue that partitioning along the conditioning vector is the structurally correct factorization to serve as a proxy for gradient alignment.
• Implementation: A static routing function $\pi: \mathcal{C} \to \{1, \dots, K\}$ assigns each sample to one of $K$ experts based on its condition.
  • Medical Context (MIMIC-CXR-LT): Samples are partitioned into semantic groups (Head, Medium, Tail) and further split by a domain-specific prior (Healthy vs. Pathological), resulting in four experts.
  • General Context (ImageNet-LT): Partitioning is achieved via text-embedding clustering (bisecting K-means) on the frozen text encoder, without domain-specific priors.
• Robustness: The method relies on "coarse" priors (e.g., frequency bins or vanilla clustering) rather than perfect semantic separation. The goal is to ensure that, on average, samples within an expert share more homogeneous gradient dynamics than those across experts.

2. Guided Residual Adapters

Once samples are routed, the architecture utilizes group-specific residual adapters inserted into the feedforward layers of the frozen pre-trained backbone (Flux).

• Mechanism: For a sample $x_i$ routed to expert $k = \pi(c_i)$, the layer output is $h_{l+1} = F_l(h_l) + A_{k,l}(h_l)$, where $F_l$ is the shared base transformation and $A_{k,l}$ is the expert-specific adapter.
• Gradient Isolation: Parameters within adapter $k$ are updated only by samples assigned to that expert. This isolates group-specific gradients, preventing destructive interference from head-class updates while allowing the shared backbone to accumulate transferable structure.
• Design: The adapters are lightweight two-layer transformations. Unlike LoRA, they do not enforce strict dimensionality reduction, allowing for wider adapters ($d' \approx d$ or $d' > d$) to stabilize fine-tuning across large domain gaps.

Key Contributions

1. Transfer of Gradient Consistency to Flow Matching: The authors adapt the gradient-consistency objective from long-tail language modeling (Dong et al.) to text-to-image flow matching. Crucially, they replace the learned gating mechanism with a deterministic, conditioning-driven partition. They argue that in conditional flow matching, conditioning naturally indexes probability path geometry, making static partitioning a structurally correct factorization rather than a generic heuristic.
2. Non-Invasive and Composable Design: GRASP leaves the flow-matching loss and the sampler untouched, modifying only which parameters absorb which gradients. This allows it to be composed with orthogonal interventions. The authors demonstrate that combining GRASP with self-guided minority sampling at inference time yields superior results compared to either method alone.
3. Empirical Validation of Conditioning as a Gradient Proxy: Through ablation studies and gradient analysis, the paper validates that conditioning homogeneity correlates strongly with expert-wise gradient consistency, confirming that static partitions effectively approximate the ideal gradient grouping even with noisy priors.

Results

The method was evaluated on MIMIC-CXR-LT (medical chest X-rays), NIH-CXR-LT (downstream classification), and ImageNet-LT (general long-tail generation).

• Image Generation Quality & Diversity:
  • On MIMIC-CXR-LT, GRASP reduced the overall FID by up to 80% compared to full fine-tuning and increased tail-class coverage by up to 44%.
  • It outperformed state-of-the-art baselines, including full fine-tuning, generic adapters, and learned-routing Mixture-of-Experts (MoE) variants like MoCLE.
  • Qualitative analysis showed GRASP recovered class-specific features for rare pathologies (e.g., Pneumoperitoneum) that baselines smoothed out or collapsed.
• Downstream Classification:
  • When synthetic images generated by GRASP were used to train a DenseNet-121 classifier on NIH-CXR-LT, the model achieved a macro F1 score matching that of real-data training.
  • Crucially, GRASP enabled nonzero F1 scores on 9 out of 13 classes, whereas full fine-tuning only achieved this on 3 classes.
• Generalization: Results on ImageNet-LT confirmed that the mechanism is not tied to medical inductive biases and generalizes to broad long-tail generation tasks.

Significance and Claims

The paper positions GRASP as a controlled augmentation and analysis tool, explicitly stating that synthetic images are not a substitute for clinically grounded cohorts or expert review. The authors emphasize that synthetic data can preserve or amplify biases if used without provenance and validation.

The primary significance claimed is that static routing is the default choice for long-tail text-to-image generation whenever even coarse priors over the conditioning are available (a common case in medical imaging). Unlike learned gating, which fails in extreme long-tail regimes due to head-class dominance, static partitioning guarantees that every tail sample updates its assigned expert. This provides a stable, efficient, and non-invasive mechanism to mitigate optimization-level pathologies without modifying the core generative objective.