A Fast, Robust Elliptical Slice Sampling Method for Truncated Multivariate Normal Distributions

Authors: kaiwen  wu
jacob r gardner
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 07:32:45

Abstract: Elliptical slice sampling, when adapted to linearly truncated multivariate normal distributions, is a rejection-free Markov chain Monte Carlo method. At its core, it requires analytically constructing an ellipse-polytope intersection. The main novelty of this paper is an algorithm that computes this intersection in $\mathcal{O}(m \log m)$ time, where $m$ is the number of linear inequality constraints representing the polytope. We show that an implementation based on this algorithm enhances numerical stability, speeds up running time, and is easy to parallelize for launching multiple Markov chains.

Main paper: 
 introduction: 
 Let x ∼ N (0, I) be a d-dimensional standard normal random variable. This paper is concerned with sampling from the truncated multivariate normal distributionp(x) = 1 Z ϕ(x) x ∈ D, 0 x / ∈ D,where ϕ(x) ∝ exp -1 2 x ⊤ x is the standard normal density, Z = x∈D ϕ(x) dx is a normalization constant, and the domain D = {x ∈ R d : Ax ≤ b} is a polytope defined by m linear inequalities with A ∈ R m×d and b ∈ R m . We assume the polytope domain has a non-empty interior but is not necessarily bounded. The standard normal assumption is without loss of generality, since non-standard normal distributions can be handled by a change of variable, as shown in §A.Truncated normal sampling has numerous applications in machine learning and statistics, with recent ones in skew Gaussian processes (e.g., Benavoli et al., 2021) and preferential Bayesian optimization (Benavoli et al., 2021; Takeno et al., 2023) . In addition, truncated normal sampling is a key building block of sophisticated numerical methods estimating integrals related to truncated normal distributions (Gessner et al., 2020) .Linear elliptical slice sampling (e.g., Murray et al., 2010; Fagan et al., 2016; Gessner et al., 2020 ) is a rejection-free tuning-free Markov chain Monte Carlo method for truncated normal distributions. Each iteration analytically constructs the intersection of an ellipse and the polytope domain, from which the next iterate is sampled. In principle, this method is particularly suitable for high dimensional truncated normal distributions thanks to its rejection free property.However, the devil is in the details. Despite its conceptual simplicity, constructing the ellipsepolytope intersection is easier said than done. We will show that all existing implementations share a worst-case time complexity of O(m 2 ), which scales poorly as the number of constraints increases. Moreover, existing implementations have complex control flows, which makes it hard, if not impossible, to parallelize on GPUs. Indeed, to the best of our knowledge, there is no batch implementation of elliptical slice sampling to this date, which is likely due to the programming complexity. Contributions. We develop a new algorithm computing the ellipse-polytope intersection that has a better time complexity and is easier to implement. The algorithm runs in O(m log m) time and is faster than the existing implementations. Moreover, this algorithm has a simple control flow and is particularly amenable to GPU parallelism. As a result, we are able to parallelize thousands of independent Markov chains easily. Experiments show that our implementation accelerates truncated normal sampling massively in high dimensions.2 Elliptical Slice Sampling for Truncated Normal SamplingIn the t-th iteration of linear elliptical slice sampling (e.g., Fagan et al., 2016; Gessner et al., 2020) , we sample a multivariate normal random variable ν t ∼ N (0, I) and form an ellipseEQUATIONThe next iterate x t+1 is sampled from the ellipse-polytope intersection E ∩ D, i.e., the parts of the ellipse that lie inside the polytope domain. This intersection can be constructed analytically by exploiting the polytope structure, and thus no rejection sampling is needed. See Figure 1 for an illustration and Algorithm 1 for the pseudocode. The arising questions are, of course, how to "analytically construct" the ellipse-polytope intersection E ∩ D and how to do it efficiently.Note that the polytope domain itself is the intersection of m halfspaces:D = m i=1 H i = m i=1 {x ∈ R d : a ⊤ i x ≤ b i },where a i is the i-th row of A. Thus, the ellipse-polytope intersection E ∩ D reduces to constructing each ellipse-halfspace intersection E ∩ H i , which does admit an analytical construction.The intersection of the ellipse and the i-th halfspace H i is an elliptical arc. The end points of the elliptical arc are identified by the intersection angles, i.e., the roots of the trigonometry equationEQUATIONwhich typically indeed has two distinct roots α i and β i in closed-forms (see §B). Without loss of generality, we assume all intersection angles α i and β i are converted into [0, 2π] by, if necessary, Algorithm 1: Elliptical Slice Sampling for Truncated Multivariate Normal Distributions Input:1 Initialize x 0 ∈ D 2 for t = 1, 2, • • • do 3 sample ν t ∼ N (I i = [0, α i ] ∪ [β i , 2π] with α i < β i for i = 1, 2, • • • , m Output: I act = ∩ m i=1 I i 1 sort {α i } m i=1 in ascending order: 0 ≤ α i1 ≤ α i2 ≤ • • • ≤ α im ≤ 2π 2 compute γ k = max{β i1 , β i2 , • • • , β i k } for k = 1, 2, • • • m 3 return [0, α i1 ] ∪ m k=2 [γ k-1 , α i k ] ∪ [γ m , 2π] // define [γ k-1 , αi k ] = ∅ if γ k-1 > αi kadding or subtracting a multiple of 2π. In addition, we assume α i is strictly smaller than β i . A simple observation is that the ellipse-halfspace intersection E ∩ H i is precisely represented by the union of two disjoint intervals:I i = [0, α i ] ∪ [β i , 2π]. Note that the interval I i has two disjoint segments due to periodicity: The point x t is represented by two distinct angles 0 and 2π. Intersecting all I i 's gives the interval representation of the ellipsepolytope intersection:I act = m i=1 I i ,which we call the active intervals. Note that the plural form is used because I act may consist of several disjoint intervals, each of which is an active interval. There is an one-to-one correspondence between angles in the active intervals (except the repetition of θ = 0 and θ = 2π) and points in the ellipse-polytope intersection.Below we discuss existing methods constructing the active intervals.Brute Force. The most straightforward algorithm follows the definition of the active intervals:I (0) act = [0, 2π], I (i) act = I (i-1) act ∩ [0, α i ] ∪ [β i , 2π] .In the end, I (m) act equals to the desired active intervals I act . However, this brute force intersection is tedious to implement and hard to parallelize. Moreover, its worst-case time complexity is as slow as O(m 2 ). See §D for why it is not O(m), contrary to how it may appear.Angle Perturbation. In Figure 1 , θ = 2 5 π and θ = 4 5 π are not active, since they do not contribute to the active intervals. The hard part of computing the active intervals precisely lies in identifying those active intersection angles. Gessner et al. (2020) identify the active angles by angle perturbations: An intersection angle θ is active if and only if exactly one of the perturbations θ i -ϵ and θ i + ϵ corresponds a point within the domain D. The time complexity is O(m 2 ). 
a simple method for the ellipse-polytope intersection: 
 Sorting α i 's in ascending order yieldsEQUATIONNote that β i k is not necessarily monotonic in k. Our algorithm is based on the observation below. Proposition 1. Let α i < β i for all i ∈ [m] and let {α i k } m k=1 be sorted in ascending order as in (3). Then, the active intervalsI act = ∩ m i=1 [0, α i ] ∪ [β i , 2π] have an equivalent representationI act = [0, α i1 ] ∪ m k=2 [γ k-1 , α i k ] ∪ [γ m , 2π],whereγ k = max{β i1 , β i2 , • • • , β i k } is the cumulative max until β i k . We interpret the interval [γ k-1 , α i k ] as an empty set if γ k-1 > α i k .Proposition 1 gives a closed-form expression for the active intervals I act , which yields Algorithm 2. Despite a somewhat lengthy proof, the idea and the final expression are both very simple. The time complexity is O(m log m), faster than the brute force algorithm and likelihood testing. In addition, it is simple to program, amenable to GPU parallelism, and easy to batch, since the sorting and cumulative max operations are well supported in every popular machine learning package nowadays.  
experiments: 
 All simulations are run on a single machine with a RTX 3090 GPU using single precision floating points. We use the linear elliptical slice sampling implementation in BoTorch V0.11.1 as the baseline. Our code will be released publicly.We run elliptical slice sampling with Algorithm 2 on two univariate truncated normal distributions: N (0, 1) truncated by -1 ≤ x ≤ 3 and 15 ≤ x ≤ 16, respectively. We draw 10 5 samples from each distribution by running 2000 independent Markov chains in parallel with 500 iterations of burn-in and a thinning of 10. Hence, the total number of steps is 2000 × 500 + 10 × 10 5 = 2 × 10 6 . The mean and variance estimates are accurate at least to the second digit after the decimal point (see §E).We demonstrate Algorithm 2 accelerates high dimensional truncated normal sampling, especially when the number of inequality constraints m is large. We generate a set of random instances with varying dimensions as follows. First, we generate a d × d random matrix A whose entries are i.i.d. samples from a univariate standard normal distribution. Second, we generate a random vector x 0 drawn from a d-dimensional standard normal distribution, which will used as the initialization the Markov chain. Third, we set b = Ax 0 + u, where u is a random vector drawn uniformly from the hypercube [0, 1] d . By construction, the initialization x 0 lies in the interior of the domain. Note that the number of constraints m = d increases as the number of dimensions increases.In Figure 2 , we draw 1000 samples from the general instances of truncated normal distributions and compare the running time against BoTorch's implementation. BoTorch's implementation runs a single Markov chain for 1000 steps. Our implementation runs either a single Markov chain for 1000 steps or 10 chains in parallel for 100 steps. Both of them use no burn-in and no thinning. No rejection occurs when running our implementation on these high dimensional distributions. With a single Markov chain, our implementation is over 10x faster than BoTorch's implementation in high dimensions, e.g., d ≥ 1000 on CPU and d ≥ 4000 on GPU. This speed-up solely comes from the improved per iteration complexity O(m log m). Furthermore, running 10 Markov chains on GPU in parallel yields an additional 10x speed-up in high dimensions. 
conclusion: 
 We have presented a O(m log m) algorithm computing the active intervals in linear elliptical slice sampling for linearly truncated normal distributions. We hope our algorithm and implementation unlock the full potential of elliptical slice sampling for linearly truncated normal distributions, and enable new applications that are previously bottlenecked by the speed of sampling.We end this paper by mentioning two extensions. First, it is interesting to adapt elliptical slice sampling to handle linear equality constraints, in which case the Markov chain has to run in the null space of the linear equality constraints. Second, it is interesting to support differentiable samples by adapting the idea of Zoltowski et al. (2021) . 
 
 Appendix: The standard normal assumption is without loss of generality, since non-standard normal distributions reduce to the standard one by a change of variable. Let x ∼ N (µ, Σ) and let LL ⊤ = Σ be the Cholesky decomposition. Let u ∼ N (0, I) be a standard normal variable. Truncating x by D = {x ∈ R d : Ax ≤ b} is the same as truncating u by a transformed domainD ′ = {u ∈ R d : ALu ≤ b -Aµ}.Thus, we can sample from the truncated standard normal, truncated by D ′ , and then apply a linear transformation u → Lu + µ.This section solves the trigonometry equationa ⊤ x cos θ + a ⊤ ν sin θ = b.Define p = a ⊤ x, q = a ⊤ ν, and r = p 2 + q 2 . WLOG, we assume r ̸ = 0, otherwise the corresponding inequality constraint is either invalid (b < 0) or a tautology (b ≥ 0). Note that the ratio b/r ≥ -1, otherwise we have b < -r ≤ a ⊤ x. This causes a contradiction since x ∈ D is a feasible point satisfying the linear inequality a ⊤ x ≤ b.Dividing both sides by r gives p r cos θ + q r sin θ = b r .There exists a unique angle τ ∈ [-π, π] (ignoring the repetition at the boundary) such that cos τ = p r and sin τ = q r . In practice, τ is given by arctan2(q, p), a function implemented in many libraries. Applying the angle sum formula givescos(θ -τ ) = b r .It is clear that the ratio b/r determines the number of roots. When -1 < b/r < 1, the two distinct roots are given byEQUATIONA multiple of 2π has to be added to the roots, if necessary, to make sure the angles fall in into [0, 2π].Note that (4) is not the only form of the roots. For instance, Gessner et al. (2020) used the roots of the form θ = ± arccos b r + 2 arctan q r + p .Another root formula, used by Benavoli et al. (2021) , is of the formtan 1 2 θ = q ± √ r 2 -b 2 b + p .Proving these root formulas is left as an exercise for the readers. Despite their equivalence, we recommend using our root formula (4). This is an unbiased opinion, since we arrive at this conclusion after trying all formulas. The other two root formulas need to check additional edge cases when r + p ≈ 0 and b + p ≈ 0, which do happen annoyingly in certain extreme situation in practice.C Proofs Proposition 1. Let α i < β i for all i ∈ [m] and let {α i k } m k=1 be sorted in ascending order as in (3). Then, the active intervalsI act = ∩ m i=1 [0, α i ] ∪ [β i , 2π] have an equivalent representationI act = [0, α i1 ] ∪ m k=2 [γ k-1 , α i k ] ∪ [γ m , 2π],whereγ k = max{β i1 , β i2 , • • • , β i k } is the cumulative max until β i k . We interpret the interval [γ k-1 , α i k ] as an empty set if γ k-1 > α i k .Proof. The sorted angles{α i k } m k=1 divide [0, 2π] into m + 1 disjoint segments: [0, 2π] = [0, α i1 ] ∪ (α i1 , α i2 ] ∪ • • • ∪ (α im-1 , α im ] ∪ (α im , 2π].The active intervals I act are constructed by computing the intersection of I act with each segment. That is, we use the trivial identity Part One. The intersection with the first segment is easy to compute:I act = I act ∩ [0, 2π] = I act ∩ [0, α i1 ] part one ∪ m-1 k=1 I act ∩ (α i k , α i k+1 ] part two ∪ I act ∩ [α im , 2π]I act ∩ [0, α i1 ] = m i=1 I i ∩ [0, α i1 ] = m i=1 I i ∩ [0, α i1 ] = [0, α i1 ],where the first equality uses the definition of the active intervals I act ; the second equality swaps the order of intersections; the third equality uses this observation: [0, a i1 ] is a subset of all I i for i ∈ [m] since α i1 is the smallest angle among all α i 's and β i 's.Part Three. Similarly, the intersection with the last segment is also easy to compute:I act ∩ (a im , 2π] = m i=1 I i ∩ (α im , 2π] = m i=1 [β i , 2π] ∩ (α im , 2π] = max 1≤i≤m β i , 2π ,where the first equality plugs in the definition I act = ∩ m i=1 I i ; the second equality is because α im is the largest angle among all α i 's and therefore we can ignore [0, α i ]; the last equality is due toα im < β im ≤ max 1≤i≤m β i and thus the chunk [α im , max 1≤i≤m β i ) is removed from [a im , 2π].Part Two. Now we deal with the remaining segments(a i k-1 , a i k ] for k = 2, 3, • • • , m.We assume a i k-1 is strictly smaller than a i k for now and defer the case a i k-1 = a i k to the end. For a fixed k, we must computeI act ∩ (a i k-1 , a i k ] = m i=1 I i ∩ (a i k-1 , a i k ] = m i=1 I i ∩ (a i k-1 , a i k ] .Now we split the intersection index i into two cases {i ∈ [m] :α i ≥ α i k } and {i ∈ [m] : α i ≤ α i k-1 }. For the first case, notice that {i∈[m]:αi≥αi k } I i ∩ (a i k-1 , a i k ] = (α i k-1 , α i k ],because the choice of the index i implies(α i k-1 , α i k ] ⊆ [0, α i ] ⊆ I i . For the second case, we have {i∈[m]:αi≤αi k-1 } I i ∩ (a i k-1 , a i k ] = {i∈[m]:αi≤αi k-1 } [0, α i ] ∪ [β i , 2π] ∩ (α i k-1 , α i k ] = {i∈[m]:αi≤αi k-1 } [β i , 2π] ∩ (α i k-1 , α i k ] = max{β i1 , β i2 , • • • , β i k-1 }, 2π ∩ (α i k-1 , α i k ],where the first equality plugs in the definition of I i ; the second equality is because the index i is specifically chosen such that α i ≤ α i k-1 ; the last equality is because α i 's are sorted: the indices such that αi ≤ α i k-1 are precisely i 1 , i 2 , • • • , i k-1 . Define γ k = max{β i1 , β i2 , • • • , β i k-1 }.Combining the two cases, we obtainI act ∩ (a i k-1 , a i k ] = [γ k , α i k ] if γ k ≤ α i k ∅ otherwise for k = 2, 3, • • • , m.Finally, we come back to the edge case a i k-1 = a i k , which implies (a i k-1 , a i k ] = ∅ by convention. Thus, the intersection I act ∩ (a i k-1 , a i k ] is automatically empty. One can verify that the above expression outputs an empty set as well, since γ k ≥ β i k-1 > α i k-1 = α i k . This section shows that the time complexity of Algorithm 3 is at least Ω(m 2 ). To do so, we construct a worst-case input on which Algorithm 3 takes at least Ω(m 2 ) operations.Without loss of generality, we will work with intervals of the form I i = [0, α i ] ∪ [β i , 1]. Consider the following intersection intervalsI i = 0, 1 3 i ∪ 2 1 3 i , 1 , i = 1, 2, • • • , m.By induction, it is easy to show that the intersection of the first k intervals is k i=1I i = 0, 1 3 k ∪ k i=1 2 1 3 i , 1 3 i-1.In particular, ∩ k i=1 I i is an union of k + 1 intervals. Thus, the k-th inner loop of Algorithm 3 takes Ω(k) operations. As a result, the algorithm takes Ω(m 2 ) operations in total. It is also not hard to verify that presorting {α i } m i=1 , or {β i } m i=1 , does not reduce the time complexity. The worse-case inputs for these variants can be constructed similarly.The follow table presents the statistics of the MCMC runs in §4.truncation [-1, 3] [15, 16] true µ 0.2828 15.0661 estimates μ 0.2820 15.0661 true σ 2 0.6161 0.0043 estimates σ2 0.6148 0.0043 



Reviews: review: title: title: The paper proposed a O(n logn) algorithm for computing the analytical active intervals in linear elliptical slice sampling.
main_review: main_review: **Pros:**

1. The manuscript is clearly written, easy to understand, and the proposed method is straightforward to implement.
2. The proofs are well-structured and of high quality.
3. The method effectively improves the complexity of high-dimensional slice sampling, which is a significant contribution for MCMC method.


**Questions and Suggestions:**

1. On line 15, the authors said the polytope is not necessarily bounded, it's a confusing claim. I think the definition of  "polytope" is the bounded (and convex) polyhedron. If the authors only refer to the solution of a set of a finite number of linear constraints, "polyhedron" might be a more acccurate term.

2. Is there a way to improve the algorithm when $m = d$, where $m$ is the number of linear constraints and $d$ is the dimension? It seems that inactive intersections mainly or only occur when $m > d$, i.e., the matrix $A$ is not full-rank.

3. In Section 4.1, the authors evaluate the performance of slice sampling in a one-dimensional case. Given that the focus of the manuscript is on multi-dimensional sampling, it would be more informative to assess the method's performance in two-dimensional cases.

4. In Section 4.2, the polytope with non-empty interior is constructed by generating a random matrix, adding a nonzero uniformly distributed random vector, and setting $m = d$. Does this setup tend to create a polytope with certain properties? For instance, if matrix $A$ is filled with i.i.d. standard normal random variables, it is almost always full rank. Could this influence the structure of the active and inactive intervals? If so, it would be helpful to mention this or include a broader range of polytopes in the experiment. As I mentioned above, checking inactive intersections is the time-consuming part of computing the analytical active intervals, and this seems to occur primarily, or only, when $m > d$ (i.e., when matrix $A$ is not full rank). Therefore, designing an experiment to show how the number of inactive intersections affects the efficiency of standard elliptical slice sampling, and how the proposed method speeds up this process, would better highlight the advantages of the algorithm.

5. I think it would help demonstrate the algorithm’s benefits in real-world scenarios by including some practical applications, such as Bayesian optimization. For instance, in what contexts would sampling from a 1000-dimensional truncated multivariate normal distribution be necessary? Mostly we don't have to sample from a real high dimensional space, therefore giving some real cases would be more convincinng.
title: title: The review identifies key areas for improvement in the paper, including the need for stronger theoretical justification, more comprehensive experimental validation with real-world data, and broader comparisons with alternative methods. The paper could benefit from a deeper analysis of accuracy, convergence, and numerical stability, as well as a clearer explanation of the algorithm’s pseudocode and potential GPU performance limitations.
main_review: main_review: 1. While the paper introduces a faster \(O(m \log m)\) algorithm for the ellipse-polytope intersection, the theoretical justification for this improvement could be strengthened. There is a lack of detailed analysis explaining why this method significantly outperforms existing \(O(m^2)\) algorithms beyond empirical results. 

2. The experiments focus mainly on synthetic datasets with large dimensions, which does validate the speed-up, but they are not very diverse. Testing the algorithm on real-world datasets or applications would provide a better understanding of its practical relevance. The lack of results from domains such as Bayesian optimization or Gaussian processes, which are cited as use cases, leaves a gap in practical validation.

3. The comparison is limited to BoTorch's elliptical slice sampling implementation. While this is useful, it would be more informative to compare the new algorithm against other alternative methods, such as rejection sampling or Hamiltonian Monte Carlo (HMC). This would provide a broader context for understanding the algorithm’s efficiency, especially in scenarios where other methods might also be applied.

4. The paper focuses primarily on speed improvements but provides little discussion on the accuracy and convergence properties of the algorithm. How does the method ensure that the samples drawn are representative of the truncated multivariate normal distribution? 

5. The \(O(m \log m)\) time complexity is theoretically advantageous, but the practical implementation might still face bottlenecks depending on the nature of the constraints. The authors could clarify under which conditions the method performs optimally and when it might face challenges.

6. The authors briefly mention other slice sampling methods in the introduction but do not explore how their method compares in terms of complexity, ease of implementation, and adaptability. 

7. While the paper shows promising results using GPU parallelism, there is no discussion on the limitations of GPU scalability. How well does the method scale when moving from consumer-grade GPUs, like the RTX 3090, to larger setups? 

8. The paper touches on numerical stability in the introduction but does not delve into how the new algorithm ensures stability in practice, especially in high-dimensional cases. Numerical instabilities, particularly when handling many constraints, can cause issues in Monte Carlo methods. 

9. Although the pseudocode in Algorithm 2 is concise, the explanation accompanying it lacks detail in certain areas. For example, the transition between sorting the angles and computing the cumulative max is under-explained.