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

1 Introduction

• Reward-guided fine-tuning risks early reward-model bias and collapses onto narrow high-reward regions [1] [2]; SMC methods suffer weight degeneracy and reduced diversity [3] [4] [2], while tree samplers struggle to find diverse low-likelihood targets [6] [7] [8].Source: §1 · pp. 1–2
• Figure 1 contrasts LATS with SMC and DTS on MNIST targets 7 and 9, depicting tail-mode discovery without off-manifold samples and with intra-target diversity.The caption attributes noisy reward-gradient guidance to SMC and Gaussian bias to DTS.Source: Figure 1 · p. 2
• The goal is to discover diverse target-aligned samples from long-tailed data using limited interactive feedback while preserving sample fidelity.Targets may have heterogeneous properties, reside in the tail, and be revealed only through feedback.Source: §1 · p. 2
• LATS combines Lévy diffusion’s exploration of diverse modes with tree-based pruning of unlikely branches to concentrate sampling on target-aligned modes.The introduction presents the two mechanisms as complementary: Lévy diffusion explores, while the tree sampler exploits feedback by pruning branches unlikely to yield preferred samples.Source: §1 · p. 2
• LATS is presented as discovering low-likelihood target modes while preserving sample quality and generating diverse samples across long-tail preference modes.The introduction reports validation across diverse settings, including material science.Source: §1 · p. 2

Fine-Tune Diffusion model with RL

• Online feedback distinguishes this setting from methods that assume static rewards: target properties are initially unknown and revealed sequentially.The cited prior online learning framework [18] is not tailored to diffusion models; a diffusion fine-tuning method [1] uses a separate online reward model and remains vulnerable to early bias and mode collapse.Source: §2 · p. 3

Inference-time Search Approaches

• SMC uses particle weights and repeated resampling, which can reduce diversity under weight variance or inaccurate value estimates [2] [21] [22].In online settings, poorly calibrated early reward models can bias resampling and degrade generated-sample quality.Source: §2 · p. 3
• Tree-based samplers backpropagate terminal rewards to avoid intermediate reward evaluations [7] [23] [8], but Gaussian bias constrains access to preferred long-tail modes.Particle-correction methods [4] [5] require reward gradients at every resampling step, which must be approximated in the described setting.Source: §2 · p. 3

3 Problem Formulation

• The setup assumes access to a long-tailed dataset of i.i.d. samples from an unknown data distribution.The paper frames this as a model of real-world data, where frequent patterns coexist with many rare ones.Source: §3 · p. 3
• Target properties may be heterogeneous and multimodal, including long-tail modes, and are revealed through costly interactive feedback under a fixed query budget.The objective is to allocate the budget efficiently to learn target structure and generate aligned samples.Source: §3 · p. 3
• Unlike direct fine-tuning, the approach keeps the initial generative model trained on the long-tail dataset unchanged during interaction.The unknown probabilistic reward is 1 when a sample satisfies a target property and 0 otherwise; the sampler explores and exploits information gathered from feedback.Source: §3 · pp. 3–4

α-Stable Distributions and Lévy Process

• α-stable distributions generalize Gaussian laws and arise as limits of sums of i.i.d. variables with infinite variance.The background contrasts this with the Gaussian limit for i.i.d. variables with finite variance.Source: §4 · p. 4
• For an α-stable distribution, α controls tail heaviness: α = 2 gives a Gaussian with variance 2σ^2, while α < 2 yields power-law tails and infinite variance.The stated asymptotic tail decay for α < 2 is |x|^−(1+α).Source: §4 · p. 4

Isotropic α-stable distribution and relation to Gaussian Noise

• An isotropic α-stable random variable in R^d is described by a characteristic function involving the location parameter and scale σ.The passage identifies σI_d as playing the role of a covariance matrix.Source: §4 · p. 4

• The forward process replaces Gaussian noise with α-stable noise to better capture long-tail structures in the denoising model.Source: §5 · p. 5
• α-stable kernels retain an algebraic minority-mode signal over distance, unlike Gaussian kernels’ exponentially suppressed signal, keeping training sensitive to rare classes.Source: §5 · p. 5
• Under positively correlated reward weighting, online fine-tuning drives the generated distribution toward the majority class and the tail-class probability decays exponentially.Source: §5 · p. 6
• The sampler builds a finite-horizon tree over Lévy transitions, with each node storing its state, timestep, visit count, and Monte Carlo soft value.Source: §5 · p. 6
• At each interaction, SelectUCT guides tree traversal, β balances exploration and exploitation, and classifier feedback supplies target-class probabilities for terminal rewards and upward value backups.Source: §5 · p. 6
• At B = 250, Table 1 reports FID, MMD, and QADAS diversity across two target sets for each dataset.Source: Table 1 · p. 7

Evaluation Metrics

• QADAS rewards sample quality and diversity across target classes, addressing how success rate alone can overlook collapse onto one high-likelihood mode.The evaluation also uses FID and MMD, with reference distributions drawn equally from target classes.Source: §6 · p. 7

Dataset Details and Evaluation Setting

• The evaluation spans MNIST-LT, CIFAR-LT, ImageNet-LT, and MICRO2D, across diverse target sets and sampling budgets.MICRO2D is described as a large-scale heterogeneous microstructure informatics dataset.Source: §6 · p. 7

Results with Target Set that only includes Rare classes

• Across rare-only S2 settings, LATS consistently surpasses baselines, preserving sample quality while revealing broader high-utility target modes.The reported MICRO2D table covers Set S2.Source: §6 · p. 7, Table 2 · p. 7

Results with Target Set that is a Mixture of Frequent and Rare classes

• On mixed head–tail S1 targets, LATS surpasses all baselines across evaluation metrics, while baselines collapse onto the high-utility, high-likelihood head class.SMC underperforms because early reward-model bias degrades image quality and alignment; DTS guidance suppresses long-tail exploration and diversity.Source: §6 · p. 8

Effect of Lévy Process on Exploration

• Lévy exploration reaches high-utility tail modes that Gaussian dynamics fail to access, and tree-based online value backup amplifies their sampling frequency.The comparison with DTS is reported using Tables 1 and Figure 5.Source: §6 · p. 8

Performance Comparison at different Sampling Budgets

• Across datasets and metrics, LATS converges fastest in quality and target discovery while sustaining markedly higher sample diversity as the sampling budget increases.Figure 4 compares methods on MNIST-LT, CIFAR-10-LT, and ImageNet-LT for Set S2.Source: §6 · p. 8, Figure 4 · p. 8

A PROOF OF PROPOSITION 1

• N denotes the total number of training images, while N^i counts images belonging to class c_i.These counts give the class-frequency weights in the rewritten objective.Source: A PROOF OF PROPOSITION 1 · p. 11
• Across diffusion steps, the rewritten DDPM objective weights each class-specific KL denoising term by N^i/N.The derivation factors the objective into a sum over classes, each with its own expected denoising loss.Source: A PROOF OF PROPOSITION 1 · p. 11
• The proof expresses the original DDPM objective as an expectation over each class, weighted by that class's fraction of the dataset.The displayed sum separates the class weights from the class-conditional KL losses.Source: A PROOF OF PROPOSITION 1 · p. 11

B PROOF OF THEOREM 2

• The induced distribution is proportional to the weighted data distribution multiplied by an exponential factor of the negative loss.The proof defines the weighted distribution as w(x_0)P(x_0), then writes the induced distribution as proportional to that quantity times exp(-γL(x_0)).Source: B PROOF OF THEOREM 2 · pp. 11–12
• Assuming the Linduced(x0; θ) loss converges to 0 after the fine-tuning step, P^1_θ(x0) is proportional to w(x0)P(x0), and normalization by Z gives P^1_θ(x0)=w(x0)P(x0)/Z.Source: B PROOF OF THEOREM 2 · p. 12
• Repeating the update for H online interaction steps yields a distribution proportional to w(x_0)^H P(x_0), normalized by Z^H.The proof obtains this expression by iterating the one-step relation.Source: B PROOF OF THEOREM 2 · p. 13
• After N epochs, the majority-to-tail probability ratio diverges, so tail-class mass tends to zero relative to majority-class mass.The proof concludes that this exacerbates the initial long-tail imbalance with every fine-tuning epoch.Source: B PROOF OF THEOREM 2 · p. 13

C.1 The Non-Local Score Gradient

• Lévy's polynomially decaying score signal preserves a connection between distant modes, unlike the Gaussian signal that can vanish at minority modes.The proof links this non-local influence to reverse-process sampling that remains aware of disjoint modes.Source: §C.1 · pp. 13–14

D Additional Discussion on α-Stable Distribution and comparison with Gaussian

• Mode collapse is defined as failure to assign mass to part of the target distribution when a low-density region separates it from the rest.The discussion frames movement between modes in terms of a diffusion transition kernel.Source: D Additional Discussion on α-Stable… · pp. 14–15

1. Gaussian Case (α = 2)

• For Gaussian DDPM, long-range dependencies are exponentially suppressed, making transitions between distant modes less likely.The text states that for large separation, the reverse process cannot jump between modes and may collapse into the nearest local optimum.Source: §1 · p. 15

2. DLM Case (α < 2)

• For any finite distance, some α in (0, 2) gives DLM a higher probability of reaching a distant mode than Gaussian DDPM.As L grows, the DLM-to-DDPM reach-probability ratio is proportional to L^-α/e^-L^2 and tends to infinity.Source: §2 · p. 15

E Impact of UCT-based Exploration on Performance

• At budget 250, β=0 disables exploration and significantly reduces performance, β=0.1 recovers it, and β=1 achieves the strongest overall results.Table 3 evaluates mixed frequent and rare target classes in Set S1 across datasets.Source: E Impact of UCT-based Exploration on… · p. 15, Table 3 · p. 15
• Removing UCT-style exploration lowers QADAS, while removing online value backup also yields lower QADAS and weakens feedback-aware global search.Table 4 evaluates mixed frequent and rare target classes in Set S1 with a sampling budget of 1000.Source: E Impact of UCT-based Exploration on… · p. 16, Table 4 · p. 16

F.1 Datasets

• MNIST-LT and CIFAR10-LT use an exponential long-tail profile, while ImageNet-LT and MICRO2D use specified class subsets with progressively smaller class counts.The ImageNet-LT subset uses 10 classes at 64 × 64 resolution; MICRO2D uses seven microstructure classes.Source: §F.1 · p. 16

F.2 Model Hyperparameters

• The diffusion UNet, training schedule, and reward classifier vary by dataset, while all models use Adam and 100-step DDIM inference.Training uses 1000 epochs for MNIST-LT, CIFAR10-LT, and ImageNet-LT, and 400 epochs for MICRO2D.Source: §F.2 · p. 16
• For each dataset, DLM and DDIM share UNet architectures and training and diffusion hyperparameters; all DLMs use α = 1.7.Source: §F.2 · p. 16

G.1 Soft Value Function (Equation 9)

• The soft value recursion expresses the current value as a log-expectation of the exponentiated next-step value under the reverse diffusion transition, with terminal value equal to reward.The derivation is included for completeness and cited to [7].Source: §G.1 · p. 17, G Derivation of Equation 9 · p. 17

H Additional Qualitative Comparisons across Different Settings

• LATS consistently produces high-quality samples across target classes, including rare, high-utility long-tail categories.The qualitative comparisons cover MNIST-LT and ImageNet-LT examples in Figures 5 and 6.Source: H Additional Qualitative Comparisons… · p. 17, Figure 5 · p. 18, Figure 6 · p. 19
• FKS often degrades sample fidelity, whereas DTS largely preserves visual quality but fails to reliably generate high-utility tail classes.Source: H Additional Qualitative Comparisons… · p. 17