Source-linked AI summary
Risk-Sensitive Reward Composition for Conditional GFlowNets
Carine Ribeiro dos Santos, Ina Pöhner
TL;DR
The paper asks how to combine multiple conformation-specific scores into one reward when populations are uncertain and some conformations are deal-breakers. It uses distributionally robust CVaR targets amortised by one conditional GFlowNet, finding stronger coverage of robust candidates and showing that exploration matters under sparsity.
Problem
GFlowNets require a single reward, but no standard aggregation rule handles uncertain conformation weights while disqualifying candidates that fail required conformations.
Method
The method composes scores with distributionally robust CVaR, tuning tail level β and ambiguity radius ρ, and trains one conditional GFlowNet across the resulting target family.
Results
Across four synthetic cases, one conditional network reaches 0.37–2.7× the sampling floor; under sparsity, state injection achieves 0.987–1.000 mode coverage versus 0.35–0.76 for reweighting.
Takeaways & Limitations
Pricing the lower tail rather than averaging moves 2–10 times more mass toward candidates that pass every conformation, while exploration is decisive when good candidates are rare.
Takeaways & Limitations
The peaked weight family remains at 3.2× the floor-case error despite tripling the budget, and conditioning on score-noise margin σ breaks the endpoint in both hard cases.
Abstract
from arXiv · showhide
Generative Flow Networks (GFlowNets) for structure-based drug design condition on one rigid protein structure. A flexible target holds several distinct structural shapes, its conformations, each occupied for a fraction of the simulation time. Scoring a candidate against all of them raises an open question: how do K scores become one reward? The designer cannot choose arbitrarily. Populations carry simulation error, and biology dictates which conformations are deal-breakers, so a candidate that fails one is disqualified, not merely ranked lower. No standard rule captures this. We compose the reward from a conditional value-at-risk (CVaR), a worst-case score rule, and an ambiguity radius expressing distrust in the stated weights. Together, these define a family of targets, amortised by a single conditional GFlowNet. We answer whether such a sampler can be trained on fully enumerable synthetic worlds, where every error is exact rather than estimated. Pricing the tail rather than averaging moves 2-10 times more mass to candidates that pass every conformation. One network covers the family to within 0.37-2.7x the error of a perfect sampler. An exact-KL oracle, a copy trained on the true target, shows if a shortfall is the optimiser's or the architecture's. When good candidates are rare, exploration decides: injecting unseen states finds 0.987-1.000 of good regions, while reweighting visited finds 0.35-0.76.
1 Introduction
The paper introduces risk-sensitive reward composition for GFlowNets targeting flexible proteins, where multiple conformational scores must encode both uncertain populations and disqualifying failures. It evaluates a conditional sampler and exploration strategies in exactly enumerable synthetic worlds.
- Motivation: Flexible protein targets require aggregating scores across conformations because some must bind, others must be avoided, and failing a required conformation disqualifies a candidate.Fixed weights can misrepresent conformational populations, while weighted sums, Pareto fronts, and action masks do not capture the stated requirements.
- Method: The reward combines CVaR, worst-case aggregation, and an ambiguity radius around stated weights to define a family of conditional targets.The CVaR tail level β and ambiguity radius ρ are included among the conditioning controls, alongside temperature and mixing weight.
- Results: 2–10× more probability mass reaches candidates satisfying every requirement when the tail is priced under distrust of the weights.This compares tail-sensitive aggregation with averaging in the reported synthetic experiments.
- Results: 0.37–2.7× the sampling-floor error is achieved by one conditional network covering the entire target family across four cases.An exact-KL oracle separates representational limits from optimisation shortfalls.
- Results: 0.987–1.000 mode coverage is maintained by state injection under sparsity, compared with 0.35–0.76 for reweighting visited states.The results identify exploration as decisive when good candidates are rare.
- Evaluation: The exact diagnostic framework uses enumerable worlds to compute targets, policies, and L1 divergence exactly across a continuous family of conditions.The synthetic worlds are intended as an evaluation instrument rather than a performance comparison against docking.
2 Related work
The related work covers score aggregation, risk-sensitive GFlowNets, molecular generation, conditioning, exact evaluation, and exploration. The paper distinguishes its contribution by placing deterministic risk and weight ambiguity inside reward construction and by isolating injection versus reweighting under sparsity.
- Aggregation in multi-objective GFlowNets: Prior aggregation methods include desirability functions, linear or rank-based rules, and preference-conditioned GFlowNets, but none prices score-distribution tails under weight uncertainty.Desirability functions encode disqualification multiplicatively but fix the aggregation shape in advance.
- Risk-sensitivity in GFlowNets: Existing risk-sensitive GFlowNets apply distortion measures to stochastic terminal-reward noise, whereas this paper uses deterministic score aggregation and conditions on risk parameters.The paper defines a fixed target before training rather than a target emerging only as an optimisation solution.
- Generation against molecular targets: Prior molecular-generation methods fix a single protein pocket or condition on which pocket to bind, while this work conditions on how multiple conformations should be aggregated.The distinction concerns the indexed target family rather than the conditioning mechanism itself.
- Conditioning and exact evaluation: The evaluation adds exact distances on held-out conditions and an exact-KL oracle trained directly on the enumerated target to attribute residual error.The oracle is used solely to distinguish representational insufficiency from optimisation residual.
- Exploration under sparsity: Prior exploration work is usually compared by sample efficiency or mode count, without distinguishing injected states from states reweighted within the policy’s existing support.The paper identifies this source-of-training-states distinction as the axis separating exploration outcomes.
3 Methods
The method builds exact, risk-sensitive targets over enumerable synthetic worlds and trains one conditional GFlowNet across a continuous family of aggregation and constraint geometries. Exact target computation enables direct L1 evaluation while preserving floors, vetoes, nesting, and uncertainty over conformation weights.
- Synthetic worlds: The experiments use enumerable synthetic spaces, including a 1,024-state calibration family and a 65,536-state sequence family with synthetic PWM and epistatic scores.Scaling variants enlarge the alphabet or sequence length while retaining exact enumeration.
- Exact evaluation: Enumerability makes the target p∗ and terminating policy computable in closed form, so reported errors are exact distances rather than Monte Carlo estimates.Deployment reward evaluation does not require enumerating X, but exact Zc does, so the synthetic setting supports exact diagnostics.
- Risk-sensitive aggregation: DRO-CVaR aggregates K conformation scores by pricing the lower β-tail and taking the infimum over reweightings within an ambiguity ball of radius ρ.β = 1 recovers the weighted average, while ρ = 0 trusts the nominal weights exactly.
- Constraint geometries: Four geometries encode distinct requirements: Case A binds every conformation, Case B floors an aggregate, Case C vetoes disallowed scores, and Case D nests grouped conformations.Case B removes 82% of the reference space and 92% at s = 4; Case C creates 2% isolated dead points.
- Conditional sampling: The condition c indexes target temperature, auxiliary-objective weight, and risk parameters, while FiLM modulates the conditional GFlowNet over this continuous family.A single sampler is trained to serve varying radii because calibrated uncertainty can differ across targets.
- Target construction: The exact construction is evaluated against the target over all states, with threshold crossings producing contour discontinuities and inadmissible tail-level cells reported rather than dropped.This exact target is the reference object against which each learned sampler is measured.
4 Results
The composed aggregation improves joint satisfaction over simpler weighting, while exact diagnostics show when conditional sampling shortfalls arise from optimisation or representation. Under sparsity, exploration by state injection preserves mode coverage, whereas visited-state reweighting can fail despite better density fidelity.
- 4.1 The family is worth indexing: 0.603 joint-satisfaction mass on floor case B versus 0.058 for simpler aggregation, a tenfold gain.The probe cell β = 0.25, ρ = 0.5 reproduces the headline at 0.601.
- 4.1 The family is worth indexing: Weight distrust alone beats the Boltzmann pole in all 32 case–world pairs, but neither single dial reaches the composed interior.The composed regime therefore contributes beyond either dial alone.
- 4.1 The family is worth indexing: 2–2.5× satisfaction gains on Cases A and C cost essentially no diversity, while the nested case instead benefits from aimed distrust.For the nested case, assigning the larger per-origin radius to the noisier origin improves tail stress by Δ = +0.012.
- 4.2 One network serves the family: All four cases meet the 3× sampling-floor endpoint; the veto case is optimisation-limited, while the floor case has a representational deficit on grid geometry.The veto oracle reaches 0.44× the floor versus 1.7× for the student; the floor oracle reaches 1.18×, though 0.95× on the deployment-scale family.
- 4.3 Under sparsity, exploration decides the outcome: The ϵ-plateau creates a degenerate on-policy trap, and contrastive replay can spend up to 93% of samples on its Hamming-1 boundary without crossing it.State injection avoids entering the trap because live states force log Z upward.
- 4.4 What it costs and where it stops: 1.4× to 2.1× floor error with increasing resolution is recoverable, and tripling training restores 2.9–3.0× after depth-related breaches.Resolution costs the veto case, while depth costs the floor case as a budget artifact.
- 4.4 What it costs and where it stops: 3.2× floor error under peaked weights remains unmoved by tripling the budget, and score-noise conditioning makes the hard families partly unrepresentable.Pre-shifting scores by σ roughly doubles error to the true target, from TV 0.20 to 0.38 at 10% corruption.
5 Conclusion
The construction transfers risk-sensitive reward composition to structure-based drug design, where one conditional GFlowNet serves targets indexed by robustness preferences. The results support viable discovery of robust candidates, while distinguishing transferable construction and diagnostics from docking performance.
- One GFlowNet learns the whole family within 0.37–2.7× the error a perfect sampler would still make at the evaluation sample size.
- When good candidates are rare, whether the sampler visits the right states matters more than how well it fits a target.
- Docking scores against conformations become the K scores, populations become nominal weights, and β and ρ control tail strictness and weight distrust.
- The scores are analytic rather than docking output, so the transferable contribution is the construction and diagnostics, not a performance claim.
- The framework is positioned as a viable strategy for discovering robust candidates across heterogeneous structural ensembles in challenging drug discovery problems.
A The risk block is not order-preserving
Temperature conditioning can reweight a target family but preserves state rankings, whereas the risk block can invert them. Therefore, conditioning on risk parameters is necessary for one sampler to represent the family.
- Temperature conditioning preserves the ranking over terminal states because its logit-scale terms differ only by a state-independent constant.
- The risk block changes rankings: candidate A wins at the Boltzmann pole, while uniformly mediocre candidate B wins with a deepened tail or widened radius.
- A single network serving both ranking cells requires risk-block conditioning, since logit scaling alone cannot invert their ordering.
B Notation and the exact target
The exact target is defined before sampling by applying the selected aggregation to every state, assigning floors or vetoes, and producing p∗c. The implementation supports exact, geometry-specific evaluation of ambiguity radii.
- The target p∗c exists before training and is the reference object for every L1 distance in the evaluation.
- The reward family exposes every dial and its effect, while Algorithm B.1 specifies the exact target computation.
- For each state, the algorithm computes an aggregation Ψ(x), combines it with the auxiliary objective, and assigns ε to floored or vetoed states.
C Configuration
The configuration section makes reported runs reconstructable by recording defaults, overrides, and inert settings from the executed configuration.
- The configuration records settings needed to read the results, with unlisted values taking defaults or remaining inert across every reported battery.
- Table C.1 derives configuration values from the dataclasses that ran and records per-battery overrides and inert flags in the released results tree.
D Exploration under sparsity
Under sparsity, exploration determines whether GFlowNet arms discover the target’s modes, especially in the floor case. The online collapse diagnostic exposes this failure during training, while L1 fidelity and mode coverage can rank arms differently.
- Exploration under sparsity: The online collapse diagnostic combines a rising share of samples on the dead set’s boundary with a coverage stall.Unlike exact L1, this signal remains available when the state space cannot be enumerated.
- Exploration under sparsity: At s = 4, the floor case produces the clearest trap: on-policy training starts at L1 = 2 and has the widest spread across seeds.The spread reflects some seeds escaping the trap while others do not.
- Exploration under sparsity: At s = 1, all arms recover the target with tightly clustered L1, so the reference regime does not discriminate between arms.The smooth case spans L1 0.112–0.149, while the floor case already removes 82% of the space without collapse.
- Exploration under sparsity: At s = 4, the floor case is uniquely collapse-prone; on the other geometries, on-policy L1 is comparable to or better than injecting arms.On the veto case, on-policy attains L1 0.148 despite covering the fewest modes, showing that fidelity and coverage can disagree.
- Exploration under sparsity: In the nested case at s = 4, replay attains the best L1 (0.095) but reaches only 0.530 in coverage.This separates density fidelity from mode discovery.
E The price list
The difficulty analysis separates four axes and identifies the case binding each one. Its price list reports held-out L1 relative to the sampling floor, while noting that size and depth costs have different interpretations.
- The price list: Four difficulty axes are separated by character, with each axis associated with a binding case.Table E.1 collects the corresponding difficulty entries.
- The price list: The size price does not flatten with resolution because the veto boundary grows with the space and must be visited.The depth price is measured while the student is still descending, making it a budget artifact rather than a limit.
- The price list: Held-out L1 is reported as a multiple of the sampling floor for each difficulty axis and its binding case.This is the organizing quantity in Table E.1.