Bridging Maximum Likelihood and Optimal Transport for Efficient Inference and Model Selection in Stochastic Block Models

Abstract: We study inference in stochastic block models (SBMs) through the lens of optimal transport (OT). We first establish that maximum likelihood variational inference (MLVI) can be interpreted as a semi-relaxed Gromov-Wasserstein (srGW) projection with entropic regularization. While this formulation yields accurate clustering, the entropic regularization prevents transport plans to be sparse, hindering intrinsic model selection. Consequently, we investigate unregularized srGW estimators, and prove that they consistently recover both the SBM connectivity matrix and latent cluster assignments in the asymptotic regime. However, this asymptotic property does not translate into reliable model selection in finite samples, and calls for additional mechanisms to promote sparsity in the inferred cluster proportions. We empirically show that such a regularized formulation yields estimators that simultaneously recover model parameters and select the number of clusters in a single optimization problem, thereby avoiding costly grid search or heuristic model selection procedures.

• The paper establishes a formal equivalence between maximum likelihood variational inference for SBMs and a regularized semi-relaxed Gromov-Wasserstein divergence, enabling efficient model selection.
• It introduces a sparse srGW estimator that accurately recovers latent cluster assignments and block connectivity with rigorous consistency guarantees.
• Empirical evaluations on synthetic SBMs show superior clustering performance and significant speed-ups compared to exhaustive grid search methods.

Bridging Maximum Likelihood and Optimal Transport for Efficient Inference and Model Selection in Stochastic Block Models

Introduction and Problem Formulation

This work establishes a comprehensive connection between variational inference in stochastic block models (SBMs) and optimal transport—specifically the semi-relaxed Gromov-Wasserstein (srGW) divergence. SBMs provide a fundamental statistical framework for modeling networks with latent group structure, and their inference relies on estimating both node-to-cluster assignments and block-wise connectivity parameters. Traditionally, maximum likelihood estimation and its variational approximations have dominated this field, but challenges persist regarding efficient model selection (determining the number of latent clusters) and scalability.

Optimal Transport (OT), and its generalization to the Gromov-Wasserstein (GW) distance, provides a principled approach for comparing structured data, such as graphs, based on their internal geometric relationships. Recent advances in OT have yielded efficient algorithms and regularization schemes, enabling new perspectives on mixture modeling and clustering. However, the intersection of GW OT with latent variable inference in SBMs remains underexplored.

The core contributions of this paper are:

• A formal equivalence between the maximum likelihood variational inference (MLVI) objective for SBMs and the entropic srGW divergence.
• The introduction and analysis of a regularized (sparse) srGW estimator that enables both accurate inference and automatic model selection in SBMs, thus bypassing exhaustive grid search over cluster numbers.
• Theoretical analysis demonstrating consistency of the srGW-based estimators for both cluster assignments and connectivity matrices, with rigorous asymptotic guarantees.
• Empirical validation on synthetic SBMs, benchmarking against state-of-the-art clustering and model-based competitors across multiple SBMs regimes.

Gromov-Wasserstein OT and Variational Inference in SBMs

The GW distance extends OT to settings where two distributions may reside in different metric spaces, operating by aligning pairwise relationships rather than direct correspondences. For graphs, this is particularly relevant as the adjacency matrix encapsulates pairwise connections, and aligning to a latent block-structured “prototype” under GW yields a soft clustering objective.

The srGW divergence allows relaxation of marginal constraints—making inference scalable when only one marginal (source distribution) is prescribed. Within variational inference for SBMs, the mean-field Evidence Lower Bound (ELBO) can be reparameterized such that its principal term is exactly a GW objective with a negative log-likelihood inner loss, plus additional entropy-based regularization terms.

Crucially, entropic regularization promotes dense solutions, impeding the vanishing of unneeded clusters—hence, hindering intrinsic model selection. The sparse (regularized) srGW estimator adds an explicit sparsity penalty on node-to-cluster assignments, encouraging support only on truly present clusters and thus enabling automatic selection of cluster number.

Figure 1: Left: Sample from a Bernoulli SBM with five communities. Right: Corresponding connectivity matrix $\bTheta$

, showing the archetypal block structure exploited by OT.

Consistency and Theoretical Properties

The paper rigorously proves:

• Consistency of srGW-based estimators for both cluster membership recovery (node assignments) and block connectivity parameter estimation, as $N \to \infty$.
• The objective based on negative log-likelihood or broader Bregman divergences delivers asymptotically correct minimizers, under identifiability assumptions standard for SBMs.
• The srGW loss tightly matches (up to small uniform errors) the SBM log-likelihood, ensuring that minimizing srGW corresponds to maximizing likelihood in the large-$N$ regime.

However, unregularized srGW, like classical variational inference, is unable to deactivate redundant clusters at finite $N$ due to the inherent smoothness of the solution—necessitating the explicit $\ell_{1/2}$ penalty for model selection.

Algorithmic Framework

An efficient block coordinate descent algorithm alternates between updating the node-to-cluster transport matrix and estimating the latent block connectivity matrix. Initialization via spectral clustering on the adjacency matrix, followed by srGW minimization, enables rapid convergence. The regularized objective, incorporating a sparsity-inducing term on cluster proportions, is amenable to efficient solution via majorization-minimization schemes and retains the scalability advantages of OT solvers.

Numerical Experiments

The approach is validated on synthetic graphs generated from SBMs under assortative, hub, and disassortative regimes. Evaluations compare the srGW-based estimators (with both negative log-likelihood and squared loss objectives) against widely used methods such as Louvain, Infomap, Greed (hierarchical ICL maximization), and variational EM.

Partitioning performance is judged by Adjusted Rand Index (ARI) over varying block separation (parameter $\alpha$). srGW with negative log-likelihood demonstrates competitive or superior cluster recovery across all regimes, especially where non-assortative structure undermines specialized community-detection algorithms. Notably, the srGW-NLL estimator retains strong performance as $\alpha$ increases, showing robust label recovery as signal strength grows.

Figure 2: ARI scores across varying $\alpha$ (block separability) for different SBM regimes and algorithms, highlighting srGW-NLL's strong performance and stability

.

A sweep over the sparsity hyperparameter reveals clear plateaus where the estimated number of nonempty clusters matches the ground-truth, validating the efficacy of the regularization approach for model selection.

Figure 3: Estimated number of clusters versus the sparsity regularization strength $\lambda$, showing accurate intrinsic model selection when $\alpha$ is sufficiently large.

Additionally, srGW-based methods exhibit orders-of-magnitude faster running times on CPU than full grid search-based model selection via MLE or variational inference, further accelerated by GPU implementation. This enables large-scale analysis previously infeasible with classical SBM inference.

Implications, Practical and Theoretical

This framework advances latent graph inference in several respects:

• Provides a statistical foundation for OT-based clustering and parameter learning in SBMs, justifying the use of GW objectives for community detection and blockmodel estimation.
• Demonstrates that OT-based methods, augmented with appropriate regularization, naturally unify clustering, parameter estimation, and model selection, circumventing ad-hoc procedures and enabling principled one-shot inference.
• Opens avenues for extension beyond Bernoulli SBMs to count, weighted, or attributed network models, provided the inner loss is a proper Bregman divergence or negative log-likelihood.
• Suggests new directions for theoretical analysis of minimax optimality and robustness of OT estimators—potentially tying together high-dimensional inference, optimal transport, and empirical process theory.

Practically, the algorithm provides a tractable route to large-network analysis with credible statistical guarantees and interpretable outputs.

Conclusion

This work bridges a fundamental gap between classical statistical inference in latent graph models and modern optimal transport methodology. By reformulating the variational inference problem for SBMs as a (regularized) srGW problem, the authors provide a unified approach that is provably consistent, scalable, and enables intrinsic model selection. Future research may extend these results to broader classes of random graph and attributed network models, refine sparsity regularization schemes for optimal model selection, and further elaborate the statistical properties of GW-based estimators.