Kaitlyn Chan

EECS Department, University of California, Berkeley

Technical Report No. UCB/EECS-2023-90

May 10, 2023

http://www2.eecs.berkeley.edu/Pubs/TechRpts/2023/EECS-2023-90.pdf

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.

Advisors: Sayeef Salahuddin

---

BibTeX citation:

@mastersthesis{Chan:EECS-2023-90,
    Author= {Chan, Kaitlyn},
    Title= {Solving Ising Model Optimization Problems with Restricted Boltzmann Machines on Google Tensor Processing Units (TPUs)},
    School= {EECS Department, University of California, Berkeley},
    Year= {2023},
    Month= {May},
    Url= {http://www2.eecs.berkeley.edu/Pubs/TechRpts/2023/EECS-2023-90.html},
    Number= {UCB/EECS-2023-90},
    Abstract= {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.},
}

EndNote citation:

%0 Thesis
%A Chan, Kaitlyn 
%T Solving Ising Model Optimization Problems with Restricted Boltzmann Machines on Google Tensor Processing Units (TPUs)
%I EECS Department, University of California, Berkeley
%D 2023
%8 May 10
%@ UCB/EECS-2023-90
%U http://www2.eecs.berkeley.edu/Pubs/TechRpts/2023/EECS-2023-90.html
%F Chan:EECS-2023-90