Tensor-Train Joint Modeling for Few-Step Masked Diffusion

[Project Page] [Arxiv] [Paper]

Setup

conda env create -f requirements.yaml
conda activate ttd

Reproducing Experiments

This repository includes code and scripts to reproduce the experiments on MDLM(+CPD,TTD) with OpenWebText and QM9 datasets. The scripts directory contains the necessary scripts. Check the train.sh, generate.sh, and time.sh scripts for training, generation, and inference time measurement, respectively.

Before running the scripts, you should change dataset.cache_dir in conf/config.yaml to the path where you wish to store the dataset files.

The pretrained checkpoints are available here.

hf download ssamt/tensor-train --local-dir checkpoints

Further Usage

If you wish to apply this paper's tensor decomposition method to your own work, check the tensor directory, which is entirely independent of the rest of the codebase. It contains the implementation of the tensor decomposition methods as outlined in the paper, along with other general-purpose queries for both CPD and TTD. tests/test_decompositions.py includes unit tests for it, which you can run with python -m unittest.