Diffusion-Based Data-Driven Assortment Optimization

Diffusion-Based Data-Driven Assortment Optimization proposes D3AO, a model-agnostic framework for assortment optimization in the offline setting, based on guided discrete diffusion. The paper addresses the fundamental problem of assortment optimization in revenue management, where a seller selects a subset of products to offer to maximize expected revenue, noting that traditional parametric choice models like the multinomial logit (MNL) are sensitive to model misspecification and struggle to capture complex customer behavior.

The paper formalizes the offline assortment optimization problem, where only historical interaction data D = (Si, Ai) ni=1 is available, with Si being previously offered assortments and Ai the observed customer choices. The authors assume the behavior policy follows a Boltzmann (softmax) distribution µ(s) ∝ exp(β0 R(s)), which they justify through Theorem 2.1, showing that such policies arise as solutions to entropy-regularized revenue maximization problems.

The proposed D3AO framework consists of three stages: (i) learning a neural choice model pθ(as) from offline data via maximum likelihood estimation, (ii) constructing a model-based reward estimator Rθb(s) = Σa∈s∪ 0 r(s,a) pθb(as), and (iii) performing guided stochastic optimization via a diffusion-based generative model. The diffusion model serves as a data-driven prior over feasible assortments, while the reward signal acts as a guiding force that progressively steers samples toward high-reward regions.

The method represents assortments as binary vectors and uses a discrete-time Bernoulli corruption process for the forward diffusion. The reverse process is parameterized by a neural network that predicts clean assortments from noisy states. For reward guidance, the paper introduces a coordinate-wise approximation based on local perturbations, defining the marginal reward difference ∆i Rθb(st):= Rθb(st(i=1)) - Rθb(st(i=0)), which measures the incremental value of including item i given the current context. The guided transition modifies the reverse logits as geϕ(i)(st, t) = gϕ(i)(st, t) + λt ∆i Rθb(st), with a time-dependent guidance strength λt = λmax(1 - (t-1)/(T-1))γ that suppresses guidance for highly noisy states and applies strongest guidance near the end of the reverse process.

The paper provides several theoretical properties. Proposition 4.1 shows that the forward corruption process approaches the uniform distribution at a controlled finite-time rate, with ∥qt(·s0) - u∥TV ≤ (N/2)ᾱt. Proposition 4.3 characterizes the population minimizer of the denoising objective as the posterior clean-sample marginal induced by the forward corruption process. Theorem 4.4 shows that the reward-guided logit update is equivalent to a KL-regularized local policy improvement step, where the guided transition is the unique optimizer of maxq∈∆(0,1 N) Ex∼q[lst(x)] - (1/λt)KL(q ∥ pϕ(·st, t)).

Empirically, the paper evaluates D3AO on synthetic datasets generated from three choice models: MNL, Markov Chain Choice Model (MCCM), and Mixed Multinomial Logit (MMNL), with product counts N ∈ 20, 40, 60, 80, 100 and offline datasets of size n = 10,000. The method is compared against parametric baselines MNL-MLE and MCCM-EM, which are correctly specified for MNL and MCCM data respectively but misspecified for MMNL data.

Results show that under MNL data, guided diffusion methods (GAsN-guided and RAsN-guided) remain highly competitive with the correctly specified MNL-MLE baseline, achieving maximum optimal ratios consistently close to 100%. Under MCCM data, both guided variants match or slightly exceed the oracle baseline MCCM-EM. Under MMNL data with model misspecification, guided diffusion methods significantly outperform the misspecified parametric baselines, demonstrating robustness to model mismatch.

The paper also evaluates exact recovery of optimal assortments and diversity of generated solutions. Exact recovery rates decrease markedly as N increases, but guided variants retain nontrivial recovery probabilities in low- and medium-dimensional regimes (e.g., 68.16% at N=20 and 32.73% at N=40 for GAsN-guided on MCCM data with β=1.0). Diversity metrics, including unique sample ratio and average pairwise Hamming distance, tend to increase as N grows, indicating that guided sampling becomes more diverse in larger combinatorial spaces.

Sample efficiency experiments show that under MNL data, regret is uniformly low and generally decreases with sample size, while under MMNL data, the misspecified MNL-MLE baseline exhibits substantially larger regret across all sample sizes, with diffusion-based methods remaining markedly stronger. Distribution-shift experiments, where historical assortments are generated from a different policy (uniform size then subset procedure), show that while performance drops noticeably, maximum optimal ratios remain relatively high, indicating the method preserves strong best-case solution quality under distribution shift.

The paper concludes that D3AO consistently produces high-quality assortments, remains effective under model misspecification, and naturally yields a diverse set of near-optimal solutions. The authors note that performance depends on the structure of offline data and quality of reward estimation, and suggest future work including developing methods that reduce dependence on data regularity, extending the framework to more general constraints, and exploring connections with continuous relaxation approaches such as Langevin dynamics.

Improvements to AI Systems Based on This Paper:

1. Robust Offline Decision-Making Under Model Misspecification

• What: Replace parametric choice models (e.g., MNL) with a neural choice model trained via maximum likelihood, then use a guided discrete diffusion process for optimization.

• Improved AI system: Can make optimal assortment recommendations (or similar combinatorial decisions) directly from historical interaction data without assuming a specific customer behavior model, remaining accurate even when the true choice process is unknown or complex (e.g., mixed MNL, Markov chain). This is critical for real-world revenue management where customer behavior is rarely captured by simple parametric forms.

2. Guided Generative Optimization with Uncertainty-Aware Steering

• What: Use a diffusion model as a data-driven prior over feasible solutions, with a time-dependent guidance strength λt that suppresses noise early and applies strong reward-based steering late in the reverse process.

• Improved AI system: Can generate a diverse set of near-optimal solutions (e.g., product assortments) that balance exploration and exploitation, even in high-dimensional combinatorial spaces (N up to 100). The system can explicitly control the trade-off between solution quality and diversity, producing multiple valid alternatives—useful for scenarios requiring flexibility (e.g., supply chain constraints, customer segmentation).

3. Sample-Efficient Learning from Offline Data

• What: The framework learns directly from a fixed dataset of past assortments and choices (n=10,000) without online interaction, using a diffusion-based generative model to approximate the data distribution.

• Improved AI system: Can be deployed in settings where online experimentation is costly or impossible (e.g., healthcare, retail pricing). It achieves low regret even with limited data, and outperforms misspecified parametric baselines under distribution shift—meaning it can generalize to new customer segments or market conditions without retraining.

4. Theoretical Guarantees for Guided Diffusion Optimization

• What: Leverage Theorem 4.4, which shows the guided logit update is equivalent to a KL-regularized local policy improvement step.

• Improved AI system: Can be designed with provable convergence properties—each reverse step is a principled optimization step that improves expected reward while staying close to the data distribution. This enables safe deployment in high-stakes settings where guarantees on solution quality are required, and allows for principled hyperparameter tuning (e.g., λmax, γ) based on theoretical insights.

5. Exact Recovery and Diversity in Combinatorial Spaces

• What: The method achieves high exact recovery rates (e.g., 68% at N=20, 33% at N=40) and increasing diversity (Hamming distance) as N grows.

• Improved AI system: Can be used for tasks beyond assortments, such as feature selection, portfolio optimization, or network design, where the system must find the exact optimal subset while also offering diverse alternatives for human decision-makers to review. The diversity metric can be used as a built-in diagnostic for solution robustness.

6. Distribution-Shift Resilience

• What: The method maintains strong best-case solution quality even when the historical data comes from a different policy (e.g., uniform size then subset procedure).

• Improved AI system: Can be deployed in non-stationary environments where past data was collected under a legacy or exploratory policy. The system can still identify high-reward solutions without needing to re-collect data, making it suitable for rapid adaptation to new market conditions or user preferences.

7. Model-Agnostic Integration with Existing Pipelines

• What: The framework is model-agnostic—it only requires a differentiable reward estimator and a binary representation of decisions.

• Improved AI system: Can be plugged into existing recommendation engines, pricing systems, or logistics planners as a drop-in optimizer, replacing brittle parametric models with a data-driven generative approach. It can also be extended to handle constraints (e.g., budget, capacity) by modifying the diffusion prior, enabling broader industrial adoption.