LATS: Levy Adaptive Tree Sampling for Feedback-Driven Diverse Target Discovery

Binglin Ji; Anindya Sarkar; Hengchang Lu; Lecheng Kong; Yixin Chen; Yevgeniy Vorobeychik

cs.LGSep 6, 2026

LATS: Levy Adaptive Tree Sampling for Feedback-Driven Diverse Target Discovery

Authors: Binglin Ji, Anindya Sarkar, Hengchang Lu, Lecheng Kong, Yixin Chen, Yevgeniy Vorobeychik

Organizations: Department of CSE, Washington University in St.Louis, USA

Abstract

While diffusion models excel at capturing complex data distributions, scientific discovery often requires steering generation toward specific, uncharacterized regions that maximize a target objective. These high-utility modes frequently reside in low-likelihood tail regions and are only revealed sequentially through interactive feedback. Existing diffusion samplers fail in this regime: they inherit the pre-trained model's bias toward high-density regions, leaving rare yet promising phenomena underexplored. Conversely, exploration-heavy samplers ensure broad coverage but fail to efficiently exploit high-utility modes when constrained by a strict sampling budget. To resolve this dilemma, we introduce Levy Adaptive Tree Search (LATS), a principled sampling framework for online feedback-driven search. LATS leverages heavy-tailed exploration coupled with tree-based value backpropagation to progressively uncover preferred modes. By maintaining broad distributional coverage, LATS successfully discovers low-likelihood, high-utility regions while preserving sample fidelity and structural diversity. Experiments across diverse benchmarks, including materials science, demonstrate that LATS significantly outperforms baselines in target discovery efficiency.

Figures & tables

Figure 1: Unlike SMC (which suffers from noisy reward gradient guidance) and DTS (which is constrained by Gaussian bias), LATS discovers high-utility tail modes (Viz., using MNIST digits 7 and 9 as targets, which are the two least represented categories during training) while avoiding off-manifold sample generation. LATS leverages online feedback to enrich intra-target class diversity actively, thereby avoiding mode collapse.

Panel of Figure 2 — Figure 2: Overview of the Lévy Adaptive Tree Search (LATS) Framework. (Stage 1) The generative model uses heavy-tailed Lévy noise rather than Gaussian noise, extending structural coverage to the data manifold’s rare long-tail modes. (Stage 2) LATS constructs a dynamic search tree built upon the trained Lévy transition dynamics to systematically explore the latent space. (Stage 3) Guided by feedback, a soft Bellman value backup mechanism propagates reward signals upward, allowing the $\text{Select}^{\text{UCT}}$ score to aggressively prune low-value paths while continuously amplifying the discovery of high-utility targets without fine-tuning the Lévy model.

Target	Method	MNIST-LT			CIFAR-10-LT			ImageNet-LT
Target	Method	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$
Set 1	SMC	1.8137	13.748	11.8700	3.7810	3.8231	9.9300	2.1463	2.8476	12.9748
	DTS	1.2532	9.4025	14.8631	0.9021	3.1071	13.0262	0.9444	1.7802	18.8476
	FKS	1.0699	7.5111	14.0710	0.8150	2.4711	11.4592	1.4609	2.0285	17.4755
	LATS	0.7264	5.5532	18.7452	0.7223	1.7286	14.0967	0.7588	0.7692	20.1346
Set 2	SMC	4.0168	18.9856	6.7324	3.7970	6.8242	0.0000	3.4637	6.7032	5.2677

Table 1: Quantitative Comparison with

\mathcal{B}=250

. We report FID (lower is better), MMD (lower is better), and the QADAS diversity metric (higher is better) across two distinct target sets for each dataset. LATS consistently outperforms baselines in maintaining sample quality while uncovering diverse target modes.

Method	MICRO2D ( $\mathcal{B}=250$ )
Method	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$
SMC	4.2154	29.8165	0.0000
DTS	2.8699	17.7681	4.6196
FKS	3.5499	27.7336	0.0000
LATS	0.9464	4.2721	6.1041

Table 2: MICRO2D for Set S2.

Panel of Figure 3 — Figure 3: Qualitative comparison of generation diversity and quality across target sets and datasets. While DTS and FKS tend to focus on dominant modes, LATS uncovers diverse, high-utility samples by leveraging Lévy-driven exploration and tree-based online value backup, aiding history-aware exploitation while maintaining sample quality. long-tail high-utility classes are denoted as blue .

Panel of Figure 4 — Figure 4: Quantitative Comparisons across Sampling Budgets. Comparison of LATS against baselines across MNIST-LT, CIFAR-10-LT, and ImageNet-LT for target Set S2. LATS achieves the fastest convergence in quality and target discovery while maintaining significantly higher sample diversity as the sampling budget increases.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Method	CIFAR-10-LT			ImageNet-LT
Method	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$
LATS w $\beta$ = 0	0.8657	2.2700	12.6108	0.8976	0.8637	18.7534
LATS w $\beta$ = 0.1	0.8264	2.0361	13.3479	0.7846	0.8211	18.9304
LATS w $\beta$ = 1	0.7223	1.7286	14.0967	0.7588	0.7692	20.1346

Appendix

Table 3: Analysis of Different Exploration Factor

\beta

. We evaluate performance under different Exploration factor

\beta

with a sampling budget of

250

. For this study, the target sets on each dataset include a mix of frequent and rare classes (i.e., Set

S1

Method	CIFAR-10-LT			ImageNet-LT
Method	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$	FID $\downarrow$	MMD $\downarrow$	QADAS $\uparrow$
LATS w/o Exploration term in node selection mechanism	0.8657	2.2700	12.6108	0.8976	0.8637	18.7534
LATS w/o Online Value Backup	1.0877	4.0201	6.8011	1.6677	6.1229	8.5721
LATS	0.7223	1.7286	14.0967	0.7588	0.7692	20.1346

Appendix

Table 4: Analysis of Different Components of LATS. We evaluate performance under a sampling budget of 1000. For this study, the target sets on each dataset include a mix of frequent and rare classes (i.e., Set

S1

Panel of Figure 5 — Figure 5: Qualitative comparison of generation diversity and quality across target sets and datasets (MNIST-LT, ImageNet-LT examples). We compare LATS against the DTS, FKS baseline on MNIST-LT and ImageNet-LT on representative target sets.

Panel of Figure 6 — Figure 6: Qualitative comparison of generation diversity and quality across target sets and datasets (ImageNet-LT examples). We compare LATS against the DTS, FKS baseline on ImageNet-LT on representative target sets.

Adaptive Sampling

Adaptive sampling allocates limited data, measurements, or computation where they are expected to be most useful.

See all 202 papers

BFMT: Enhancing Search Capabilities of Tree Sampler via Bootstrap Flow-Map Tree
Jul 3, 2026Binglin Ji, Anindya Sarkar, Hengchang Lu +4Inference-Time Search Image Generation
Department of CSE, Washington University in St.Louis, USA · Department of Information Technology, Uppsala University, Sweden
Bootstrap Stochastic Interpolant Tree for Inference-time Search Under Sparse Reward
Oct 16, 2026Binglin Ji, Anindya Sarkar, Yevgeniy VorobeychikInference-Time Search Stochastic Interpolants
Department of Computer Science & Engineering, Washington University in St. Louis, USA
Adaptive teachers for amortized samplers
Oct 2, 2024Minsu Kim, Sanghyeok Choi, Taeyoung Yun +7Amortized Inference RL Exploration
Mila, KAIST · KAIST · Recursion +1

Jul 3, 2026J/K move · Enter open · S save

BFMT: Enhancing Search Capabilities of Tree Sampler via Bootstrap Flow-Map Tree

Binglin Ji, Anindya Sarkar, Hengchang Lu, Kunyu Wang+3

Department of CSE, Washington University in St.Louis, USA · Department of Information Technology, Uppsala University, Sweden

Inference-Time Search Image Generation Adaptive Sampling Molecule Design

Open paper AskPDF arXiv

The ability to efficiently navigate high-dimensional spaces to identify optimal candidates-whether designing targeted molecular structures or generating images with preferred attributes-stands as a central pillar of modern scientific progress. Recent tree-based samplers enable highly scalable inference-time search by eliminating the reward-gradient bottleneck inherent to particle-based methods. Despite their potential, existing tree-based samplers are fundamentally constrained by the massive number of function evaluations (NFEs) required for node valuation, crippling their utility in exploration-heavy tasks. Furthermore, their reliance on small, uniform transitions at each depth precludes dynamic, adaptive search capabilities. To address this, we introduce Bootstrap Flow-Map-Tree (a.k.a BFMT), a computationally efficient sampling framework that enables full tree-path construction from any tree depth using a single function evaluation, drastically reducing computational overhead while providing critical foresight for sequential sampling. By enabling dynamic transition time-step scheduling, BFMT efficiently allocates its sampling budget, smoothly transitioning from broad global exploration to fine-grained local refinement of high-utility modes discovered through exploration. Extensive experiments and ablation studies across diverse domains-from high-dimensional images to large-scale molecules-demonstrate BFMT's superiority over baseline approaches in both search and alignment tasks.

Abstract

Adaptive Sampling

Explore similar work

BFMT: Enhancing Search Capabilities of Tree Sampler via Bootstrap Flow-Map Tree

Bootstrap Stochastic Interpolant Tree for Inference-time Search Under Sparse Reward

Adaptive teachers for amortized samplers