Technical Summary: Manifold Sampling via Entropy Maximization (MASEM)

Problem Statement

Sampling from constrained distributions is a fundamental challenge in machine learning, with critical applications in Bayesian inference, molecular design, and robotics (e.g., trajectory optimization and behavior cloning). In many practical scenarios, the feasible set $\Sigma$ is defined implicitly by smooth equality and inequality constraints ($h(x)=0, g(x) \leq 0$) and often decomposes into multiple disconnected components.

Existing sampling paradigms, primarily based on Markov Chain Monte Carlo (MCMC) methods like Metropolis–Hastings, Langevin dynamics, or Hamiltonian Monte Carlo, typically rely on local transition kernels. While these methods can ensure feasibility through projection or landing mechanisms, they are fundamentally limited to moving within a single connected region of the feasible set. Consequently, they struggle to traverse "infeasibility barriers" between disconnected components. This leads to a failure in correctly allocating probability mass across components, resulting in undersampling of certain regions and an inability to estimate component volumes or provide unbiased initialization for arbitrary distributions.

Methodology: MASEM

The authors propose Manifold Sampling via Entropy Maximization (MASEM), a framework designed to sample uniformly from a constrained manifold with an unknown number of disconnected components. The core objective is to maximize the entropy of the empirical distribution, which corresponds to a uniform distribution over the feasible set.

Core Mechanism

MASEM operates as a modular resampling scheme that can be combined with existing local constrained samplers (e.g., Non-linear Hit-and-Run or OLLA). The algorithm proceeds iteratively:

1. Initialization: Particles are initialized via random sampling followed by projection onto the feasible set using a Gauss-Newton method on a squared slack function.
2. Local Rejuvenation: A feasibility-preserving kernel $K$ (the local sampler) is applied to ensure particles are approximately uniformly distributed within their current components.
3. Entropy-Based Resampling: This is the key innovation. The algorithm estimates the local density of particles using $k$-nearest neighbor ($k$-NN) distances.
   • Let $\epsilon_{i,k}$ be the distance from particle $x_i$ to its $k$-th nearest neighbor.
   • The local density estimator is $\hat{\rho}(x_i) \propto \epsilon_{i,k}^{-p}$, where $p$ is the intrinsic dimension.
   • To maximize entropy, particles in low-density regions (large $\epsilon_{i,k}$) should be up-weighted. The resampling weights are defined as:
     $$\bar{w}_i = \epsilon_{i,k}^{\tau}$$
     where $\tau$ is a temperature parameter controlling the aggressiveness of the redistribution.
   • Particles are resampled according to normalized weights $w_i = \bar{w}_i / \sum \bar{w}_j$.
4. Iteration: The resampling step is followed by further local rejuvenation to mix particles within components.

The method is designed to be lightweight and minimally invasive, adding only a small runtime overhead while enabling global mass redistribution across disconnected components.

Theoretical Contributions

The paper provides rigorous theoretical guarantees for MASEM in the mean-field regime (as the number of particles $N \to \infty$):

• Convergence Rate: The authors prove that the resampling operator contracts the Kullback-Leibler (KL) divergence between the empirical distribution and the uniform target distribution at a geometric rate of $(1 - \tau/p)$ per iteration.
• Exponential Convergence: Consequently, the KL divergence decreases exponentially with the number of resampling steps.
• Component Coverage: Under mild assumptions (Lipschitz-continuous derivatives, Linear Independent Constraint Qualifier), the initialization step ensures that all connected components are populated with probability approaching 1 as $N \to \infty$.
• Extinction Bounds: The paper analyzes the probability of "component loss" (where a component receives no particles after resampling) and shows this probability decays exponentially with the number of samples, suggesting that even a small number of chains per component (e.g., 4) is sufficient in practice.

Experimental Results

MASEM was evaluated on synthetic benchmarks and robotics applications, comparing it against standard local samplers (NHR, OLLA), soft-constraint annealing (SCMC), and a clustering-based importance sampling baseline (Cluster-NHR).

• Synthetic Benchmarks: On problems with disconnected components (e.g., disconnected disks on a sphere, sine curves with gaps, Swiss rolls), MASEM outperformed all baselines by an order of magnitude in the squared Sinkhorn distance ($W_2^2$). Standard local samplers (NHR, OLLA) failed to allocate mass correctly, often undersampling smaller or isolated components.
• Scalability: MASEM demonstrated robust performance as the ambient dimension and the number of constraints increased. It maintained competitive wall-clock times compared to baselines that evaluate constraint Jacobians, indicating the resampling overhead is negligible.
• Robotics Applications:
  • Motion Planning: MASEM successfully sampled trajectories through narrow gaps and across multiple homotopy classes in obstacle courses, whereas standard NHR often collapsed to local solutions.
  • Grasping: In a 18-dimensional grasp sampling problem, MASEM achieved higher feasible entropy and lower constraint violations compared to alternatives.
• Metrics: The paper introduces the use of KL divergence on pairwise distance histograms as a sensitive metric for detecting mass allocation errors, showing it is more effective than Energy Distance or Sinkhorn distance in distinguishing subtle distributional differences.

Significance and Claims

The authors position MASEM as a significant advancement for constrained sampling in scenarios where the feasible set is disconnected. The primary significance lies in its ability to:

1. Enable Global Mixing: Unlike local MCMC kernels, MASEM explicitly addresses the "infeasibility barrier" problem, allowing particles to redistribute mass across disconnected components without requiring explicit knowledge of the number of components or their geometry.
2. Modularity: It acts as a plug-in enhancement for existing constrained samplers, improving their global exploration capabilities without replacing the underlying local dynamics.
3. Efficiency: It achieves superior sampling quality (lower Sinkhorn distance) with competitive runtime, making it scalable to high-dimensional robotics problems.

The paper concludes that while MASEM introduces new hyperparameters (temperature $\tau$, neighbor count $k$) and a small runtime overhead, the trade-off is justified by the dramatic improvement in sample quality and the ability to handle complex, disconnected feasible sets that previously caused standard methods to fail. The authors note that the current focus is on uniform sampling (entropy maximization) and suggest future work could explore generalizations to non-uniform target distributions.