Clustered Randomized Smoothing for Stochastic Prediction Functions

Clustered Randomized Smoothing for Stochastic Prediction Functions proposes a framework to address the mode collapse problem in randomized smoothing for stochastic, multi-modal predictors. The paper introduces clustered α-smoothing, which partitions noisy prediction samples using a clustering algorithm, applies α-smoothing locally within each cluster, and combines the results into a mixture distribution. The authors derive a lower bound on the probability that the smoothed prediction lies within a union of compact regions corresponding to distinct modes. They evaluate the framework on two benchmarks: stochastic trajectory prediction on a driving simulator dataset and quadrotor control. In trajectory prediction, the approach achieves a 27% lower Wasserstein distance to the ground-truth distribution compared to α-smoothing. In quadrotor control, the method reduces the collision rate by 81% relative to state-of-the-art randomized smoothing.

The paper begins by noting that stochastic predictors such as variational autoencoders, normalizing flows, and Bayesian neural networks can model rich, multi-modal outcome distributions, but this expressive power creates challenges in ensuring robust predictions in safety-critical domains. Randomized smoothing is a leading technique for improving robustness against adversarial perturbations, but in stochastic multi-modal regression settings, it often fails due to mode collapse, yielding averaged predictions that do not reflect the underlying distribution.

To address this limitation, the authors propose clustered α-smoothing, a framework that: (1) partitions noisy samples using an arbitrary clustering algorithm, (2) applies α-smoothing locally within each cluster, and (3) combines the resulting predictions into a mixture distribution. By interpreting the smoothing distribution as a mixture of α-smoothers, they derive a lower bound on the probability that the smoothed prediction lies within a union of compact regions corresponding to distinct modes.

The main contributions are: (i) introducing a framework for randomized smoothing for multi-modal stochastic predictors via clustering and α-trimming, (ii) proving that the output of the clustered α-smoothed predictor with high probability falls within a small coverage region per cluster for all inputs in a radius, and (iii) demonstrating the effectiveness of the framework in benchmarks from stochastic trajectory prediction in traffic and a multi-modal RL controller for a quadrotor, showing improved performance compared to state-of-the-art.

The paper defines a stochastic predictor as a deterministic function hw with a parameter vector w that is a random variable, encompassing models like BNNs, VAEs, and normalizing flows. Randomized smoothing transforms the predictor to H(x) = Eε[hw(x + ε)], but this approach is fundamentally limited for stochastic neural network predictors because it collapses multi-modal structure by averaging over noisy outputs. The α-smoothing technique from prior work allows trimming outliers, but it still inadvertently collapses modes when the predictor distribution is multi-modal.

Clustered α-smoothing is formally defined as follows: given a partition V = V1,..., VM of the output space Y, for each partition cell Vm, the set of indexes Im identifies which prediction samples belong to that cell. Then H̃N,α,Vm applies α-trimming average to each partition cell individually. The clustered α-smoother H̃N,α,V is constructed by randomly selecting each of the H̃N,α,Vm with probability proportional to how many samples are in each Vm. When V = Y, the definition is equivalent to standard α-smoothing.

The paper proves Theorem 3.1, which provides robustness certification for clustered α-smoothing. For any perturbation δ with norm ∥δ∥2 ≤ r, the probability that the smoothed prediction belongs to the union of coverage regions R̃ is lower bounded by an expression involving binomial probability mass functions and the Neyman-Pearson Lemma. The bound accounts for the probability that samples fall within each partition cell and each coverage region, adjusted for the perturbation radius r and the noise level σ.

For the case L = 1, the bound simplifies to a tractable expression. For the general case L > 1, the optimization problem is non-convex, so the authors show how to lower bound it by a linear program via anchor points with a controllable error, as presented in Proposition 3.2.

The paper also presents Algorithm 1 for constructing the partition V and coverage regions R from samples, using clustering algorithms like DBSCAN, and Algorithm 2 for robustness certification. Proposition 4.1 provides high-confidence probability bounds for sets using the Clopper-Pearson Lemma with a union-bound argument.

In the trajectory prediction experiments, the authors use the TrajFlow prediction model on the L-GAP driving simulator dataset. They conduct a safety analysis comparing their approach to α-smoothing and RS-Reg. The risk rate for their method is 2.5%, compared to 7.5% for RS-Reg and 16.5% for α-smoothing. They also compute the 2-Wasserstein distance between the output distributions, finding that clustered α-smoothing reduces the mean distance by 27% compared to α-smoothing (3.80 vs. 5.18). A parameter analysis shows that certified lower bounds are higher when the robustness radius is smaller and when the trimming parameter α is higher.

In the quadrotor control experiments, the task is navigation with obstacles, which is inherently multi-modal as obstacles admit passage on either side. The nominal policy is bi-modal, but α-smoothing destroys the multi-modality and predicts the trimmed average, which corresponds to navigating directly at obstacles. The clustered α-smoothed policy preserves the modes and only tightens the clusters. In simulation, the nominal policy reaches the goal in 72 of 100 trajectories and crashes in 28, the α-smoothed policy crashes in 48, while the clustered α-smoothed policy only crashes in 9 trajectories.

The paper acknowledges limitations, including the reliance on repeated inference, the dependence on the choice of clustering algorithm, the restriction to convex and disjoint sets for R, the fact that certification is with respect to the smoothed predictor rather than the base regressor, the need to fix r/σ to maintain probability bounds, and the potential conservativeness of lower bounds with the inclusion of new clusters.

In conclusion, the paper addresses a key limitation of randomized smoothing for regression by proposing clustered α-smoothing, which preserves multi-modality by applying α-trimmed smoothing within clusters and combining the results as a mixture. The authors establish probabilistic robustness guarantees and demonstrate significant empirical gains, reducing the Wasserstein distance by 27% in trajectory prediction and the collision rate by 81% in quadrotor control. Future work directions include extending the framework to allow input-dependent smoothing and partitions V(x).

Improvements to AI systems:

1. Multi-modal robust regression for safety-critical control

• The improved AI system can perform certified robust predictions that preserve multiple distinct output modes (e.g., left/right obstacle avoidance, multiple future trajectories) instead of collapsing them into a single averaged prediction.

• This enables autonomous vehicles, drones, or robots to maintain safe, diverse action plans under adversarial input perturbations, reducing collision rates by up to 81% compared to standard randomized smoothing.

2. Certified robustness with distributional guarantees for stochastic predictors

• The system can provide formal probabilistic lower bounds (via Theorem 3.1 and Proposition 3.2) that the smoothed output lies within a union of mode-specific coverage regions for any input perturbation within a given radius.

• This allows AI systems to certify safety in deployment, e.g., guaranteeing with high confidence that a trajectory predictor’s output remains within a safe set of possible futures, even under sensor noise or adversarial attacks.

3. Adaptive clustering-based smoothing for multi-modal regression

• The AI system can automatically partition the output space into meaningful clusters (using DBSCAN or similar) and apply local α-trimming within each cluster, then combine them into a mixture distribution.

• This improves distributional fidelity—reducing the 2-Wasserstein distance to the ground truth by 27% in trajectory prediction—so the AI’s output distribution better matches the true stochasticity of the environment, enabling more realistic uncertainty estimation for downstream planning.

4. Outlier-robust stochastic prediction with mode preservation

• By using clustered α-smoothing, the system can trim outliers within each mode without destroying the mode structure, unlike standard α-smoothing which trims across modes and loses multi-modality.

• This makes AI systems more reliable in noisy, multi-modal environments (e.g., traffic scenes with multiple plausible driver behaviors) by maintaining sharp, distinct predictions while discarding spurious samples.

5. Certified multi-modal reinforcement learning policies

• The improved AI system can train and certify stochastic RL policies that exhibit multiple optimal behaviors (e.g., passing an obstacle on either side) while remaining robust to input perturbations.

• This yields policies that are both safer (lower crash rate) and more flexible (preserving multiple viable strategies) than those produced by standard smoothing, which often degrade to a single, unsafe averaged behavior.

6. Efficient high-confidence robustness certification with union-bound guarantees

• The system can compute high-confidence lower bounds on robustness probabilities using Clopper-Pearson intervals and union bounds (Proposition 4.1), enabling practical certification without exhaustive sampling.

• This allows AI systems to be deployed with a quantifiable safety margin, even when the number of modes or clusters is large, by providing tractable linear-programming-based bounds (with controllable error) for non-convex certification problems.

Modern stochastic predictors can model rich, multi-modal outcome distributions. However, this expressive power comes with challenges in ensuring robust predictions - a critical requirement in safety-critical domains. Randomized smoothing is a leading technique for improving robustness, particularly against adversarial perturbations. Yet, in stochastic multi-modal regression settings, randomized smoothing often fails due to mode collapse, yielding averaged predictions that do not reflect the underlying distribution. To address this limitation, we propose clustered alpha-smoothing, a framework that (1) partitions noisy samples using an arbitrary clustering algorithm, (2) applies alpha-smoothing locally within each cluster, and (3) combines the resulting predictions into a mixture distribution. By interpreting the smoothing distribution as a mixture of alpha-smoothers, we derive a lower bound on the probability that the smoothed prediction lies within a union of compact regions corresponding to distinct modes. We empirically evaluate our framework on two benchmarks, demonstrating substantial improvements over state-of-the-art methods. In stochastic trajectory prediction on a driving simulator dataset, our approach achieves, on average, a 27% lower Wasserstein distance to the ground-truth distribution compared to alpha-smoothing. In quadrotor control, where modes correspond to distinct feasible paths to a target, our method reduces the collision rate by 81% relative to the state-of-the-art randomized smoothing.