P-GRAFT: Partial Generalized Rejection Sampling Fine-Tuning

• P-GRAFT is a framework that leverages partial generalized rejection sampling to fine-tune models, integrating multi-task deep learning for survival analysis and RL-inspired reward shaping for diffusion models.
• In survival analysis, it employs a deep neural network with dual heads to predict event rankings and survival curves, optimizing a combined Cox partial likelihood and isotonic-regression loss.
• For diffusion models, it refines intermediate noise levels via reward-based fine-tuning, achieving notable performance gains across tasks like text-to-image synthesis and molecule generation.

P-GRAFT (Partial Generalized Rejection Sampling Fine-Tuning) denotes a family of methods for distribution shaping and fine-tuning in probabilistic generative models, most notably for diffusion models and in the context of multi-task learning for individualized survival analysis. The term encompasses two distinct frameworks in recent literature: (1) a multi-task deep learning approach for survival curve prediction in kidney graft outcome modeling (Luck et al., 2017), and (2) an RL-inspired, reward-based fine-tuning method that shapes intermediate distributions in diffusion models, extending and generalizing rejection sampling fine-tuning (GRAFT) (Anil et al., 3 Oct 2025). Both frameworks aim to enable individualized, robust predictions by optimizing appropriate objective functions under complex, high-dimensional data distributions.

1. Formal Definitions and Frameworks

1.1 P-GRAFT for Survival Analysis

In personalized survival analysis, P-GRAFT describes a multi-task deep neural network architecture. Given preprocessed patient or donor–recipient features $x\in\mathbb R^D$, the model comprises:

• Shared representation: $h = \phi(x)$, traversing several fully connected layers with batch normalization, ReLU activation, dropout, and weight regularization.
• Rank prediction head: $s^{(1)}(x) = w_r^\top h + b_r$, producing an event risk score for ranking under the Cox partial log-likelihood.
• Survival curve prediction head: $s^{(2)}(x) = \sigma(W_s h + b_s) \in [0,1]^T$, directly estimating discrete time survival probabilities $\hat S(t|x) \approx P(T \geq t \mid x)$.

The network is trained via a weighted sum of Cox partial likelihood and isotonic-regression-based time-to-event ranking losses, with standard regularizers.

1.2 P-GRAFT for Diffusion Model Fine-tuning

In the generative modeling context, P-GRAFT refers to Partial Generalized Rejection Sampling, targeting intermediate noise levels in reverse diffusion processes to steer the learned data distribution toward reward-optimality. Letting $\bar p_{T:0}$ denote a pretrained reference diffusion process ($N$ timesteps), and $r(x_0)$ a scalar reward on denoised samples, the algorithm proceeds by:

• Simulating $M$ full denoising trajectories under $\bar p_{T:0}$.
• Computing rewards $h = \phi(x)$0 on final samples $h = \phi(x)$1.
• Accepting each trajectory with probability determined by an acceptance function $h = \phi(x)$2.
• Collecting the partially denoised states $h = \phi(x)$3 (intermediate noise).
• Fine-tuning model parameters for steps $h = \phi(x)$4 only, using the accepted trajectories.

The resulting fine-tuned model is stitched with the reference model at test time, using the refined early steps and original late-stage denoising.

2. Objectives and Loss Formulations

2.1 Multi-Task Loss in Survival Analysis

The total loss is:

$h = \phi(x)$5

where $h = \phi(x)$6 trades off Cox-style risk ranking and survival curve accuracy:

• $h = \phi(x)$7: Cox partial log-likelihood with Efron's tie correction for ranks.
• $h = \phi(x)$8: Isotonic-regression term enforcing correct event orderings in survival curve outputs.

2.2 Distribution Shaping Objective in Diffusion Models

P-GRAFT optimizes an intermediate marginal distribution via

$h = \phi(x)$9

with the reshaped reward defined by

$s^{(1)}(x) = w_r^\top h + b_r$0

This generalizes PPO-style reward shaping and connects acceptance-based filtering with KL-regularized policy optimization.

3. Algorithmic Implementation

3.1 P-GRAFT Training for Diffusion Models

The P-GRAFT training algorithm is:

$s^{(2)}(x) = \sigma(W_s h + b_s) \in [0,1]^T$3

At inference:

$s^{(2)}(x) = \sigma(W_s h + b_s) \in [0,1]^T$4 This procedure enables flexible, parameter-efficient, and modular reward shaping at arbitrary intermediate diffusion steps (Anil et al., 3 Oct 2025).

3.2 Training and Early-Stopping in Survival Modeling

Training adopts Adam (learning rate $s^{(1)}(x) = w_r^\top h + b_r$1), batch size 32, dropout (0.2–0.5), L1/L2 penalties, gradient clipping, random search hyperparameter selection, stratified early stopping, and validation on a 20% hold-out set balanced by censoring and event distributions (Luck et al., 2017).

4. Bias–Variance Tradeoff and Theoretical Properties

P-GRAFT's partial rejection sampling entails a fundamental bias–variance tradeoff when performing distribution shaping at different noise levels in diffusion models:

• Variance increases with earlier ($s^{(1)}(x) = w_r^\top h + b_r$2) steps: If rejection or reward signals are applied deep in noise, the conditional distribution over target samples is broad, so $s^{(1)}(x) = w_r^\top h + b_r$3 is nearly independent of $s^{(1)}(x) = w_r^\top h + b_r$4, resulting in high variance of reward estimates given $s^{(1)}(x) = w_r^\top h + b_r$5.
• Bias decays exponentially at later, less noisy steps (small $s^{(1)}(x) = w_r^\top h + b_r$6): Shaping the distribution later leads to a harder score estimation problem, since the model must learn sharp deviations from the reference, but the functional gap (integral squared-error of the score) decreases as $s^{(1)}(x) = w_r^\top h + b_r$7.

A practitioner selects the intermediate noise cutoff $s^{(1)}(x) = w_r^\top h + b_r$8 to optimize this tradeoff for empirical performance and computational efficiency (Anil et al., 3 Oct 2025).

5. Empirical Results

5.1 Diffusion Models

P-GRAFT achieves notable improvements across text-to-image, layout, molecule, and image generation benchmarks:

Task	Baseline	GRAFT	P-GRAFT (cutoff)	Metric	Result
Text-to-Image (SD2)	Base: 66.87 VQA	70.51	71.94 (0.25N)	VQA Score	+7.6%
Layout Gen. (PubLayNet)	Base: .094/.088 Alignment	.064/.064	.053/.064 (0.25N)	Alignment/FID	lower
Molecule Gen. (QM9)	Base: 90.50% stable	90.76%	92.61% (0.25N)	%Stable	+2.11pp
Uncond. Image (CelebA-HQ)	Pre: FID 11.93 (1000 st.)	—	8.02 (200+200)	FID	lower

These represent consistent, domain-generalizable relative gains of 7–11% over base models, particularly for T2I generation (Anil et al., 3 Oct 2025).

5.2 Survival Analysis

On the SRTR kidney graft data ($s^{(1)}(x) = w_r^\top h + b_r$9), P-GRAFT outperforms standard Cox models:

• Cox+Efron (baseline): C-index = 0.6504
• MLP (Cox loss only): C-index = 0.6535
• MLP (ranking loss only): C-index = 0.6302
• P-GRAFT (combined): C-index = 0.6550

This improvement, though numerically modest ($s^{(2)}(x) = \sigma(W_s h + b_s) \in [0,1]^T$0), is statistically meaningful in large medical cohorts and is supported by improved AUROC per-year survival prediction (Luck et al., 2017).

6. Practical and Clinical Implications

P-GRAFT's direct modeling of individualized survival curves enables clinicians to communicate patient-specific, temporalized risk (e.g., "80% five-year survival"), optimize post-transplant care, and donor-recipient matching strategies. The absence of proportional hazards or parametric constraints permits full nonparametric flexibility in survival prediction (Luck et al., 2017).

In generative modeling, P-GRAFT's parameter-efficient distribution shaping raises the practical ceiling for reward-driven fine-tuning of large diffusion models, exhibiting architectural modularity, easy integration into UNet pipelines, and offering new pathways for computationally tractable RL-based generative modeling. The only essential hyperparameter is the cutoff noise level $s^{(2)}(x) = \sigma(W_s h + b_s) \in [0,1]^T$1 (Anil et al., 3 Oct 2025).

7. Related Methodologies and Theoretical Context

P-GRAFT generalizes earlier Generalized Rejection Sampling (GRAFT), which unifies RAFT-type methods (accept/reject at $s^{(2)}(x) = \sigma(W_s h + b_s) \in [0,1]^T$2) with PPO-style KL-regularized distribution shaping. By introducing partial (intermediate) shaping, P-GRAFT interpolates between the high-variance/low-bias regime of early rejection and the low-variance/high-bias regime of late rejection. This theoretical insight motivates the architecture and hyperparameterization of most recent reward fine-tuning in diffusion models. Additionally, the framework can be further extended to parameter-efficient correction methods ("inverse noise") for flow models without explicit rewards, leveraging similar bias–variance principles (Anil et al., 3 Oct 2025).