---
title: 'P-GRAFT: Partial Generalized Rejection Sampling Fine-Tuning'
url: https://www.emergentmind.com/topics/p-graft
type: topic
---

# P-GRAFT: Partial Generalized Rejection Sampling Fine-Tuning

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 [1705.10245], 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) [2510.02692]. 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 $R_i$ on final samples $x_0^{(i)}$.
- Accepting each trajectory with probability determined by an acceptance function $A(R_i,F_R(R_i),x_0^{(i)},\hat P_{X_0})$.
- Collecting the partially denoised states $x_{t_{N_I}}$ (intermediate noise).
- Fine-tuning model parameters for steps $T \to t_{N_I}$ 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:
$$
L(w) = \alpha \, \ell_1(s^{(1)}(w)) + (1-\alpha)\ell_2(s^{(2)}(w)) + \lambda_1 \Vert w \Vert_1 + \lambda_2 \Vert w \Vert_2^2,
$$
where $\alpha$ trades off Cox-style risk ranking and survival curve accuracy:
- $\ell_1(s^{(1)})$: Cox partial log-likelihood with Efron's tie correction for ranks.
- $\ell_2(s^{(2)})$: 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
$$
p_t^\star = \arg\max_{q}\left\{\mathbb{E}_{x \sim q}\left[\hat r(x)\right] - \alpha\,\mathrm{KL}(q \| \bar p_t)\right\}
$$
with the reshaped reward defined by
$$
\frac{\hat r(x)}{\alpha} = \log \mathbb{E}\bigl[ A(r(x_0),\hat F_R(r(x_0)),x_0,\hat P_{X_0}) \mid x_t = x \bigr].
$$
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:

```python
# Algorithm 1: P-GRAFT (Training)
Input: reference p̄, trainable model p_θ, reward r, acceptance A,
       #rounds N_S, intermediate step N_I, samples per round M
Initialize dataset D ← ∅
for j in 1..N_S:
    for i in 1..M:
        draw x_{T:0}^{(i)} ~ p̄_{T:0}
        compute R_i ← r(x₀^{(i)})
    for i in 1..M:
        with probability A(R_i, F̂(R_i), x₀^{(i)}, P̂_{X₀}) accept x_{t_{N_I}^{(i)}}
    add all accepted x_{t_{N_I}} to D
Train p_θ on dataset D for denoising steps T → N_I only
Return fine-tuned p_θ
```

At inference:

```python
# Algorithm 2: P-GRAFT (Inference)
Input: fine-tuned p̂, reference p̄, intermediate N_I
x_T ~ 𝒩(0,I)
for n = N .. N_I + 1:
    x_{t_{n-1}} ← denoise p̂-step(x_{t_n})
for n = N_I .. 1:
    x_{t_{n-1}} ← denoise p̄-step(x_{t_n})
Return x₀
```
This procedure enables flexible, parameter-efficient, and modular reward shaping at arbitrary intermediate diffusion steps [2510.02692].

### 3.2 Training and Early-Stopping in Survival Modeling

Training adopts Adam (learning rate $10^{-5}$), 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 [1705.10245].

## 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 ($t \approx T$) steps: If rejection or reward signals are applied deep in noise, the conditional distribution over target samples is broad, so $r(x_0)$ is nearly independent of $x_t$, resulting in high variance of reward estimates given $x_t$.
- **Bias** decays exponentially at later, less noisy steps (small $t$): 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 $e^{-2t}$.

A practitioner selects the intermediate noise cutoff $N_I$ to optimize this tradeoff for empirical performance and computational efficiency [2510.02692].

## 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 [2510.02692].

### 5.2 Survival Analysis

On the SRTR kidney graft data ($n=131,\!709$), 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 ($\Delta\sim 0.005$), is statistically meaningful in large medical cohorts and is supported by improved AUROC per-year survival prediction [1705.10245].

## 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 [1705.10245].

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 $N_I$ [2510.02692].

## 7. Related Methodologies and Theoretical Context

P-GRAFT generalizes earlier Generalized Rejection Sampling (GRAFT), which unifies RAFT-type methods (accept/reject at $t=0$) 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 [2510.02692].

Source: https://www.emergentmind.com/topics/p-graft