Statistical attribute alignment for black-box generative AI via output post-processing

Kevin Jiang; Morgane Austern; Edgar Dobriban; Jason M. Klusowski

stat.MESep 25, 2026

Statistical attribute alignment for black-box generative AI via output post-processing

Authors: Kevin Jiang, Morgane Austern, Edgar Dobriban, Jason M. Klusowski

Abstract

Generative AI systems are increasingly used, but aligning their outputs with user requirements poses a continuing challenge. Here, we aim to ensure that the distribution of an attribute of an AI-generated output aligns with a user-specified target. This is motivated by examples such as fairness, where we want to ensure that a protected attribute (e.g., gender, race, or age categories) follows a desired distribution, and synthetic data generation, where we want the generated data to be representative of a target distribution. We study the practically important black-box access setting, where a user can repeatedly query a generative AI model. The goal is to return $m\ge 1$ outputs whose joint attribute distribution is as close as possible to this target. For both exact and approximate alignment, we develop algorithms that minimize the expected number of queries to the generator, and we further demonstrate their optimality as the number of requested outputs $m \rightarrow \infty$ . Experiments on text-to-image generation and geocoded persona generation tasks show that our post-processing algorithms improve statistical attribute alignment, complementing prompting-based interventions.

Figures & tables

Figure 1: High-level illustration of statistical attribute alignment via black-box output post-processing, using demographic attributes of generated headshots as an example. A user queries a black-box generator to form an annotated candidate batch. A black-box post-processing method then filters the candidate batch by selecting outputs that match a user-specified target distribution (e.g., uniform) as closely as possible, subject to what was observed in the batch.

T_{\nu}^{\star}\coloneq\inf\{t\geq m:\vec{A}^{\star}\in\Omega_{t}\}.

Algorithm 1 Random Demand Coupon Collector (RDC)

\sigma\coloneqq\begin{cases}+\infty,&\text{A-RDC},\\ \tau_{f,\varepsilon},&\text{TA-RDC},\\ n,&\text{CA-RDC},\end{cases}

Algorithm 2 Anytime RDC algorithms

Figure 2: Illustration of the A-RDC algorithm of Definition 16 . There are

k=2

attributes,

m=4

outputs, and the target joint law is

\nu=q^{\otimes m}

for uniform

q=(1/2,1/2)

\vec{C}(t)

, the source-count process of Equation 3 , records the histogram of the M-sampled attributes. At time

t\geq m

, if the realized clock time

g_{t}<g_{\min}\coloneq\min_{1\leq s<t}g_{s}

, we uniformly sample

\widehat{A}_{t}

from the newly feasible attribute sequences (in green). Otherwise, we keep the returned attributes from time

t-1

, with the convention that

g_{\min}=+\infty

for

t=1

. Once the terminal draw

A^{\star}\sim\nu

enters the feasible support, the A-RDC algorithm stops and returns

\vec{A}^{\star}

Figure 3: Cost–discrepancy profiles under KL divergence for the

k=8

non-uniform source and target distributions with output sizes

m\in\{100,500,1000\}

. The horizontal axis reports normalized unconditional KL discrepancy,

D_{\mathrm{KL}}(P_{\mathrm{Alg}_{m}}\|q^{\otimes m})/m

, and the vertical axis reports normalized M-cost,

\mathbb{E}_{p}[N]/m

. The nominal TA-RDC certificate and the first-order optimal frontier are included for comparison. RDC and Oracle RS attain zero discrepancy, although Oracle RS requires knowledge of the source rates. DS denotes direct-sampling, which returns the first

m

candidate samples and has normalized M-cost one.

M-sampling method	Compliance rate	Target-attribute coverage	Attributes observed
HPS	–	3/16	F/Y/W, F/Y/B, F/Y/H
TSI	–	5/16	F/Y/B, F/Y/H, F/O/B, F/O/A, F/O/H
Demographic HPS	0.974	16/16	All

Table 1: Pool-level diagnostics for the individual-headshot prompt sources using the HiDream-O1-Image model. Attribute abbreviations are gender/age/race: F/M, Y/O, and W/B/A/H.

Figure 4: Statistical alignment discrepancies under KL divergence normalized by

m=50

, the number of returned images. The images in the pool are generated using demographic HPS and are annotated into

k=16

target attributes. The panels correspond to the three T2I models.

Model	$\widehat{\alpha}_{\rm same}^{(2)}$	$\widehat{\alpha}_{\rm diff}^{(2)}$	Oracle	CA-RDC	PM
Flux	$3.726\%$	$0.0395\%$	$0.270\%$	$0.257\%\;(0.005\%)$	$0.205\%\;(0.001\%)$
Qwen	$83.019\%$	$0.449\%$	$5.610\%$	$5.295\%\;(0.112\%)$	$4.089\%\;(0.000\%)$
HiDream	$94.662\%$	$0.520\%$	$6.404\%$	$6.099\%\;(0.156\%)$	$4.670\%\;(0.000\%)$

Table 2: Synthetic pairwise audit task at cosine threshold

t=0.45

. Same- and different-attribute match rates are estimated from all valid image pairs in the frozen pool of 1000 demographic HPS images. Oracle is an infeasible benchmark with output attribute law equal to the target joint law

\nu=q^{\otimes m}

. Parentheses report standard errors across ten sampled batches.

Figure 5: Examples of HiDream group-photo generations. The top row shows on- and off-support TSI generations; captions report the annotated observed count vector (O). The bottom row shows hard-compliant, on-support non-compliant, and off-support demographic HPS generations; captions report the prompted target count vector (T) and annotated observed count vector (O). Abbreviations are F/M for female/male, Y/O for young/old, and W/B/A/H for White/Black/Asian/Hispanic.

M-sampling method	On-support rate	Compliance rate	Attributes observed
TSI	0.028	–	6/8
Demographic HPS	0.624	0.219	8/8

Table 3: Pool-level diagnostics for the group-photo prompt sources using the HiDream-O1-Image model.

Figure 6: Various metrics for the group-photo setting with demographic HPS and the HiDream model. There are

k=8

target attributes and each post-processing method returns

m=50

images. Because hard prompting generates images whose annotated attribute lies outside

\textrm{supp}(q)

, we report the feasible-run discrepancies

D_{\mathrm{KL}}(P_{\textrm{CA-RDC}}(\widehat{A}_{n}\mid\alpha_{n}>0)\|q^{\otimes m})

and the infeasibility probability

\beta_{n}\coloneq 1-\Pr(\alpha_{n}>0)

(cf. Section 5.2.1 ) in panels (B) and (C), respectively.

Pool-level diagnostics				Returned personas after post-processing
M-sampling method	On-support rate	Target-attribute coverage	Compliance rate	Post-processing method	Setting	$\text{M-cost}/m$	Housing- cost burden	No-vehicle prevalence	Uninsured prevalence	Top-5-city–state share
Hard SVI prompting	$1.00$	$4/4$	$0.35$	RDC	–	$21.43\;(7.86)$	$0.35\;(0.07)$	$0.39\;(0.07)$	$0.11\;(0.04)$	$0.42\;(0.06)$
				TA-RDC	$d=0.07$	$11.97\;(4.57)$	$0.41\;(0.07)$	$0.46\;(0.07)$	$0.12\;(0.05)$	$0.37\;(0.05)$
				CA-RDC	$n=600$	$11.74\;(1.02)$	$0.41\;(0.07)$	$0.46\;(0.07)$	$0.12\;(0.05)$	$0.38\;(0.05)$
				DS	–	$1.00\;(0.00)$	$0.60\;(0.07)$	$0.62\;(0.07)$	$0.20\;(0.06)$	$0.39\;(0.06)$
County-seeded prompting	$0.98$	$4/4$	–	RDC	–	$1.74\;(0.41)$	$0.27\;(0.06)$	$0.31\;(0.06)$	$0.17\;(0.05)$	$0.15\;(0.03)$

Table 4: SVI4 persona-generation results at

m=50

. Pool-level diagnostics are reported once per M-sampling method, while the remaining columns report means and standard deviations over sets of returned persona outputs. The three socioeconomic rates give the proportions of returned personas, among those linked to the 2022 national ZCTA-level SVI data, whose ZCTAs lie above the 75th national percentile for the corresponding indicator. The top-five city–state share is the proportion belonging to the five most frequently represented city–state pairs within each returned set. Anytime RDC methods are selected by choosing the nearest hyperparameter setting (normalized tolerance

d=\varepsilon/m

or cap

n

for TA-RDC and CA-RDC, respectively) whose M-cost is halfway between direct sampling and RDC.

Figure 7: Various metrics for the SVI persona generation task with

k=4

target quartiles and

m=50

returned personas. Panel (A) reports the KL discrepancy with hard SVI prompting for both TA-RDC and CA-RDC, since the annotated attributes of all 1000 generated personas lie inside

\textrm{supp}(q)

. In contrast, county-seeded prompting produces

18

personas with attributes outside of

\textrm{supp}(q)

(cf. Table 4 ), so panel (B) plots the cost-discrepancy profile of only TA-RDC. For CA-RDC, we report the feasible-run discrepancies

D_{\mathrm{KL}}(P_{\textrm{CA-RDC}}(\widehat{A}_{n}\mid\alpha_{n}>0)\|q^{\otimes m})

and the infeasibility probability

\beta_{n}\coloneq 1-\Pr(\alpha_{n}>0)

(cf. Section 5.2.1 ) in panels (C) and (D), respectively.

Figure 8: Source-dependent cost–discrepancy profiles for direct-sampling (DS), RDC, and TA-RDC, with

m=100

and target joint law

\nu=q^{\otimes m}

with uniform target rates

q=(0.5,0.5)

: (A)

p=q

; (B)

p=(0.8,0.2)

. In this small-scale setting, M-cost and unconditional KL discrepancy are evaluated explicitly through enumeration of count state spaces (cf. Section A.6 ). In (A), DS attains both zero discrepancy and minimum cost. For tolerance values

\varepsilon>m\log 2

, TA-RDC simplifies to DS, while for

\varepsilon=0

, it reduces to RDC. In either of these two cases, it achieves zero unconditional discrepancy. Small irregularities for intermediate tolerances

0<\varepsilon<m\log 2

reflect the discrete nature of the count vectors used for the stopping time (cf. Definition 18 ). In (B), TA-RDC reduces discrepancy relative to DS with additional M-cost.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

\mathcal{D}_{i}\coloneqq\left\{w\in\mathbb{Z}:\max\{0,h-R_{i+1}\}\leq w\leq\min\{h,B_{i}\}\right\}.

Appendix

Algorithm 3 MultinomialCDF

(h,q,\vec{r})

: (upper-rectangular) Multinomial c.d.f. probability via Hayter (2014) recursion

Setting	$k$	Source distribution $p$	$\kappa=\max_{i}q_{i}/p_{i}$
mild skew	3	$(0.20,0.30,0.50)$	1.67
one rare attribute	4	$(0.08,0.22,0.30,0.40)$	3.12
graded source	8	$(0.04,0.06,0.08,0.10,0.12,0.16,0.20,0.24)$	3.12
graded source	12	$p_{i}=(i+1)/90,\ i=1,\ldots,12$	3.75
two mild groups	16	$((5/64)^{\times 8},(3/64)^{\times 8})$	1.33
four rare attributes	16	$((1/64)^{\times 4},(5/64)^{\times 12})$	4.00

Appendix

Table 5: Additional synthetic source settings. In every setting, the target joint law is

\nu=q^{\otimes m}

with the target rates uniform,

q_{i}=1/k

, and the output size is

m=300

Panel of Figure 9 — Figure 9: Cost–discrepancy profiles for the additional synthetic source settings with a uniform target distribution; see Table 5 for exact source rate specifications.

Figure 10: Finite output sample size performance of various post-processing methods in the main text setting with

k=8

attributes and non-uniform source and target rates

p=(0.03,0.05,0.08,0.12,0.16,0.18,0.18,0.20)

and

q=(0.14,0.12,0.13,0.15,0.14,0.12,0.11,0.09)

. First-order optimal procedures – TA-RDC and CA-RDC for KL divergence, and RDC/DS mixture for TV distance – closely track their respective frontiers even for moderate output sizes

m=100

and

200

Figure 11: Cost–discrepancy profiles for approximate rejection sampling (ARS) as described in Section A.3 . Here, the setting is exactly that of the synthetic experiments in the main text, with $k=8$ attributes, non-uniform source and target rates, and output sizes $m\in\{100,500,1000\}$ (cf. Section 5.1 ). The top row reports unconditional KL discrepancy normalized by the output size $m$ and the bottom row reports unconditional TV distance. The orange points of ARS vary the threshold $\gamma$ of the relative width stopping time rule, with smaller thresholds corresponding to more pilot samples and thus a larger M-cost.

Figure 12: Race/ethnicity hard-prompt failure diagnostics for individual-headshot demographic HPS. Heatmaps show annotated race conditional on prompted race for each T2I model, and image panels show representative off-diagonal race/ethnicity annotation failures. Race abbreviations are W/B/A/H for White, Black, Asian, and Hispanic.

Figure 13: Output-size sensitivity analysis for individual-headshot demographic HPS. Columns correspond to three T2I models, while the top and bottom rows report normalized KL discrepancy and unconditional TV discrepancy, respectively. Colors distinguish

m\in\{25,50,100\}

; solid and dashed curves report the actual discrepancies of TA-RDC and CA-RDC, respectively, and dotted curves report the nominal TA-RDC certificates.

Figure 14: Examples of on- and off-support group-photo generations for each image model for the TSI method. Captions report the annotated observed count vector (O). The demographics are abbreviated as follows: F/M stands for female/male; W/B/A/H stands for White, Black, Asian, and Hispanic; and Y/O stands for young/old. The number preceding each letter represents the number of foreground figures identified as the associated demographic.

Figure 15: Examples of hard-compliant, on-support non-compliant, and off-support group-photo generations for each image model under demographic HPS. Captions report the prompted target count vector (T) and annotated observed count vector (O). The demographic count vectors are abbreviated using the same protocol as Figure 14 .

Figure 16: Output-size sensitivity analysis for group-photo demographic HPS. Columns correspond to three T2I models, while the top row reports unconditional KL discrepancy for TA-RDC and feasible-run KL discrepancy for CA-RDC. The bottom row reports unconditional TV discrepancy for both methods. Colors distinguish

m\in\{10,25,50\}

; solid and dashed curves report the actual discrepancies of TA-RDC and CA-RDC, respectively, and dotted curves report the nominal TA-RDC certificates.

Observed label	Age band	City, state, ZIP code	Care setting	Request type
SVI Q1	18–29	Utica, KS 67584	Cardiology	Medication question
SVI Q2	45–64	Lidgerwood, ND 58053	Cardiology	Symptom question
SVI Q3	45–64	Alpharetta, GA 30009	Cardiology	Symptom question
SVI Q4	65+	Algodones, NM 87001	Cardiology	Symptom question

Appendix

Table 6: Representative generated persona extracts from the county-seeded source, one from each observed SVI quartile. The observed label is assigned after generation by matching the generated ZIP to the CDC/ATSDR SVI file.

Figure 17: Mean overall and theme-specific SVI percentiles for all matched ZCTAs in the 2022 CDC/ATSDR file, grouped by overall SVI quartile.

Figure 18: Hard SVI prompting confusion matrix over the full frozen hard-prompt candidate pool of 1000 personas. Each cell reports the integer count, with the row percentage in parentheses. Cell shading encodes the row percentage.

Figure 19: Output-size sensitivity analysis for county-seeded SVI persona generation. The left panel reports unconditional KL discrepancy for TA-RDC and feasible-run KL discrepancy for CA-RDC, while the right panel reports unconditional TV discrepancy for both methods. Colors distinguish

m\in\{25,50,100\}

; solid and dashed curves report the actual discrepancies of TA-RDC and CA-RDC, respectively, and dotted curves report the nominal TA-RDC certificates.

Figure 20: Example pairwise decisions on generated headshots, grouped by whether the annotated demographic attributes agree. The thresholds are $t=0.45$ for HiDream and $t=0.35$ for Flux and Qwen , as shown in the panels.

Figure 21: Threshold sensitivity on the cosine threshold for the pairwise audit task of Section 5.2.3 . The gap in the false match rate between the oracle and CA-RDC (solid) and proportion matching (dashed) are plotted across all three T2I models. The vertical line denotes the cosine threshold

t=0.45

reported in the main text. The gap between the Oracle and PM procedure remains large through non-trivial intermediate ranges of the threshold. In contrast, CA-RDC better captures pairwise interactions in the same- and different-attribute class counts by targeting a joint law. This results in better estimates of the false-match rate.

Constrained Generative Modeling

Constrained generative modeling develops methods that produce samples satisfying structural, physical, semantic, or user-defined requirements.

See all 192 papers

Inference-Time Attribute Distribution Alignment for Unconditional Diffusion
May 8, 2026Hao Luan, See-Kiong Ng, Chun Kai LingControllable Image Generation Inference-Time Guidance
School of Computing, National University of Singapore · Institute of Data Science, National University of Singapore
Who Should Be Generated? Justifying Demographic Targets in Open-Ended Generation
Aug 3, 2026Zeshen Zheng, Yujia He, Qianmian Lin +2Language Model Bias Evaluation Demographic Bias in Language Models
Sun Yat-sen University, China
Counting Worlds Branching Time Semantics for post-hoc Bias Mitigation in generative AI
Apr 21, 2026Alessandro G. Buda, Giuseppe Primiero, Leonardo Ceragioli +1Algorithmic Fairness Formal Verification

May 8, 2026J/K move · Enter open · S save

Inference-Time Attribute Distribution Alignment for Unconditional Diffusion

Hao Luan, See-Kiong Ng, Chun Kai Ling

School of Computing, National University of Singapore · Institute of Data Science, National University of Singapore

Controllable Image Generation Inference-Time Guidance Diffusion Models

Open paper AskPDF arXiv

Inference-time controllable generation is essential for real-world applications of unconditional diffusion models. However, most existing techniques focus on individual samples, struggling in applications that require the sample population to follow specific attribute distributions (e.g., demographic balance or semantic proportions). We formalize this setting as the inference-time attribute distributional alignment problem for pretrained unconditional diffusion models. To address this, we cast inference-time attribute distributional alignment as an optimal control problem over the reverse diffusion process, viewing the process as the rollout of a dynamical system and augmenting it with additive, time-dependent perturbations as control. We solve for the perturbations using an optimal-control-based algorithm to optimize a differentiable distribution-matching objective while penalizing control effort to preserve data fidelity. Experiment results in image generation demonstrate that our proposed plug-and-play approach can better align attribute distributions to diverse and flexible test-time targets compared to baselines, without retraining or finetuning the pretrained diffusion model.

Abstract

Constrained Generative Modeling

Explore similar work

Inference-Time Attribute Distribution Alignment for Unconditional Diffusion

Who Should Be Generated? Justifying Demographic Targets in Open-Ended Generation

Counting Worlds Branching Time Semantics for post-hoc Bias Mitigation in generative AI