Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Xi Wang; Wenbo Lu; Shengjie Wang

cs.LGFeb 28, 2026

Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Authors: Xi Wang, Wenbo Lu, Shengjie Wang

Organizations: Courant Institute School of Mathematics, Computing, and Data Science, New York University

Abstract

Generative Flow Networks (GFlowNets) enable fine-tuning large language models to approximate reward-proportional posteriors, but they remain prone to mode collapse, manifesting as prefix collapse and length bias. We attribute this to two factors: (i) weak credit assignment to early prefixes, and (ii) biased replay that induces a shifted, non-representative training flow distribution. We propose Rooted absorbed prefix Trajectory Balance RapTB, an objective that anchors subtrajectory supervision at the root and propagates terminal rewards to intermediate prefixes via absorbed suffix-based backups, providing dense prefix-level learning signals. To mitigate replay-induced distribution shift, we further introduce SubM, a submodular replay refresh strategy that promotes both high reward and diversity. Empirically, on tasks such as molecule generation with LLM using SMILES strings, RapTB combined with SubM consistently improves optimization performance and molecular diversity while preserving high validity. The code is released on https://github.com/ComDec/ChemGFN.

Figures & tables

Figure 1: Training objectives for LLM-GFlowNets. TB uses only terminal reward $\log R(s_{\tau})$ ( $O(1)$ ). SubTB adds $O(N^{2})$ windowed consistency constraints. RapTB replaces prefix stop-rewards with suffix-absorbed targets $u_{j}$ and applies $O(N)$ rooted prefix constraints. ( $N$ : trajectory length). The right column shows examples of generated molecules trained with different losses. QED (Quantitative Estimate of Drug-likeness) is a comprehensive metric for measuring a molecule’s drug-likeness and is used as the task reward during training.

Method	Acc $\uparrow$	Score $\uparrow$	Entropy $\uparrow$	FPDiv $\uparrow$	Len
PPO	1.000	0.604	$\approx$ 0	–	–
GRPO	0.997	0.661	0.98	–	10.0
TB	0.998	0.717	2.503	0.807	3.065
SubTB	0.328	0.755	2.127	0.836	8.354
RapTB	0.996	0.740	2.448	0.860	6.142
RapTB + SubM	0.988	0.844	2.726	0.898	7.435

Table 1: SMILES generation performance. Unless specified, metrics are computed on valid samples. Len represents average token length. Results are averaged over seeds; see Appendix A.1 for 95% CIs and per-length breakdowns.

Panel of Figure 2 — Figure 2: Length-stratified analysis on SMILES ( $L_{\max}=10$ ). (a) Distribution of valid generation lengths. (b) Mean Score and FPDiv conditioned on length.

Figure 3: Prefix-collapse diagnostics on SMILES (

L_{\max}=10

). Metrics vs. prefix length

k

computed on correct samples: prefix survival (fraction of samples reaching length

k

), prefix entropy (diversity of prefix), and top-1 mass (frequency of the most common prefix).

	Performance		Length Dist.			Prefix Diagnostics			Diversity
Method	Acc	Score	0–5 $\downarrow$	6–10	11+ $\uparrow$	Surv $\uparrow$	PefEnt $\uparrow$	Top1 $\downarrow$	MacroFP $\uparrow$	FPDiv $\uparrow$
TB	0.999	0.716	0.858	0.129	0.013	0.207	2.99	0.303	0.653	0.813
SubTB	0.636	0.742	0.286	0.442	0.271	0.561	4.82	0.085	0.716	0.770
RapTB	0.988	0.768	0.113	0.318	0.568	0.681	5.59	0.084	0.793	0.810
RapTB+SubM	0.972	0.849	0.094	0.205	0.701	0.751	5.32	0.071	0.805	0.868

Table 2: Longer-horizon stress test on SMILES (

L_{\max}=15

). We report length fractions of valid samples (Length Dist.) and diversity metrics. Prefix diagnostics are averaged over

k

: Surv (Prefix survival), PefEnt (prefix entropy), and Top1 (Top-1 mass). MacroFP macro-averages FPDiv across length bins 0–5, 6–10, and 11+. RapTB+SubM achieves the best balance of long-horizon coverage and diversity. CIs omitted for space (Appendix A.1 ).

		Diversity		Quality	Distributional Fidelity
Replay	Objective	Unique ✓ $\uparrow$	NormCov $\uparrow$	Acc $\uparrow$	KL( $\pi\!\to\!p^{*}$ ) $\downarrow$	KL( $p^{*}\!\to\!\pi$ ) $\downarrow$	JS ${}_{\text{tok}}$ $\downarrow$
N/A	PPO	1	0.000	0.003	–	–	–
N/A	GRPO	1	0.000	0.002	–	–	–
PRT	TB	103.7	0.016	0.999	1.105	7.803	0.292
	SubTB	292.0	0.046	0.311	0.424	0.672	0.107
	RapTB	129.3	0.020	0.992	0.908	5.538	0.230

Table 3: Expr24 results under different replay schemes. Results are averaged over different random seeds; 95% confidence intervals are reported in Appendix A.2 .

RP Replay
Method	Acc $\uparrow$	NormCov $\uparrow$	$\log p_{\text{term}}(\tau)$	$\log Z$
TB	0.999	0.001	-0.000	0.063
SubTB	0.229	0.051	-79.638	–
RapTB	0.991	0.039	-0.065	0.062
RootSubTBLogZ	0.999	0.023	-0.068	0.062
Oracle Replay

Table 4: Termination/length calibration diagnostic on Expr24. More negative values indicate overly suppressed termination.

\log p_{\text{term}}(\tau)

is the termination log-probability at the sampled stop step, computed from the model’s raw

q_{\theta}(\top\mid s_{0:\tau})

Method	Entropy $\uparrow$	BLEU-4 $\uparrow$	Len	$\Delta\log p_{\text{term}}$
TB	2.966	5.95	13.86	-1.34
SubTB	3.719	24.39	20.00	-28.32
RapTB	3.933	11.75	15.63	-0.94
RapTB+SubM	4.102	33.23	11.83	4.89

Table 5: CommonGen performance.

\Delta\log p_{\text{term}}

measures the deviation of the learned policy’s termination logits from the pre-trained reference anchor.

Variant	Score $\uparrow$	FPDiv $\uparrow$	Entropy $\uparrow$	Len
RapTB	0.740	0.860	2.448	6.142
RapTB w/o reward absorb	0.716	0.805	2.031	5.296
RapTB absorb max-only	0.821	0.775	1.716	7.431
RapTB absorb soft-only	0.819	0.748	1.516	7.710
RapTB + SubM (Full)	0.844	0.898	2.726	7.435
+ length-only SubM	0.773	0.885	1.914	6.569

Table 6: Ablation study on SMILES generation.

Appendix figures & tables28 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Acc	Score	Entropy	FPDiv	Len μ	Len 50	Len 90	Frac[0–2]	Frac[3–5]	Frac[6–8]	Frac[9–10]
TB	0.998 $\pm$ 0.001	0.717 $\pm$ 0.001	2.503 $\pm$ 0.026	0.807 $\pm$ 0.003	3.06 $\pm$ 0.02	2.15 $\pm$ 0.03	6.42 $\pm$ 0.13	0.601 $\pm$ 0.004	0.252 $\pm$ 0.007	0.084 $\pm$ 0.005	0.064 $\pm$ 0.004
TB+SubM	0.996 $\pm$ 0.000	0.842 $\pm$ 0.001	2.775 $\pm$ 0.002	0.889 $\pm$ 0.002	6.56 $\pm$ 0.09	7.23 $\pm$ 0.07	9.84 $\pm$ 0.06	0.143 $\pm$ 0.006	0.186 $\pm$ 0.015	0.348 $\pm$ 0.013	0.323 $\pm$ 0.008
RapTB	0.996 $\pm$ 0.001	0.740 $\pm$ 0.004	2.448 $\pm$ 0.017	0.860 $\pm$ 0.001	6.14 $\pm$ 0.03	6.58 $\pm$ 0.03	9.01 $\pm$ 0.04	0.134 $\pm$ 0.007	0.204 $\pm$ 0.009	0.466 $\pm$ 0.017	0.197 $\pm$ 0.015
RapTB+SubM	0.988 $\pm$ 0.003	0.844 $\pm$ 0.001	2.726 $\pm$ 0.017	0.898 $\pm$ 0.001	7.44 $\pm$ 0.05	7.91 $\pm$ 0.03	9.84 $\pm$ 0.02	0.046 $\pm$ 0.005	0.119 $\pm$ 0.003	0.433 $\pm$ 0.012	0.402 $\pm$ 0.007
SubTB	0.328 $\pm$ 0.016	0.755 $\pm$ 0.004	2.127 $\pm$ 0.037	0.836 $\pm$ 0.003	8.35 $\pm$ 0.06	9.22 $\pm$ 0.05	9.97 $\pm$ 0.03	0.006 $\pm$ 0.001	0.047 $\pm$ 0.005	0.076 $\pm$ 0.005	0.871 $\pm$ 0.001
SubTB+SubM	0.298 $\pm$ 0.006	0.736 $\pm$ 0.005	2.165 $\pm$ 0.006	0.851 $\pm$ 0.002	8.73 $\pm$ 0.06	9.59 $\pm$ 0.08	9.99 $\pm$ 0.01	0.005 $\pm$ 0.001	0.029 $\pm$ 0.003	0.057 $\pm$ 0.001	0.909 $\pm$ 0.003

Appendix

Table 7: All-length averaged SMILES performance and induced length distribution (

L_{\max}=10

; mean

\pm

95% CI over 6 runs). Acc is computed over all samples; Score/Entropy/FPDiv are computed on valid samples only; Frac[

\cdot

] is computed over all samples.

$L$	TB				SubTB				RapTB				TB+SubM				SubTB+SubM				RapTB+SubM
	Acc	Score	Frac	Count	Acc	Score	Frac	Count	Acc	Score	Frac	Count	Acc	Score	Frac	Count	Acc	Score	Frac	Count	Acc	Score	Frac	Count
1	1 $\pm$ 0	0.708 $\pm$ 0.004	0.215 $\pm$ 0.007	687 $\pm$ 22.7	1 $\pm$ 0	0.762 $\pm$ 0.007	0.002 $\pm$ 0.001	1.7 $\pm$ 0.7	1 $\pm$ 0	0.651 $\pm$ 0.004	0.042 $\pm$ 0.002	133 $\pm$ 4.9	1 $\pm$ 0	0.712 $\pm$ 0.007	0.055 $\pm$ 0.007	175 $\pm$ 22.6	1 $\pm$ 0	0.766 $\pm$ 0.001	0.004 $\pm$ 0.003	3.7 $\pm$ 2.6	1 $\pm$ 0	0.702 $\pm$ 0.006	0.02 $\pm$ 0.002	64.3 $\pm$ 7.5
2	0.999 $\pm$ 0.001	0.752 $\pm$ 0.002	0.387 $\pm$ 0.008	1235.7 $\pm$ 28.1	1 $\pm$ 0	0.748 $\pm$ 0.003	0.017 $\pm$ 0.003	17.7 $\pm$ 3.5	0.999 $\pm$ 0.002	0.725 $\pm$ 0.004	0.092 $\pm$ 0.006	294.7 $\pm$ 17.4	1 $\pm$ 0	0.79 $\pm$ 0	0.088 $\pm$ 0.003	282 $\pm$ 7.9	1 $\pm$ 0	0.75 $\pm$ 0.014	0.012 $\pm$ 0.004	11.7 $\pm$ 4	1 $\pm$ 0	0.783 $\pm$ 0.005	0.026 $\pm$ 0.003	83.3 $\pm$ 10.8
3	0.998 $\pm$ 0.002	0.727 $\pm$ 0.002	0.138 $\pm$ 0.003	439.3 $\pm$ 10.5	1 $\pm$ 0	0.774 $\pm$ 0.01	0.009 $\pm$ 0.004	9.3 $\pm$ 4	1 $\pm$ 0	0.754 $\pm$ 0.001	0.056 $\pm$ 0.002	177 $\pm$ 6.8	1 $\pm$ 0	0.787 $\pm$ 0.003	0.058 $\pm$ 0.008	184.3 $\pm$ 26.7	1 $\pm$ 0	0.759 $\pm$ 0.008	0.008 $\pm$ 0.003	8 $\pm$ 3	0.996 $\pm$ 0.008	0.759 $\pm$ 0.006	0.03 $\pm$ 0.006	95 $\pm$ 19.3
4	0.998 $\pm$ 0.003	0.722 $\pm$ 0.001	0.062 $\pm$ 0.002	199.3 $\pm$ 6.8	1 $\pm$ 0	0.787 $\pm$ 0.003	0.035 $\pm$ 0.003	36.3 $\pm$ 2.4	0.995 $\pm$ 0.006	0.748 $\pm$ 0.003	0.058 $\pm$ 0.004	185.7 $\pm$ 13.6	0.997 $\pm$ 0.005	0.793 $\pm$ 0.003	0.046 $\pm$ 0.005	145.3 $\pm$ 16.8	1 $\pm$ 0	0.789 $\pm$ 0.006	0.024 $\pm$ 0.003	23.3 $\pm$ 2.4	1 $\pm$ 0	0.813 $\pm$ 0.012	0.031 $\pm$ 0.002	99.3 $\pm$ 5.3
5	0.996 $\pm$ 0.008	0.739 $\pm$ 0.003	0.052 $\pm$ 0.002	165.3 $\pm$ 5.8	1 $\pm$ 0	0.813 $\pm$ 0.001	0.099 $\pm$ 0.011	103.7 $\pm$ 9.9	0.998 $\pm$ 0.004	0.78 $\pm$ 0.009	0.09 $\pm$ 0.005	288.3 $\pm$ 14.6	0.997 $\pm$ 0.002	0.851 $\pm$ 0.004	0.083 $\pm$ 0.004	263.7 $\pm$ 13.6	1 $\pm$ 0	0.815 $\pm$ 0.001	0.066 $\pm$ 0.011	63 $\pm$ 10.4	0.995 $\pm$ 0.006	0.836 $\pm$ 0.003	0.059 $\pm$ 0.004	185.3 $\pm$ 12.1

Appendix

Table 8: Per-length valid-only core metrics of SMILES generation (mean

\pm

95% CI,

L_{\max}=10

$L$	TB		SubTB		RapTB		TB+SubM		SubTB+SubM		RapTB+SubM
	Entropy	FPDiv	Entropy	FPDiv	Entropy	FPDiv	Entropy	FPDiv	Entropy	FPDiv	Entropy	FPDiv
1	2.39 $\pm$ 0.07	0.789 $\pm$ 0.012	0.23 $\pm$ 0.45	0.25 $\pm$ 0.283	2.28 $\pm$ 0.11	0.824 $\pm$ 0.005	1.35 $\pm$ 0.09	0.628 $\pm$ 0.014	0.17 $\pm$ 0.33	0.2 $\pm$ 0.226	1.32 $\pm$ 0.08	0.651 $\pm$ 0.027
2	2.54 $\pm$ 0.01	0.812 $\pm$ 0.005	1.05 $\pm$ 0.33	0.445 $\pm$ 0.069	2.33 $\pm$ 0.03	0.827 $\pm$ 0.003	2.12 $\pm$ 0.07	0.84 $\pm$ 0.004	0.99 $\pm$ 0.09	0.435 $\pm$ 0.056	1.82 $\pm$ 0.13	0.759 $\pm$ 0.022
3	2.44 $\pm$ 0.02	0.825 $\pm$ 0.001	1.1 $\pm$ 0.17	0.722 $\pm$ 0.022	1.99 $\pm$ 0.03	0.818 $\pm$ 0.003	2.21 $\pm$ 0.07	0.877 $\pm$ 0.002	0.87 $\pm$ 0.43	0.604 $\pm$ 0.091	1.76 $\pm$ 0.04	0.838 $\pm$ 0.002
4	2.46 $\pm$ 0.05	0.854 $\pm$ 0.002	1.35 $\pm$ 0.14	0.703 $\pm$ 0.03	2.02 $\pm$ 0.06	0.867 $\pm$ 0.005	2.39 $\pm$ 0.01	0.909 $\pm$ 0.005	1.31 $\pm$ 0.36	0.739 $\pm$ 0.05	2.19 $\pm$ 0.07	0.888 $\pm$ 0.007
5	2.41 $\pm$ 0.02	0.88 $\pm$ 0.005	1.19 $\pm$ 0.11	0.714 $\pm$ 0.01	2.21 $\pm$ 0.01	0.859 $\pm$ 0.003	2.48 $\pm$ 0.04	0.906 $\pm$ 0.004	1.15 $\pm$ 0.05	0.721 $\pm$ 0.004	2.21 $\pm$ 0.12	0.877 $\pm$ 0.015

Appendix

Table 9: Per-length valid-only diversity metrics of SMILES generation(mean

\pm

95% CI,

L_{\max}=10

$L$	TB				SubTB				RapTB				TB+SubM				SubTB+SubM				RapTB+SubM
	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol
1	29.3 $\pm$ 1.1	25.3 $\pm$ 1.5	0 $\pm$ 0	0 $\pm$ 0	1.3 $\pm$ 0.4	1.3 $\pm$ 0.4	0.8 $\pm$ 0.2	0.8 $\pm$ 0.2	17.3 $\pm$ 0.8	15.7 $\pm$ 0.4	0.1 $\pm$ 0	0.1 $\pm$ 0	13.7 $\pm$ 0.8	13.7 $\pm$ 0.8	0.1 $\pm$ 0	0.1 $\pm$ 0	1.3 $\pm$ 0.4	1.3 $\pm$ 0.4	0.5 $\pm$ 0.3	0.5 $\pm$ 0.3	8.3 $\pm$ 0.8	8.3 $\pm$ 0.8	0.1 $\pm$ 0	0.1 $\pm$ 0
2	184.7 $\pm$ 2.1	115.7 $\pm$ 3	0.1 $\pm$ 0	0.2 $\pm$ 0	6.3 $\pm$ 0.8	6 $\pm$ 0.7	0.4 $\pm$ 0	0.3 $\pm$ 0	77 $\pm$ 1.9	71 $\pm$ 2.5	0.3 $\pm$ 0	0.2 $\pm$ 0	60.3 $\pm$ 3	59.7 $\pm$ 2.9	0.2 $\pm$ 0	0.2 $\pm$ 0	5 $\pm$ 0.7	5 $\pm$ 0.7	0.4 $\pm$ 0	0.4 $\pm$ 0	27.3 $\pm$ 1.5	27.3 $\pm$ 1.5	0.3 $\pm$ 0	0.3 $\pm$ 0
3	224 $\pm$ 7.5	215.7 $\pm$ 7.3	0.5 $\pm$ 0	0.5 $\pm$ 0	6.3 $\pm$ 1.7	6.3 $\pm$ 1.7	0.7 $\pm$ 0.1	0.7 $\pm$ 0.1	78.3 $\pm$ 2.2	76.3 $\pm$ 2.2	0.4 $\pm$ 0	0.4 $\pm$ 0	52.3 $\pm$ 5.8	52.3 $\pm$ 5.8	0.3 $\pm$ 0	0.3 $\pm$ 0	4.7 $\pm$ 1.8	4.7 $\pm$ 1.8	0.6 $\pm$ 0.1	0.6 $\pm$ 0.1	27.7 $\pm$ 3	27.7 $\pm$ 3	0.3 $\pm$ 0	0.3 $\pm$ 0
4	166 $\pm$ 4	164 $\pm$ 3.6	0.8 $\pm$ 0	0.8 $\pm$ 0	17 $\pm$ 1.2	16.7 $\pm$ 0.8	0.5 $\pm$ 0.1	0.5 $\pm$ 0	111 $\pm$ 3.1	107.3 $\pm$ 3	0.6 $\pm$ 0	0.6 $\pm$ 0	58 $\pm$ 0.7	57.7 $\pm$ 1.1	0.4 $\pm$ 0	0.4 $\pm$ 0	14.7 $\pm$ 2.2	13.7 $\pm$ 1.8	0.6 $\pm$ 0.1	0.6 $\pm$ 0.1	41 $\pm$ 6.1	41 $\pm$ 6.1	0.4 $\pm$ 0.1	0.4 $\pm$ 0.1
5	156 $\pm$ 5.9	153.3 $\pm$ 5.8	0.9 $\pm$ 0	0.9 $\pm$ 0	41.3 $\pm$ 3.7	38.3 $\pm$ 3.9	0.4 $\pm$ 0	0.4 $\pm$ 0	187 $\pm$ 7.6	181.7 $\pm$ 7.9	0.6 $\pm$ 0	0.6 $\pm$ 0	101.3 $\pm$ 3.5	100.7 $\pm$ 3.3	0.4 $\pm$ 0	0.4 $\pm$ 0	31.3 $\pm$ 4.2	28.3 $\pm$ 3.5	0.5 $\pm$ 0	0.4 $\pm$ 0	74 $\pm$ 8.1	73.7 $\pm$ 7.9	0.4 $\pm$ 0	0.4 $\pm$ 0

Appendix

Table 10: Per-length valid-only uniqueness metrics of SMILES generation (mean

\pm

95% CI,

L_{\max}=10

$k$	TB					SubTB					RapTB
	Survival	PefEnt	Eff	Top1	UniqueRate	Survival	PefEnt	Eff	Top1	UniqueRate	Survival	PefEnt	Eff	Top1	UniqueRate
1	1 $\pm$ 0	2.531 $\pm$ 0.022	12.57 $\pm$ 0.28	0.274 $\pm$ 0.01	0.013 $\pm$ 0.001	1 $\pm$ 0	1.539 $\pm$ 0.036	4.66 $\pm$ 0.17	0.419 $\pm$ 0.01	0.017 $\pm$ 0.001	1 $\pm$ 0	1.841 $\pm$ 0.025	6.3 $\pm$ 0.15	0.4 $\pm$ 0.004	0.01 $\pm$ 0
2	0.785 $\pm$ 0.005	4.748 $\pm$ 0.019	115.34 $\pm$ 2.25	0.077 $\pm$ 0.001	0.135 $\pm$ 0.002	0.998 $\pm$ 0	3.544 $\pm$ 0.063	34.69 $\pm$ 2.15	0.096 $\pm$ 0.004	0.096 $\pm$ 0.002	0.958 $\pm$ 0.001	3.838 $\pm$ 0.023	46.46 $\pm$ 1.05	0.174 $\pm$ 0.005	0.076 $\pm$ 0.002
3	0.398 $\pm$ 0.002	5.418 $\pm$ 0.027	225.52 $\pm$ 6.1	0.121 $\pm$ 0.005	0.45 $\pm$ 0.005	0.982 $\pm$ 0.002	4.563 $\pm$ 0.044	95.94 $\pm$ 4.19	0.088 $\pm$ 0.006	0.24 $\pm$ 0.003	0.866 $\pm$ 0.005	4.717 $\pm$ 0.031	111.95 $\pm$ 3.52	0.184 $\pm$ 0.005	0.215 $\pm$ 0.005
4	0.261 $\pm$ 0.002	5.17 $\pm$ 0.038	176.02 $\pm$ 6.81	0.167 $\pm$ 0.009	0.566 $\pm$ 0.012	0.973 $\pm$ 0.004	5.054 $\pm$ 0.035	156.78 $\pm$ 5.48	0.088 $\pm$ 0.006	0.355 $\pm$ 0.003	0.81 $\pm$ 0.006	5.45 $\pm$ 0.047	233.1 $\pm$ 11.25	0.135 $\pm$ 0.007	0.357 $\pm$ 0.01
5	0.198 $\pm$ 0.002	4.796 $\pm$ 0.041	121.2 $\pm$ 4.96	0.219 $\pm$ 0.014	0.599 $\pm$ 0.009	0.938 $\pm$ 0.005	5.575 $\pm$ 0.048	264.06 $\pm$ 12.39	0.03 $\pm$ 0.003	0.471 $\pm$ 0.004	0.752 $\pm$ 0.004	5.747 $\pm$ 0.051	313.88 $\pm$ 16.34	0.138 $\pm$ 0.007	0.474 $\pm$ 0.01

Appendix

Table 11: Prefix statistics by depth. Mean

\pm

95% CI.

$k$	TB+SubM					SubTB+SubM					RapTB+SubM
	Survival	PefEnt	Eff	Top1	UniqueRate	Survival	PefEnt	Eff	Top1	UniqueRate	Survival	PefEnt	Eff	Top1	UniqueRate
1	1 $\pm$ 0	2.715 $\pm$ 0.017	15.11 $\pm$ 0.25	0.165 $\pm$ 0.006	0.011 $\pm$ 0	1 $\pm$ 0	1.559 $\pm$ 0.037	4.76 $\pm$ 0.18	0.394 $\pm$ 0.01	0.017 $\pm$ 0.001	1 $\pm$ 0	2.563 $\pm$ 0.004	12.98 $\pm$ 0.05	0.193 $\pm$ 0.005	0.009 $\pm$ 0
2	0.945 $\pm$ 0.004	4.778 $\pm$ 0.004	118.92 $\pm$ 0.46	0.052 $\pm$ 0.003	0.094 $\pm$ 0.001	0.996 $\pm$ 0.002	3.669 $\pm$ 0.041	39.26 $\pm$ 1.58	0.094 $\pm$ 0.005	0.108 $\pm$ 0.004	0.98 $\pm$ 0.001	4.43 $\pm$ 0.015	83.96 $\pm$ 1.27	0.065 $\pm$ 0.001	0.076 $\pm$ 0.001
3	0.857 $\pm$ 0.004	5.224 $\pm$ 0.01	185.6 $\pm$ 1.81	0.033 $\pm$ 0.002	0.164 $\pm$ 0.005	0.984 $\pm$ 0.002	4.723 $\pm$ 0.024	112.55 $\pm$ 2.69	0.089 $\pm$ 0.004	0.278 $\pm$ 0.008	0.953 $\pm$ 0.003	5.019 $\pm$ 0.022	151.34 $\pm$ 3.33	0.042 $\pm$ 0.006	0.141 $\pm$ 0.001
4	0.799 $\pm$ 0.005	5.434 $\pm$ 0.014	229.16 $\pm$ 3.26	0.034 $\pm$ 0.002	0.211 $\pm$ 0.006	0.975 $\pm$ 0.004	5.275 $\pm$ 0.021	195.48 $\pm$ 4.12	0.09 $\pm$ 0.004	0.437 $\pm$ 0	0.923 $\pm$ 0.007	5.342 $\pm$ 0.036	209.02 $\pm$ 7.37	0.037 $\pm$ 0.002	0.196 $\pm$ 0.004
5	0.753 $\pm$ 0.008	5.667 $\pm$ 0.021	289.19 $\pm$ 6	0.024 $\pm$ 0.002	0.254 $\pm$ 0.007	0.951 $\pm$ 0.004	5.765 $\pm$ 0.014	318.94 $\pm$ 4.35	0.029 $\pm$ 0.001	0.563 $\pm$ 0	0.892 $\pm$ 0.006	5.641 $\pm$ 0.047	282.07 $\pm$ 13.11	0.036 $\pm$ 0.004	0.254 $\pm$ 0.006

Appendix

Table 12: Prefix statistics by depth (Continue). Mean

\pm

95% CI.

$L$	TB				SubTB				RapTB				RapTB+SubM
	Acc	Score	Frac	Count	Acc	Score	Frac	Count	Acc	Score	Frac	Count	Acc	Score	Frac	Count
1	1 $\pm$ 0	0.707 $\pm$ 0.003	0.242 $\pm$ 0.012	774.7 $\pm$ 37.9	1 $\pm$ 0	0.765 $\pm$ 0	0.002 $\pm$ 0.001	1.7 $\pm$ 0.7	1 $\pm$ 0	0.7 $\pm$ 0.005	0.033 $\pm$ 0.004	105.7 $\pm$ 11.6	1 $\pm$ 0	0.694 $\pm$ 0.012	0.018 $\pm$ 0.001	56.5 $\pm$ 2.9
2	1 $\pm$ 0	0.748 $\pm$ 0	0.34 $\pm$ 0.01	1086.7 $\pm$ 33.2	1 $\pm$ 0	0.738 $\pm$ 0.009	0.05 $\pm$ 0.002	51.3 $\pm$ 2.6	1 $\pm$ 0	0.723 $\pm$ 0.003	0.033 $\pm$ 0.005	103 $\pm$ 16.7	1 $\pm$ 0	0.798 $\pm$ 0.004	0.018 $\pm$ 0.008	57.5 $\pm$ 26.5
3	0.999 $\pm$ 0.002	0.724 $\pm$ 0.007	0.128 $\pm$ 0.005	409 $\pm$ 16.3	1 $\pm$ 0	0.723 $\pm$ 0.016	0.03 $\pm$ 0.006	31 $\pm$ 6	1 $\pm$ 0	0.747 $\pm$ 0.012	0.011 $\pm$ 0.001	34.3 $\pm$ 2.4	1 $\pm$ 0	0.824 $\pm$ 0.003	0.016 $\pm$ 0.004	50.5 $\pm$ 12.7
4	0.999 $\pm$ 0.002	0.739 $\pm$ 0	0.088 $\pm$ 0.006	281.3 $\pm$ 20.3	1 $\pm$ 0	0.796 $\pm$ 0.008	0.069 $\pm$ 0.011	70.7 $\pm$ 11.3	0.995 $\pm$ 0.01	0.758 $\pm$ 0.005	0.017 $\pm$ 0.002	55.3 $\pm$ 7.5	1 $\pm$ 0	0.853 $\pm$ 0.017	0.016 $\pm$ 0.006	51 $\pm$ 17.6
5	0.997 $\pm$ 0.003	0.749 $\pm$ 0.001	0.06 $\pm$ 0.003	192.3 $\pm$ 10.3	1 $\pm$ 0	0.795 $\pm$ 0.003	0.135 $\pm$ 0.004	137 $\pm$ 5.2	1 $\pm$ 0	0.779 $\pm$ 0.01	0.019 $\pm$ 0.008	60 $\pm$ 25.5	1 $\pm$ 0	0.888 $\pm$ 0.01	0.025 $\pm$ 0.001	77 $\pm$ 2

Appendix

Table 13: Per-length valid-only core metrics of SMILES generation (mean

\pm

95% CI,

L_{\max}=15

$L$	TB		SubTB		RapTB		RapTB+SubM
	Entropy	FPDiv	Entropy	FPDiv	Entropy	FPDiv	Entropy	FPDiv
1	2.45 $\pm$ 0.04	0.82 $\pm$ 0.011	0 $\pm$ 0	0 $\pm$ 0	1.8 $\pm$ 0.06	0.625 $\pm$ 0.009	0.92 $\pm$ 0.03	0.455 $\pm$ 0.024
2	2.57 $\pm$ 0	0.813 $\pm$ 0.004	0.83 $\pm$ 0.21	0.344 $\pm$ 0.071	2.07 $\pm$ 0.12	0.802 $\pm$ 0.004	1.04 $\pm$ 0.2	0.661 $\pm$ 0.025
3	2.49 $\pm$ 0	0.834 $\pm$ 0	1.49 $\pm$ 0.05	0.793 $\pm$ 0.005	1.6 $\pm$ 0.14	0.706 $\pm$ 0.02	1.21 $\pm$ 0.24	0.682 $\pm$ 0.027
4	2.18 $\pm$ 0.03	0.791 $\pm$ 0.009	1.21 $\pm$ 0.08	0.673 $\pm$ 0.012	1.27 $\pm$ 0.23	0.748 $\pm$ 0.04	1.17 $\pm$ 0.06	0.653 $\pm$ 0.026
5	2.42 $\pm$ 0.04	0.868 $\pm$ 0.003	1.25 $\pm$ 0.02	0.748 $\pm$ 0.007	1.72 $\pm$ 0.05	0.825 $\pm$ 0.008	1.32 $\pm$ 0.19	0.793 $\pm$ 0.017

Appendix

Table 14: Per-length valid-only diversity metrics of SMILES generation (mean

\pm

95% CI,

L_{\max}=15

$L$	TB				SubTB				RapTB				RapTB+SubM
	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol	UniqStr	UniqMol	UniqRateStr	UniqRateMol
1	31 $\pm$ 1.2	24.7 $\pm$ 3	0 $\pm$ 0	0 $\pm$ 0	1 $\pm$ 0	1 $\pm$ 0	0.7 $\pm$ 0.2	0.7 $\pm$ 0.2	14 $\pm$ 1.4	11.3 $\pm$ 1.5	0.1 $\pm$ 0	0.1 $\pm$ 0	5.5 $\pm$ 0.6	5.5 $\pm$ 0.6	0.1 $\pm$ 0	0.1 $\pm$ 0
2	189 $\pm$ 3.6	128.7 $\pm$ 4.3	0.2 $\pm$ 0	0.3 $\pm$ 0	7.3 $\pm$ 0.4	6.3 $\pm$ 0.4	0.1 $\pm$ 0	0.1 $\pm$ 0	35.7 $\pm$ 4.2	33 $\pm$ 4.5	0.3 $\pm$ 0	0.3 $\pm$ 0	8.5 $\pm$ 1.7	8.5 $\pm$ 1.7	0.1 $\pm$ 0	0.1 $\pm$ 0
3	215.7 $\pm$ 2.7	203.7 $\pm$ 1.1	0.5 $\pm$ 0	0.5 $\pm$ 0	14.3 $\pm$ 0.8	14.3 $\pm$ 0.8	0.5 $\pm$ 0.1	0.5 $\pm$ 0.1	20.7 $\pm$ 1.5	19.7 $\pm$ 1.8	0.6 $\pm$ 0	0.6 $\pm$ 0	8.5 $\pm$ 0.6	8.5 $\pm$ 0.6	0.2 $\pm$ 0	0.2 $\pm$ 0
4	170 $\pm$ 9.3	164.7 $\pm$ 8.6	0.6 $\pm$ 0	0.6 $\pm$ 0	22 $\pm$ 1.4	20.7 $\pm$ 1.8	0.3 $\pm$ 0	0.3 $\pm$ 0	32 $\pm$ 4.5	31.3 $\pm$ 4.8	0.6 $\pm$ 0.1	0.6 $\pm$ 0.1	15.5 $\pm$ 1.7	15.5 $\pm$ 1.7	0.3 $\pm$ 0	0.3 $\pm$ 0
5	178.7 $\pm$ 6.6	176 $\pm$ 6.8	0.9 $\pm$ 0	0.9 $\pm$ 0	47 $\pm$ 1.2	42.3 $\pm$ 0.4	0.3 $\pm$ 0	0.3 $\pm$ 0	51.3 $\pm$ 11	48.7 $\pm$ 7.7	0.9 $\pm$ 0	0.8 $\pm$ 0.1	23 $\pm$ 0	23 $\pm$ 0	0.3 $\pm$ 0	0.3 $\pm$ 0

Appendix

Table 15: Per-length valid-only uniqueness metrics of SMILES generation (mean

\pm

95% CI,

L_{\max}=15

$k$	TB					SubTB
	Survival	PefEnt	Eff	Top1	UniqueRate	Survival	PefEnt	Eff	Top1	UniqueRate
1	1 $\pm$ 0	2.637 $\pm$ 0.021	13.97 $\pm$ 0.3	0.236 $\pm$ 0.009	0.013 $\pm$ 0	1 $\pm$ 0	1.235 $\pm$ 0.016	3.44 $\pm$ 0.05	0.508 $\pm$ 0.007	0.012 $\pm$ 0.001
2	0.758 $\pm$ 0.007	4.775 $\pm$ 0.027	118.61 $\pm$ 3.24	0.082 $\pm$ 0.005	0.152 $\pm$ 0.008	0.998 $\pm$ 0	3.085 $\pm$ 0.009	21.87 $\pm$ 0.2	0.164 $\pm$ 0.009	0.07 $\pm$ 0.003
3	0.418 $\pm$ 0.005	5.495 $\pm$ 0.048	243.8 $\pm$ 11.8	0.108 $\pm$ 0.008	0.447 $\pm$ 0.01	0.948 $\pm$ 0.002	3.987 $\pm$ 0.031	53.93 $\pm$ 1.7	0.098 $\pm$ 0	0.179 $\pm$ 0.004
4	0.29 $\pm$ 0.006	5.452 $\pm$ 0.052	233.74 $\pm$ 12.43	0.145 $\pm$ 0.008	0.591 $\pm$ 0.016	0.917 $\pm$ 0.002	4.466 $\pm$ 0.035	87.04 $\pm$ 3	0.082 $\pm$ 0.004	0.262 $\pm$ 0.003
5	0.202 $\pm$ 0.004	5.178 $\pm$ 0.092	178.4 $\pm$ 16.91	0.208 $\pm$ 0.014	0.682 $\pm$ 0.02	0.848 $\pm$ 0.007	5.041 $\pm$ 0.093	155.54 $\pm$ 14.12	0.058 $\pm$ 0.008	0.375 $\pm$ 0.016

Appendix

Table 16: Prefix statistics by depth on SMILES generation (mean

\pm

95% CI,

L_{\max}=15

): TB vs. SubTB.

$k$	RapTB					RapTB+SubM
	Survival	PefEnt	Eff	Top1	UniqueRate	Survival	PefEnt	Eff	Top1	UniqueRate
1	1 $\pm$ 0	1.693 $\pm$ 0.014	5.44 $\pm$ 0.08	0.375 $\pm$ 0.004	0.008 $\pm$ 0	1 $\pm$ 0	1.938 $\pm$ 0.001	6.94 $\pm$ 0.01	0.342 $\pm$ 0	0.007 $\pm$ 0
2	0.967 $\pm$ 0.002	3.438 $\pm$ 0.033	31.14 $\pm$ 1.03	0.204 $\pm$ 0.006	0.051 $\pm$ 0.002	0.982 $\pm$ 0.001	3.659 $\pm$ 0.011	38.83 $\pm$ 0.43	0.125 $\pm$ 0.009	0.039 $\pm$ 0
3	0.934 $\pm$ 0.001	4.542 $\pm$ 0.04	93.95 $\pm$ 3.74	0.106 $\pm$ 0.006	0.138 $\pm$ 0.006	0.963 $\pm$ 0.005	4.227 $\pm$ 0.033	68.53 $\pm$ 2.26	0.126 $\pm$ 0.01	0.09 $\pm$ 0.003
4	0.923 $\pm$ 0.001	5.249 $\pm$ 0.026	190.42 $\pm$ 4.86	0.071 $\pm$ 0.002	0.238 $\pm$ 0.004	0.947 $\pm$ 0.003	4.829 $\pm$ 0.009	125.11 $\pm$ 1.11	0.094 $\pm$ 0.003	0.148 $\pm$ 0.002
5	0.906 $\pm$ 0.001	5.822 $\pm$ 0.019	337.73 $\pm$ 6.41	0.07 $\pm$ 0.005	0.345 $\pm$ 0.003	0.931 $\pm$ 0	5.184 $\pm$ 0.001	178.44 $\pm$ 0.24	0.095 $\pm$ 0.003	0.204 $\pm$ 0.002

Appendix

Table 17: Prefix statistics by depth on SMILES generation (mean

\pm

95% CI,

L_{\max}=15

): RapTB vs. RapTB+SubM.

Replay	Objective	Unique ✓	NormCov	Acc	KL( $\pi\!\to\!p^{*}$ )	KL( $p^{*}\!\to\!\pi$ )	JS ${}_{\text{tok}}$
PRT	TB	103.7 $\pm$ 3.2	0.016 $\pm$ 0.001	0.999 $\pm$ 0.000	1.105 $\pm$ 0.002	7.803 $\pm$ 0.060	0.292 $\pm$ 0.001
	SubTB	292.0 $\pm$ 2.9	0.046 $\pm$ 0.000	0.311 $\pm$ 0.002	0.424 $\pm$ 0.010	0.672 $\pm$ 0.077	0.107 $\pm$ 0.003
	RapTB	129.3 $\pm$ 0.4	0.020 $\pm$ 0.000	0.992 $\pm$ 0.001	0.908 $\pm$ 0.003	5.538 $\pm$ 0.005	0.230 $\pm$ 0.001
RP	TB	5.3 $\pm$ 0.4	0.001 $\pm$ 0.000	1.000 $\pm$ 0.000	1.297 $\pm$ 0.001	11.403 $\pm$ 0.282	0.339 $\pm$ 0.000
	SubTB	324.7 $\pm$ 2.7	0.051 $\pm$ 0.000	0.229 $\pm$ 0.005	0.455 $\pm$ 0.005	0.865 $\pm$ 0.083	0.109 $\pm$ 0.002
	RapTB	246.7 $\pm$ 7.1	0.039 $\pm$ 0.001	0.991 $\pm$ 0.000	0.561 $\pm$ 0.001	4.480 $\pm$ 0.002	0.147 $\pm$ 0.000

Appendix

Table 18: Expr24 results under four replay schemes. All the experiments are run under 3 different random seeds with 95% CI. Per-run sample size is 6400.

Replay	Objective	$\ell=3$	$\ell=5$	$\ell=7$	$\ell=9$
RP	TB	–	–	-0.000	-0.001
	SubTB	-5.220	-1.292	-1.064	-79.638
	RapTB	–	-2.451	-2.319	-0.065
	RootSubTBLogZ	-0.709	-0.606	-0.411	-0.068
Oracle	TB	–	-4.391	-1.779	-0.441
	SubTB	-1.017	-2.803	-4.530	-86.415

Appendix

Table 19: Per-length

\log p_{\text{term}}

on Expr24.

TB	SubTB	RapTB	RapTB+SubM
(Len=19) The cat likes to sit on the couch with a pet that lives in a cat, because it.	(Len=20) My pet cat likes to sleep on a couch, I know, what a pet perv! This is.	(Len=15) The pet cat loves the couch, and the cat’s mom was an mascot.	(Len=10) A pet cat likes to sleep on a couch.
(Len=15) A kid loves to dance in a room that a kid likes to dance in.	(Len=20) My kid loves to dance in her room with all of her clothes on except her persnare cal.	(Len=7) I dance in the room for kid.	(Len=8) kid loves to dance in her room!!
(Len=20) There is a hockey team that plays at a rink and plays on a team on a hockey Claus.	(Len=20) The team practiced hockey at the rink, and do freeze to ourselves towards the rinks touring doors.	(Len=15) A team played at the hockey rink with the best players on their own.	(Len=11) The team practiced hockey at the rink for four hours.
(Len=20) There was a flight of stairs, and a jump, a flight of stairs, and a jump,	(Len=20) I jump up a flight of stair, landing in a little flight-recorded levy down to current it.	(Len=14) I jump up a flight of stair to get to school in the morning.	(Len=11) I walk up a stair: jump up a flight.

Appendix

Table 20: Samples from CommonGen task. SubTB often exhibits max-length / off-distribution artifacts, while RapTB+SubM yields shorter and cleaner sentences.

	SMILES				Expr24
Method	Acc $\uparrow$	Score $\uparrow$	Entropy $\uparrow$	Len	Acc $\uparrow$	Unique ✓ $\uparrow$
TB	0.998	0.717	2.503	3.06	1.000	5.3
AvgPrefixTB	1.000	0.661	0.665	2.89	0.998	142.0
RapTB	0.996	0.740	2.448	6.14	0.991	246.7
RapTB+SubM	0.988	0.844	2.726	7.44	0.994	1337.3

Appendix

Table 21: AvgPrefixTB comparison on SMILES and Expr24. All methods use the same backbone (Llama-3.2-1B, LoRA rank 16) and RP replay. SMILES metrics are computed on valid samples; Len denotes mean token length (

L_{\max}=10

Method	Perf $\uparrow$	Div $\uparrow$	Novelty $\uparrow$	Len	Steps
DynaPPO	0.938	12.12	9.31	$\sim$ 20	10K
COMs	0.761	19.38	26.47	$\sim$ 20	10K
GFlowNet	0.868	11.32	15.72	$\sim$ 20	10K
TB	0.927	7.39	10.65	17.4	10K
SubTB †	0.897	21.37	28.68	49.3	10K
RapTB	0.919	8.83	14.44	22.4	5K

Appendix

Table 22: AMP biological sequence generation.

\dagger

: SubTB collapses to max length; its diversity/novelty is inflated by raw edit distance over unnaturally long sequences.

Scale	Method	Acc $\uparrow$	Score $\uparrow$	FPDiv $\uparrow$	Ent $\uparrow$	Len
3B (Llama)	TB	0.999	0.716	0.838	1.92	2.69
	SubTB	0.311	0.222	0.854	2.56	9.52
	RapTB	1.000	0.795	0.839	1.81	7.99
	RapTB+SubM	0.998	0.869	0.936	2.41	8.05
8B (Llama)	TB	1.000	0.715	0.775	1.84	2.98
	SubTB	0.391	0.307	0.869	2.72	9.22

Appendix

Table 23: Scaling study on SMILES generation across model sizes and architectures.

$\beta$	$\rho$	Acc $\uparrow$	Diversity $\uparrow$
1	0.0	0.999	1.010
1	0.1	0.999	0.993
1	0.5	1.000	1.015
3	0.0	0.990	0.950
3	0.1	1.000	0.769
3	0.5	0.997	0.978

Appendix

Table 24:

(\beta,\rho)

sensitivity on Expr24 (RP replay).

$\eta$	Acc $\uparrow$	Diversity $\uparrow$
0.10	0.999	0.983
0.25	0.997	0.978
0.50	0.987	1.149

Appendix

Table 25: Auxiliary weight

\eta

sensitivity on Expr24.

$k_{\min}$ variant	Acc $\uparrow$	Diversity $\uparrow$
Fixed $k_{\min}{=}3$	0.998	0.852
Schedule $7{\to}3$	0.997	0.978
Fixed $k_{\min}{=}7$	0.969	1.025

Appendix

Table 26:

k_{\min}

schedule sensitivity on Expr24.

$\beta$	$\rho$	Acc $\uparrow$	Entropy $\uparrow$
1	0.0	0.992	2.173
1	0.1	0.994	2.161
1	0.5	0.992	2.162
5	0.0	0.995	1.997
5	0.1	0.991	2.076
5	0.5	0.997	2.079

Appendix

Table 27:

(\beta,\rho)

sensitivity on SMILES (RP replay).

Method	Acc $\uparrow$	Score $\uparrow$	Entropy $\uparrow$	FPDiv $\uparrow$	Len
TB	0.998	0.717	2.503	0.807	3.065
TB w/o ref	0.381	0.601	0.418	0.425	10.000

Appendix

Table 28: SMILES generation ablations (TB reference). Unless specified, all metrics are computed on valid samples. Len denotes the mean token length of valid samples (

L_{\max}=10

	SMILES	Expr24
Trainer steps (max)	5000	5000
Precision	Bf16	Bf16
Grad clip	0.5	0.5
Grad accumulation	4	4
Optimizer	AdamW ( $\mathrm{lr}=10^{-4}$ )	AdamW ( $\mathrm{lr}=10^{-4}$ )
Sampling (train)

Appendix

Table 29: Core training and decoding hyperparameters. TB and SubTB share all hyperparameters with RapTB except the loss definition (and RapTB-specific coefficients).

	SMILES	Expr24
Aux weight $\eta$	0.25	0.25
Distance discount $\gamma$	0.99	0.99
Detach $p_{\mathrm{term}}$ in aux	True	True
Absorb gate threshold $\varepsilon_{\mathrm{ab}}$	$10^{-6}$	$10^{-6}$
Target mode	mix	mix
Mix weight $\alpha$ (max vs soft)	0.5	0.8

Appendix

Table 30: RapTB-specific hyperparameters.

k_{\min}

is linearly scheduled by training step.

ID	Concept Set	Domain Context
1	{field, look, stand}	Outdoor / Observation
2	{kid, room, dance}	Indoor / Activity
3	{cat, pet, couch}	Domestic / Animal
4	{climb, building, side}	Urban / Action
5	{climb, wall, talk}	Social / Action
6	{drive, snow, car}	Travel / Environment

Appendix

Table 31: Diagnostic subset of CommonGen. The table lists the 10 concept-sets used to monitor termination drift and prefix entropy. All sets require relational reasoning to link three distinct concepts into a coherent scenario.

Task	Factor	start	end	horizon
SMILES	replay_buffer	0.50	0.25	5000
SMILES	k_min	5	2	5000
Expr24	replay_buffer	0.50	0.25	5000
Expr24	oracle_buffer	0.75	0.25	5000
Expr24	k_min	7	3	5000

Appendix

Table 32: Factor schedulers used in our training runs. Each factor uses linear interpolation from start to end over a fixed horizon (in steps).

Hyperparameter	Value
Buffer capacity $B$	200
Similarity backend	RDKit bulk Tanimoto (Morgan r=2, 2048-bit)
Weights $(w_{\mathrm{rew}},w_{\mathrm{val}},w_{\mathrm{div}},w_{\mathrm{len}})$	$(1.0,1.0,1.0,1.0)$
Length bin size	1 token
Length alpha mode / power	uniform / 1.0
Validity gating ratio	0.0

Appendix

Table 33: Submodular replay hyperparameters (SMILES). We select a buffer of size

B

using greedy maximization of a facility-location + length objective.

Hyperparameter	Value
Buffer capacity $B$	200
Refresh period $K_{\mathrm{buf}}$	1
Similarity backend	$k$ -gram shingles + Jaccard
Weights $(w_{\mathrm{rew}},w_{\mathrm{val}},w_{\mathrm{div}},w_{\mathrm{len}})$	$(1.0,1.0,1.0,0.0)$
Length bin size	-
Length alpha mode / power	-

Appendix

Table 34: Submodular replay hyperparameters (Expr24). We select a buffer of size

B

using greedy maximization of a facility-location + length objective.

Generative Flow Networks

Generative Flow Networks learn to sample diverse objects with probabilities proportional to their rewards.

See all 56 papers

Feb 3, 2026

GFlowPO: Generative Flow Network as a Language Model Prompt Optimizer

Finding effective prompts for language models (LMs) is critical yet notoriously difficult: the prompt space is combinatorially large, rewards are sparse due to expensive target-LM evaluation. Yet, existing RL-based prompt optimizers often rely on on-policy updates and a meta-prompt sampled from a fixed distribution, leading to poor sample efficiency. We propose GFlowPO, a probabilistic prompt optimization framework that casts prompt search as a posterior inference problem over latent prompts regularized by a meta-prompted reference-LM prior. In the first step, we fine-tune a lightweight prompt-LM with an off-policy Generative Flow Network (GFlowNet) objective, using a replay-based training policy that reuses past prompt evaluations to enable sample-efficient exploration. In the second step, we introduce Dynamic Memory Update (DMU), a training-free mechanism that updates the meta-prompt by injecting both (i) diverse prompts from a replay buffer and (ii) top-performing prompts from a small priority queue, thereby progressively concentrating the search process on high-reward regions. Across few-shot text classification, instruction induction benchmarks, and question answering tasks, GFlowPO consistently outperforms recent discrete prompt optimization baselines.

Junmo Cho, Suhan Kim, Sangjune An +5

Korea Advanced Institute of Science and Technology (KAIST) · Korea University · Mila – Québec AI Institute

LLM Input Optimization

Ask PDF

Jul 15, 2026

GFlowRL: Scaling Distribution-Matching RL to Large Language Models

Generative Flow Networks (GFlowNets) offer a promising alternative to reward-maximizing reinforcement learning (RL) for large reasoning models, encouraging diverse reasoning paths by matching reward distributions rather than collapsing to dominant modes. Recent work shows promise on math and code, but scaling GFlowNet-style RL to modern post-training pipelines remains difficult: as model size, rollout horizon, reward noise, and distributed-systems complexity grow together, a learned prompt-conditional partition function becomes a source of gradient instability and engineering overhead rather than a useful normalizer. Through systematic analysis, we find that the learned partition function, previously treated as essential, can be replaced by an in-batch Monte Carlo estimate computed from the rollout group already required for training. We propose GFlowRL, a streamlined GFlowNet-style RL algorithm that removes the auxiliary partition network entirely while preserving the reward-distribution-matching objective, completed by two stabilizers: importance-sampling correction for rollout/trainer drift and asymmetric flow-gap clipping for outlier residuals. GFlowRL exceeds all counterparts on math, code, and adversarial red-teaming benchmarks, reaching a Codeforces rating of 2048 at the 14B scale (within 25 Elo of o3-mini) and attaining the highest average ASR@1 on AdvBench and HarmBench, outperforming the previous SOTA multi-turn attacker in a regime where FlowRL, a prior GFlowNet-style method, diverges. The same recipe transfers to all evaluated MoE configurations up to 235B parameters, where FlowRL again fails to converge. To our knowledge, GFlowRL is the first GFlowNet-style RL algorithm to scale stably across both dense and sparse architectures. Code will be at: https://github.com/microsoft/gflowrl

Xiaodong Liu, Michael Xu, Jack W. Stokes +3

Microsoft Research

RL for Code Generation Reinforcement Learning RL for Language Model Reasoning Generative Flow Networks RL Post-Training

Ask PDF

Jan 28, 2026

Signal from Structure: Exploiting Submodular Upper Bounds in Generative Flow Networks

Generative Flow Networks (GFlowNets; GFNs) are a class of generative models that learn to sample compositional objects proportional to their a priori unknown value, their reward. We focus on the case where the reward has a specified, actionable structure, namely that it is submodular. We show submodularity can be harnessed to retrieve upper bounds on the reward of compositional objects that have not yet been observed. We provide in-depth analyses of the probability of such bounds occurring, as well as how many unobserved compositional objects can be covered by a bound. Following the Optimism in the Face of Uncertainty principle, we then introduce SUBo-GFN , which uses the submodular upper bounds to train a GFN. We show that SUBo-GFN generates orders of magnitude more learning signal than classical GFNs for the same number of queries to the reward function. We demonstrate the effectiveness of SUBo-GFN in terms of distribution matching and high-quality candidate generation on graph submodular tasks.

Alexandre Larouche, Audrey Durand

D´epartement d’informatique et de g´enie logiciel, Universit´e Laval, Qu´ebec, Canada · Mila - Quebec AI Institute, Montr´eal, Canada · Canada CIFAR AI Chair

Generative Flow Networks Graph Generation

Ask PDF

Abstract

Generative Flow Networks

Explore similar work

GFlowPO: Generative Flow Network as a Language Model Prompt Optimizer

GFlowRL: Scaling Distribution-Matching RL to Large Language Models

Signal from Structure: Exploiting Submodular Upper Bounds in Generative Flow Networks