Towards Practical Reproduction of Stochastic Concept Bottleneck Models

TMLR Paper9412 Authors

02 Jun 2026 (modified: 03 Jun 2026)Under review for TMLREveryoneRevisionsBibTeXCC BY 4.0
Abstract: Stochastic Concept Bottleneck Models (SCBMs) model dependencies among concept logits with a joint Gaussian distribution, enabling interventions on corrected concepts to propagate to related non-intervened concepts. We reproduce the main SCBM experiments on a synthetic correlated-concept dataset and two natural image datasets, comparing SCBM with Hard CBM, autoregressive CBM, and Concept Embedding Models. Our results broadly validate the original empirical findings: SCBM remains competitive in predictive accuracy, improves concept-probability calibration, and enables more efficient interventions, requiring fewer manual concept corrections to achieve comparable concept and target accuracy. Beyond empirical reproduction, we study the practical cost of reproducing SCBM. We identify implementation bottlenecks in the official codebase and introduce a refactored, GPU-oriented pipeline with optimized data loading, batched model execution, batch-first intervention evaluation, and a vectorized Frank-Wolfe solver for dependency-aware interventions. These changes reduce the practical reproduction cost to approximately 62 wall-clock hours on a single RTX 4090. Our optimized implementation also changes the relative runtime behavior reported in the original paper, indicating that computational-efficiency claims are sensitive to implementation choices. Our study therefore supports SCBM's core methodological contribution while suggesting that its practical value should be framed primarily around dependency-aware intervention rather than raw computational efficiency.
Submission Type: Regular submission (no more than 12 pages of main content)
Assigned Action Editor: ~Atsushi_Nitanda1
Submission Number: 9412
Loading