LAORAM: A Look Ahead ORAM Architecture for Training Large Embedding Tables

Data confidentiality/privacy is becoming a significant concern, especially in the cloud computing era. Memory access patterns have been demonstrated to leak critical information such as security keys and a program's spatial and temporal information. This information leak poses an even more significant privacy challenge in machine learning models with embedding tables. Embedding tables are routinely used to learn categorical features from training data. Even knowing the locations of the embedding table entries accessed, not the data within the embedding table, will compromise categorical input data to the model. Embedding entries are privacy sensitive since they disclose valuable properties about the user. Oblivious RAM (ORAM), and its enhanced variants such as PathORAM have emerged as viable solutions to hide leakage from memory access streams. PathORAM fetches an entire path of memory blocks even if a single block is needed. Once the block is fetched a new path is randomly assigned thereby leading to substantial bandwidth and performance overheads.

In this work, we present LAORAM, an ORAM framework explicitly designed to protect user privacy during embedding table training. LAORAM exploits the unique property of training, namely the training samples that are going to be used in the future are known beforehand. LAORAM preprocesses the training samples (securely without revealing the entry values) to identify the memory blocks which are accessed together in the near future. The system tries to assign these blocks to as few paths as possible within the PathORAM infrastructure.

LAORAM does this operation by combining multiple blocks accessed together as superblocks. Thus, future accesses to a collection of blocks can be satisfied from a few paths, effectively reducing the number of reads and writes required by the framework. To further increase performance, LAORAM uses a fat-tree structure for PathORAM, i.e. a tree with variable bucket size, effectively reducing the number of background evictions required, which improves the stash usage. We have evaluated LAORAM using both a recommendation model (DLRM) and a NLP model (XLM-R) embedding table configurations. LAORAM performs 5 times faster than PathORAM on a recommendation dataset (Kaggle) and 5.4x faster on a NLP dataset (XNLI), while guaranteeing the same security guarantees as the original PathORAM.