Speeding up the ordered allocation sampler

1. Problem Statement

The paper addresses challenges in Bayesian nonparametric mixture models, specifically regarding posterior inference when the number of mixture components is unknown or infinite.

• The Trade-off: Existing Gibbs sampling methods fall into two categories:
  • Marginal Samplers: Integrate out mixing weights and parameters, evolving in the space of partitions. They possess excellent mixing properties but require a tractable predictive distribution (specifically a generalized Pólya urn scheme), which limits their applicability to specific priors (e.g., Dirichlet Process, Pitman-Yor).
  • Conditional Samplers: Condition on weights and parameters. They are widely applicable to various priors (including those without tractable predictive structures) but evolve in the space of cluster labels rather than partitions. This leads to poor mixing properties and slow convergence.
• The Specific Bottleneck: The Ordered Allocation Sampler (OAS), proposed by De Blasi and Gil–Leyva (2023), was a breakthrough conditional sampler that improved mixing by conditioning on parameters in the order of their discovery. However, it suffered from three critical issues:
  1. Data Ordering Sensitivity: The initial ordering of data points significantly affected mixing because the "least element order" constraint restricted the movement of early data points.
  2. Limited Component Variation: Due to the ordering constraint, the sampler could not easily empty early clusters or create new ones in a single scan, leading to slow mixing of the number of occupied components ($k_n$).
  3. Computational Complexity: Updating allocation variables required identifying a set of "admissible moves" at every step to preserve the least element order, making the algorithm computationally expensive and difficult to implement.

2. Methodology

The authors propose a modified Efficient Ordered Allocation Sampler (OAS) that retains the wide applicability of conditional samplers while achieving the mixing properties of marginal samplers.

Core Innovation: Unordered Updates via Permutation

The key theoretical insight is that while the OAS operates on ordered allocation variables (based on discovery order), the data is exchangeable. The authors leverage this to update allocation variables as if they were unordered, effectively allowing the sampler to evolve in the space of partitions rather than ordered partitions.

• Mechanism: Instead of updating the allocation variable $d_i$ for observation $y_i$ directly (which requires checking admissible moves), the algorithm conceptually permutes the data so that $y_i$ becomes the last observation in the sequence.
• Implementation:
  • The algorithm uses unordered allocation variables $c_{1:n}$ to represent the current clustering.
  • For each data point $y_i$, the algorithm samples $c_i$ from a distribution proportional to the weight of existing components plus the probability of a new component (similar to marginal samplers like Neal's Algorithm 8).
  • Crucially, the weights are handled via size-biased permutation (ISBP) properties. The algorithm maintains the weights in order of appearance ($\tilde{p}$) but updates them by conditioning on a permutation of the underlying weights ($p$) and their discovery indices ($\alpha$).
  • This approach eliminates the need to compute "admissible moves" sets ($D_i$) and removes the dependency on the initial data ordering.

Split-Merge Moves

Exploiting the similarity between the new OAS and marginal samplers, the authors adapt Split-Merge moves (originally by Jain and Neal) to the OAS framework.

• Purpose: To help the Markov Chain escape local modes where the sampler gets stuck in suboptimal partitions (e.g., splitting a single true component into two or merging two distinct components).
• Adaptation: The authors define a "neighbourhood" of states involving split or merge operations. They use restricted Gibbs sampling scans to propose new states within this neighbourhood.
• Generalization: The method is extended to priors without tractable size-biased weights (e.g., Exchangeable Stick-Breaking, Geometric Process) by updating the discovery indices ($\alpha$) and weights ($p$) alongside the allocation variables.

3. Key Contributions

1. Algorithmic Enhancement: A new version of the OAS that updates allocation variables sequentially as unordered variables. This resolves the mixing issues related to data ordering and the inability to vary the number of components efficiently.
2. Implementation Simplification: The new method removes the need to calculate admissible moves sets, significantly reducing computational overhead and simplifying the code.
3. Integration of Split-Merge Moves: The derivation of split-merge moves tailored for the OAS, applicable even to mixing priors that do not admit marginalization (a significant advancement over previous split-merge implementations limited to Dirichlet/Pitman-Yor processes).
4. Theoretical Justification: Proof that the new sampler evolves in the space of partitions, inheriting the mixing benefits of marginal samplers while retaining the flexibility of conditional samplers.

4. Results

The authors conducted extensive simulation studies comparing the new OAS against:

• The original OAS (oOAS).
• Marginal samplers (where applicable).
• Standard conditional samplers (e.g., Slice samplers).

Findings:

• Mixing Efficiency: The new OAS significantly outperforms the original OAS and standard conditional samplers. The Integrated Autocorrelation Time (IAT) for the number of components ($k_n$) was drastically reduced.
• Comparison with Marginal Samplers: While marginal samplers remained the most efficient (lowest IAT), the new OAS closed the gap significantly, often performing within a factor of 2 of the marginal sampler, whereas the original OAS was much slower.
• Computational Speed: The new OAS is faster per iteration than the original OAS because it avoids the admissible move calculations.
• Split-Merge Efficacy: Simulations showed that including split-merge moves during the burn-in period allowed the sampler to escape local modes in as few as 10 iterations, whereas standard runs required over 100 iterations and often failed to find the true clustering structure.
• Applicability: The method successfully handled priors like the Geometric Process and Exchangeable Stick-Breaking, where marginal samplers cannot be used.

5. Significance

This paper represents a major step forward in Bayesian nonparametric inference:

• Bridging the Gap: It effectively bridges the gap between the wide applicability of conditional samplers and the superior mixing of marginal samplers.
• Practical Utility: By enabling split-merge moves for a broader class of priors, it provides a robust tool for complex mixture modeling where standard methods fail or converge too slowly.
• Scalability: The removal of the "admissible moves" constraint and the parallelizable nature of the update steps make the algorithm more suitable for larger datasets.
• Flexibility: It allows researchers to use complex, non-standard mixing priors (without tractable predictive distributions) without sacrificing the convergence speed typically associated with such models.

In summary, the authors have transformed the Ordered Allocation Sampler from a specialized conditional method into a highly efficient, general-purpose tool for nonparametric mixture modeling, capable of handling both standard and complex priors with improved convergence and ease of implementation.