PDF

Description

The purpose of this project is to demonstrate the capability of the Restricted Boltzmann Machine (RBM) to solve NP-Hard combinatorial optimization problems at scale. It uses the JAX framework to distribute the model across a mesh of interconnected TPU nodes for sampling. The RBM's parallel sampling scheme motivates large scaling across both the batch dimension and problem size to best take advantage of the TPU's architecture for fast and large matrix multiplies. Various sampling algorithms for inference on RBMs are implemented on the TPU and compared to other hardware accelerators. It can currently run up to a 100,000-node MAXCUT problem distributed across 8 TPUv2 cores, which represents one of the largest implementations of such a problem to date.

Details

Files

Statistics

from
to
Export
Download Full History