CrossCat is a Bayesian approach to exploratory data analysis developed by MIT’s Probabilistic Computing Project (probcomp). It partitions the columns of a data matrix into “views,” each of which further partitions rows into “clusters,” discovering the latent dependency structure behind mixed continuous and categorical data in an unsupervised way. As my graduation thesis, I reimplemented it in the probabilistic programming language GenJAX and compared it against the original Python 2 implementation.
What I did for the thesis
The goal was to demonstrate two things:
- Statistical validity: the GenJAX version shows inference quality and behavior equivalent to the original, and recovers the structure of the true generative model
- Computational performance: on practical scenarios (growing rows/columns), it converges faster than the original and scales better
Implementation
- The mixed-type SBP (Stick-Breaking Process) multi-view CrossCat generative model, written in GenJAX
- Trace-based Gibbs sampling kernels updating row clusters, column views, cluster parameters, SBP sticks, and hyperparameters
- JIT compilation via
jitted_update, vectorizing batched parallel updates over all rows and views - Two Dockerized environments (GenJAX / original Python 2.7) for reproducibility
- A unified pipeline of preprocessing → benchmarking → plotting across three datasets: Synthetic, DHA, and Adult
Results
Average time per iteration (CPU):
| Dataset | Original | GenJAX |
|---|---|---|
| Adult (10,000 rows × 15 cols) | 0.75 s | 0.067 s |
| Synthetic (1,000 × 50) | 0.104 s | 0.032 s |
| Synthetic (1,000 × 10) | 0.030 s | 0.022 s |
| DHA (307 × 64) | 0.061 s | 0.045 s |
The speedup grows with data size, reaching about 11x on Adult. The upfront cost of JIT compilation is amortized over iterations and data scale, which substantially improves row-scalability in particular.
Repository
- Source: https://github.com/shohei81/GenJAX_CrossCat
- Both environments can be started with Docker Compose, and
examples/dha_genjax.pyreproduces the GenJAX DHA demo