Technical Summary: Clustered Randomized Smoothing for Stochastic Prediction Functions

1. Problem Statement

Modern stochastic predictors, such as Variational Autoencoders (VAEs), Normalizing Flows, and Bayesian Neural Networks, are capable of modeling rich, multi-modal outcome distributions. This expressiveness is critical for safety-critical domains like autonomous driving, robotics, and healthcare, where predicting multiple plausible futures (e.g., a vehicle turning left or waiting) is essential.

However, ensuring robustness against adversarial perturbations in these models presents a significant challenge. Randomized smoothing, a leading technique for certifying robustness, typically functions by injecting noise into the input and averaging the resulting predictions. While effective for unimodal classification tasks, this approach suffers from mode collapse in stochastic multi-modal regression settings. By averaging over noisy outputs, standard randomized smoothing collapses distinct modes into a single, unimodal, and often uninformative estimate. This failure is particularly dangerous in safety-critical applications where the "average" prediction may correspond to an unsafe or physically impossible state (e.g., driving directly into an obstacle rather than choosing a valid path around it).

Existing solutions like $\alpha$-smoothing (which trims outliers) mitigate the influence of extreme values but still fail to preserve the underlying multi-modal structure when the distribution is inherently multi-modal.

2. Methodology: Clustered $\alpha$-Smoothing

The authors propose Clustered $\alpha$-Smoothing, a framework designed to preserve multi-modality while providing certified robustness guarantees. The method operates in three primary stages:

1. Clustering: Given an input $x$ and a set of noisy samples $\{x + \epsilon_1, \dots, x + \epsilon_N\}$, the resulting prediction samples $\{h_w(x + \epsilon_i)\}$ are partitioned into $M$ disjoint clusters (modes) $V = \{V_1, \dots, V_M\}$ using an arbitrary clustering algorithm (e.g., DBSCAN).
2. Local $\alpha$-Smoothing: Within each cluster $V_m$, the method applies $\alpha$-trimming. This involves removing the top and bottom $\alpha$ fraction of samples (treating them as outliers) and computing the mean of the remaining samples. This yields a locally smoothed predictor $\tilde{H}_{N,\alpha,V_m}(x)$ for each mode.
3. Mixture Reconstruction: The final smoothed predictor is constructed as a mixture distribution. A categorical random variable $z$ selects a cluster $m$ with probability proportional to the cluster size ($\pi^{(m)} = |I_m|/N$), and the output is the locally smoothed prediction for that cluster.

Theoretical Guarantees

The paper derives a probabilistic lower bound on the probability that the smoothed prediction lies within a union of compact regions $\tilde{R} = \cup R_l$ corresponding to distinct modes.

• Robustness Certification: Using the Neyman-Pearson Lemma and union-bound arguments, the authors prove that for any perturbation $\delta$ with $\|\delta\|_2 \le r$, the probability that the predictor output remains within the certified coverage region is bounded from below.
• Handling Non-Convexity: The optimization problem required to compute these bounds is non-convex. The authors propose a tractable solution using anchor points to approximate the convex envelope of the objective function, converting the problem into a Linear Program (LP) with a controllable error term (residual).
• Algorithmic Implementation: The framework includes an algorithm (Algorithm 1) to construct the partition $V$ and coverage sets $R$ from samples, ensuring that the sets are contained within their respective partitions (using generalized Voronoi partitions) to satisfy theoretical requirements.

3. Key Contributions

The paper outlines three main contributions:

1. Framework Introduction: A novel randomized smoothing framework for multi-modal stochastic predictors that utilizes clustering and $\alpha$-trimming to prevent mode collapse.
2. Theoretical Proof: A proof demonstrating that the output of the clustered $\alpha$-smoothed predictor falls within a small coverage region per cluster with high probability for all inputs within a specific radius.
3. Empirical Validation: Evaluation on two benchmarks demonstrating significant improvements over state-of-the-art methods.

4. Experimental Results

The framework was evaluated on two distinct benchmarks:

A. Stochastic Trajectory Prediction (L-GAP Dataset)

• Task: Predicting whether a vehicle will "stay" (wait) or "go" (turn left) at an intersection.
• Metrics: Risk rate (probability of collision due to incorrect prediction) and Wasserstein distance ($W_2$) to the ground-truth distribution.
• Results:
  • Risk Rate: The proposed method achieved a 2.5% risk rate, significantly outperforming RS-Reg (7.5%) and standard $\alpha$-smoothing (16.5%). The authors attribute this to the method's ability to maintain behavioral uncertainty, preventing the model from confidently predicting a dangerous "go" action when the distribution is ambiguous.
  • Distribution Fidelity: The method reduced the mean 2-Wasserstein distance to the original noisy predictor by 27% compared to standard $\alpha$-smoothing (3.80 vs. 5.18), indicating better preservation of the underlying distribution's structure.

B. Quadrotor Control (Multi-modal RL)

• Task: Navigating a quadrotor through obstacles to a goal, where two distinct feasible paths exist (left or right of an obstacle).
• Results:
  • Collision Rate: The standard $\alpha$-smoothing method collapsed the two modes into a single average path, leading to a collision rate of 48%. In contrast, the clustered $\alpha$-smoothing method preserved the distinct paths, reducing the collision rate to 9% (an 81% reduction relative to the state-of-the-art smoothing).
  • Certification: The method provided certified lower bounds on robustness, though the bounds became more conservative as the perturbation radius increased.

5. Significance and Limitations

Significance:
The paper addresses a fundamental limitation of randomized smoothing in regression: the destruction of multi-modal structures essential for safety. By treating modes individually rather than averaging them globally, the method allows for robust certification without sacrificing the diversity of predictions required for complex decision-making in robotics and autonomous systems. The authors emphasize that their approach is agnostic to the underlying predictor and clustering algorithm, making it applicable to various architectures like BNNs, VAEs, and RL policies.

Limitations:
The authors explicitly identify several limitations:

• Computational Cost: Like all randomized smoothing methods, it relies on repeated inference, which can be prohibitive for large neural networks.
• Partition Constraints: The theoretical guarantees rely on defining the coverage region as a union of disjoint, convex sets. This may be restrictive for complex, non-convex clusters (e.g., "half-moon" shapes).
• Predictor vs. Smoother: The guarantees apply to the smoothed predictor, not the base predictor. If exact parity with the base regressor is required (e.g., in physics simulations), this method may not be suitable.
• Noise-Radius Trade-off: To maintain probability bounds while reducing noise, the robustness radius $r$ must be proportional to the reciprocal of the noise standard deviation $\sigma$.
• Conservatism: In highly multi-modal settings, the lower bounds can become increasingly conservative, potentially requiring larger noise levels or smaller certified radii to guarantee robustness.

The paper concludes by noting that future work could explore input-dependent smoothing distributions and dynamic partitions $V(x)$, building on recent trends in the field.