Ramba

This repository provides an implementation of the Ramba model, a hybrid architecture model for relational data modeling.

Ramba combines Transformer's attention mechanism with Mamba's state space model, using bidirectional Mamba in the column view (col view) of relational data and sparse attention mechanisms in the feature view (feat view) and neighbor view (nbr view), achieving efficient relational data modeling.

Installation

1. Install pixi https://pixi.sh/latest/#installation.

2. Clone and install the repository:

git clone https://github.com/ROOOOOOOOL/Ramba.git
cd Ramba
pixi install
# compile and install the rust sampler
cd rustler
pixi run maturin develop --uv --release

This project modifies the mamba-ssm library (main modifications in third_party/mamba/mamba_ssm/modules/mamba2.py), so mamba-ssm will be automatically compiled and installed from the local path third_party/mamba.

Data Preparation

1. Download the datasets and tasks from Relbench:

cd .. # back to the root of the repository
pixi run python scripts/download_relbench.py

2. Link the cache repository

mkdir ~/scratch
ln -s ~/.cache/relbench ~/scratch/relbench

3. Preprocessing (per database):

cd rustler
pixi run cargo run --release -- pre rel-f1

4. Text embedding (per database):

pixi run python -m RM.embed rel-f1

Steps 3. and 4. should be run for all databases:
rel-amazon, rel-avito, rel-event, rel-f1, rel-hm, rel-stack, rel-trial

Experiments

First setup wandb with pixi run wandb login or pixi run wandb disabled. The following example commands replicate the results for rel-amazon/user-churn from the paper.

• Pretrain on all datasets with rel-amazon held-out

pixi run torchrun --standalone --nproc_per_node=8 scripts/pretrain.py

• Continued pretrain from best checkpoint obtained above on all rel-amazon tasks with user-churn held-out

pixi run torchrun --standalone --nproc_per_node=8 scripts/contd_pretrain.py

• Finetune from best checkpoint obtained above on rel-amazon/user-churn task only

pixi run torchrun --standalone --nproc_per_node=8 scripts/finetune.py

Architecture Features

The main features of Ramba include:

• Hybrid Architecture: Uses sparse attention mechanisms in feature and neighbor views, and bidirectional Mamba in the column view
• Bidirectional Mamba: Captures bidirectional dependencies through the combination of forward and backward Mamba
• State Space Model: Leverages Mamba's efficient state space model to handle long-sequence relational data
• Custom Mamba Implementation: Customized modifications to the mamba-ssm library to better adapt to relational data modeling tasks