Init commit

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,67 @@
----
-license: mit
----
+---
+license: mit
+---
+<div align="center">
+
+# Categorical Schrödinger Bridge Matching (CSBM)
+
+[Grigoriy Ksenofontov](https://scholar.google.com/citations?user=e0mirzYAAAAJ),
+[Alexander Korotin](https://scholar.google.ru/citations?user=1rIIvjAAAAAJ)
+
+[![arXiv Paper](https://img.shields.io/badge/arXiv-2502.01416-b31b1b)](https://arxiv.org/abs/2502.01416)
+[![OpenReview Paper](https://img.shields.io/badge/OpenReview-PDF-8c1b13)](https://openreview.net/forum?id=RBly0nOr2h)
+[![GitHub](https://img.shields.io/github/stars/gregkseno/csbm?style=social)](https://github.com/gregkseno/csbm)
+![GitHub License](https://img.shields.io/github/license/gregkseno/csbm)
+[![WandB](https://img.shields.io/badge/W%26B-view-FFCC33?logo=wandb)](https://wandb.ai/gregkseno/csbm)
+
+</div>
+
+This repository hosts the official checkpoints for the paper "Categorical Schrödinger Bridge Matching", accepted at ICML 2025.
+
+## 📌 TL;DR
+
+This paper extends the Schrödinger Bridge problem to work with discrete time and spaces.
+
+<!-- ![teaser](./images/teaser.png) -->
+
+## 💾 Checkpoints
+
+### CSBM
+
+| Dataset       | Reference Process | $\alpha$    | $N$                   | Saved Iteration |
+| ------------- | ----------------- | ----------- | --------------------- | --------------- |
+| Colored MNIST | **gaussian**      | 0.01        | 2, 4, 10, 25, 50, 100 | 3               |
+| Colored MNIST | **uniform**       | 0.01, 0.05  | 25                    | 3               |
+| CelebA        | **uniform**       | 0.01, 0.005 | 100                   | 4               |
+| Amazon Review | **uniform**       | 0.01, 0.005 | 100                   | 5               |
+
+> [!NOTE]
+> Each experiment directory includes a `config.yaml` file with the full training configuration.
+
+### Additional Components
+
+1. `vqgan_celeba_f8_1024.ckpt` — **VQ-GAN** pretrained on the CelebA dataset
+2. `tokenizer_amazon.json` — **Tokenizer** trained on the Amazon Reviews dataset
+
+## 🎓 Citation
+
+```bibtex
+@article{ksenofontov2025categorical,
+  title={Categorical {Schr\"odinger} Bridge Matching},
+  author={Ksenofontov, Grigoriy and Korotin, Alexander},
+  journal={arXiv preprint arXiv:2502.01416},
+  year={2025}
+}
+```
+
+## 🙏 Credits
+
+- [Weights & Biases](https://wandb.ai) — experiment-tracking and visualization toolkit;
+- [Hugging Face](https://huggingface.co) — Tokenizers and Accelerate libraries for tokenizer implementation, parallel training, and checkpoint hosting on the Hub;
+- [D3PM](https://github.com/google-research/google-research/tree/master/d3pm) — reference implementation of discrete-diffusion models;
+- [Taming Transformers](https://github.com/CompVis/taming-transformers) — original VQ-GAN codebase;
+- [VQ-Diffusion](https://github.com/microsoft/VQ-Diffusion) — vector-quantized diffusion architecture;
+- [MDLM](https://github.com/kuleshov-group/mdlm) — diffusion architecture for text-generation experiments;
+- [ASBM](https://arxiv.org/abs/2405.14449) — evaluation metrics and baseline models for CelebA face transfer;
+- [Balancing the Style-Content Trade-Off in Sentiment Transfer Using Polarity-Aware Denoising](https://arxiv.org/abs/2312.14708) — processed Amazon Reviews dataset and sentiment-transfer baselines;
+- [Inkscape](https://inkscape.org/) — an excellent open-source editor for vector graphics.

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_10/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d64537aa661f8ee3e8516b525dbc0b961eff87cecba8101296bbff5117ae445e
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_10/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5870abf0ad490bbdfcc568ddc0aa82e5b087d46de7420bcb26589e68207dafd
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_10/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 10
+  num_skip_steps: 10
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: gaussian
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 5
+  inner_iterations: 20000
+  use_mini_batch: true
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_100/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5896a264f6e29453ea4cbba2c625487ad3423ddfda8f93d9a94acaf7f36e145
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_100/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e99e791d3d37ff2b4bfe1f88461526b2216cc1e5dc785917beb8c4e973e3d48d
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_100/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 100
+  num_skip_steps: 1
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: gaussian
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 10
+  inner_iterations: 20000
+  use_mini_batch: true
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_2/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca3e655484abba1645b826534921224c05db3126fb92d352ee11a89085547a3d
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_2/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b8747b4c6f4cea35f13596dc81992b70750f9a0fe9d8c7ee8c080da52bf3451
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_2/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 2
+  num_skip_steps: 50
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: gaussian
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 5
+  inner_iterations: 20000
+  use_mini_batch: true
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_25/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb36e26158e2cba3b5b65ba35b7ebeba0817a568f3d9b8eca9d413fc752e4713
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_25/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f44afb3a90eef04257fc4ea5067386bba3a5713cc667dc1f6e7e305b8d92e540
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_25/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 25
+  num_skip_steps: 4
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: gaussian
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 5
+  inner_iterations: 20000
+  use_mini_batch: true
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_4/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba218a2b25c2e7682efee04227fe1c736d6f24905d7adda7cb254ef6422787ec
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_4/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10d5cb66e1770eb54c2d872c24d584125563facf67af25799bf561edf0c64ac1
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_4/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 4
+  num_skip_steps: 25
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: gaussian
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 5
+  inner_iterations: 20000
+  use_mini_batch: true
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_50/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2390760bcb242299d49d881e2c6c869739ead285f4067634ac9c1da8a9eaa2c4
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_50/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57aad3763421081c13d9435c2dc80ace6aaf68397d3a8fc512f9ad19d85e660b
+size 139416416

checkpoints /images/cmnist/gaussian/dim_32_aplha_0.01_num_timesteps_50/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 50
+  num_skip_steps: 2
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: gaussian
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 10
+  inner_iterations: 20000
+  use_mini_batch: true
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/uniform/dim_32_aplha_0.01/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:125c3702db67b5d8c46969e8772f8ffed3fdbc399808968a317086e2d9fe34e1
+size 139416416

checkpoints /images/cmnist/uniform/dim_32_aplha_0.01/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f743431c8174d039c6f3a48a662725893eee873331a149dc0c788f3138d914d6
+size 139416416

checkpoints /images/cmnist/uniform/dim_32_aplha_0.01/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 25
+  num_skip_steps: 4
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.01
+  type: uniform
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 5
+  inner_iterations: 20000
+  use_mini_batch: false
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /images/cmnist/uniform/dim_32_aplha_0.05/checkpoints/backward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fdbd67065bb6ef3933b370fee02c4d935103d8e32a8d7d358f27fb4ba78946c
+size 139416416

checkpoints /images/cmnist/uniform/dim_32_aplha_0.05/checkpoints/forward_3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f329e98da91a5c305b820d013ec4e816d83ce54bd3c33cf42e8b70698c937a2b
+size 139416416

checkpoints /images/cmnist/uniform/dim_32_aplha_0.05/config.yaml

@@ -0,0 +1,47 @@
+data:
+  type: images
+  dataset: cmnist
+  dim: 32
+  num_categories: 256
+  num_timesteps: 25
+  num_skip_steps: 4
+  coupling_type: independent
+model:
+  in_channels: 3
+  num_channels: 64
+  num_layers: 2
+  ch_mults:
+  - 1
+  - 2
+  - 2
+  - 2
+  attention_resolution: 16
+  num_groups: 32
+  num_att_heads: 4
+  dropout: 0.1
+prior:
+  alpha: 0.05
+  type: uniform
+  eps: 1.0e-06
+train:
+  batch_size: 128
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 5
+  inner_iterations: 20000
+  use_mini_batch: false
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.9999
+  optimizer:
+    lr: 0.0002
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 2

checkpoints /quantized_images/celeba/uniform/dim_128_aplha_0.005/checkpoints/backward_4/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fee703324c0024cfa87965f8ce98ecac4d9e199b19b7f07e56189e96bdae4483
+size 372373768

checkpoints /quantized_images/celeba/uniform/dim_128_aplha_0.005/checkpoints/forward_4/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f8b9bd34a828628b91a1a32ef82c66860e7d00f6cd017ded46af500d15a35a5
+size 372373768

checkpoints /quantized_images/celeba/uniform/dim_128_aplha_0.005/config.yaml

@@ -0,0 +1,45 @@
+data:
+  type: quantized_images
+  dataset: celeba
+  dim: 128
+  latent_dim: 16
+  num_categories: 1024
+  num_timesteps: 100
+  num_skip_steps: 1
+  coupling_type: independent
+  train_test_split: 0.9
+model:
+  hidden_dim: 256
+  num_channels: 4
+  num_layers: 18
+  num_att_heads: 16
+  dropout: 0
+codec:
+  ckpt_path: checkpoints/vqgan_celeba_f8_1024.ckpt
+  config_path: configs/vqgan_celeba_f8_1024.yaml
+prior:
+  alpha: 0.005
+  type: uniform
+  eps: 1e-6
+train:
+  batch_size: 32
+  low_precision: false
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 20
+  inner_iterations: 20000
+  use_mini_batch: false
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.999
+  optimizer:
+    lr: 0.0004
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4
+  num_translations: 1

checkpoints /quantized_images/celeba/uniform/dim_128_aplha_0.01/checkpoints/backward_4/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08630e97e5a15fcdb4b51b2ad17f1a6362b48d9b8bc502cd6e517c483ef53c3d
+size 372373768

checkpoints /quantized_images/celeba/uniform/dim_128_aplha_0.01/checkpoints/forward_4/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:814c99bd6124b3ec34af0b68c395f8c6842c79ec8c01a034314e023c4b754fa0
+size 372373768

checkpoints /quantized_images/celeba/uniform/dim_128_aplha_0.01/config.yaml

@@ -0,0 +1,52 @@
+data:
+  type: quantized_images # choices: "toy", "images", "quantized_images", "texts"
+  dataset: celeba
+  dim: 128
+  latent_dim: 16
+  num_categories: 1024
+  num_timesteps: 100
+  num_skip_steps: 1
+  coupling_type: independent # choices: "independent", "prior"
+  train_test_split: 0.9
+model:
+  hidden_dim: 256
+  num_channels: 4
+  num_layers: 18
+  num_att_heads: 16
+  dropout: 0
+codec:
+  ckpt_path: checkpoints/vqgan_celeba_f8_1024.ckpt
+  config_path: configs/vqgan_celeba_f8_1024.yaml
+prior:
+  alpha: 0.01
+  type: uniform # choices: "uniform", "gaussian", "centroid_gaussian", "von_mises"
+  eps: 1e-6
+train:
+  batch_size: 32
+  low_precision: false
+  gradient_accumulation_steps: 4
+  iterations: 20
+  prior_iterations: 20
+  inner_iterations: 20000
+  use_mini_batch: false
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.999
+  optimizer:
+    lr: 0.0004
+    betas: [0.95, 0.99]
+  #   weight_decay: 4.5e-2
+  # scheduler:
+  #     factor: 0.5
+  #     patience: 1000
+  #     min_lr: 1.0e-6
+  #     threshold: 1.0e-1
+  #     threshold_mode: rel
+  #     warmup_lr: 4.5e-4 # the lr to be touched after warmup
+  #     warmup: 10000
+eval:
+  freq: 1000
+  num_samples: 25
+  num_trajectories: 4 # How many trajecotries
+  num_translations: 1 # How many times sample trajecotry from single point

checkpoints /texts/amazon/uniform/dim_100_aplha_0.005/checkpoints/backward_5/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0bf711f8a55bf39b769f4541affe5d6551975fc68f9628199ae593843b6205e
+size 399311828

checkpoints /texts/amazon/uniform/dim_100_aplha_0.005/checkpoints/forward_5/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:945bd9efd2d46e284aed78db5beaf11504f67602d9dc5f0d9c159d7829753dc1
+size 399311828

checkpoints /texts/amazon/uniform/dim_100_aplha_0.005/config.yaml

@@ -0,0 +1,49 @@
+data:
+  type: texts
+  dataset: amazon
+  dim: 100
+  num_categories: 8192
+  num_timesteps: 100
+  num_skip_steps: 1
+  coupling_type: independent
+  train_test_split: null
+model:
+  config:
+    name: small
+    type: ddit
+    hidden_size: 768
+    cond_dim: 128
+    length: 1024
+    n_blocks: 12
+    n_heads: 12
+    scale_by_sigma: true
+    dropout: 0.1
+    tie_word_embeddings: false
+tokenizer:
+  path: checkpoints/tokenizer_amazon.json
+prior:
+  alpha: 0.005
+  type: uniform
+  eps: 1.0e-20
+train:
+  batch_size: 32
+  low_precision: true
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 20
+  inner_iterations: 20000
+  use_mini_batch: false
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.999
+  optimizer:
+    lr: 0.0004
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 32
+  num_trajectories: 1
+  num_translations: 1

checkpoints /texts/amazon/uniform/dim_100_aplha_0.01/checkpoints/backward_5/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1cac9cc955db1256fa5cf9d80da014c34d30a80dedb256e47dae4de486b12c62
+size 399311828

checkpoints /texts/amazon/uniform/dim_100_aplha_0.01/checkpoints/forward_5/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6820f1d460f3a023f00c1ca89d0dbf6ab8965906e3647ec2d77418661104138c
+size 399311828

checkpoints /texts/amazon/uniform/dim_100_aplha_0.01/config.yaml

@@ -0,0 +1,49 @@
+data:
+  type: texts
+  dataset: amazon
+  dim: 100
+  num_categories: 8192
+  num_timesteps: 100
+  num_skip_steps: 1
+  coupling_type: independent
+  train_test_split: null
+model:
+  config:
+    name: small
+    type: ddit
+    hidden_size: 768
+    cond_dim: 128
+    length: 1024
+    n_blocks: 12
+    n_heads: 12
+    scale_by_sigma: true
+    dropout: 0.1
+    tie_word_embeddings: false
+tokenizer:
+  path: checkpoints/tokenizer_amazon.json
+prior:
+  alpha: 0.01
+  type: uniform
+  eps: 1.0e-20
+train:
+  batch_size: 32
+  low_precision: true
+  gradient_accumulation_steps: 1
+  iterations: 20
+  prior_iterations: 20
+  inner_iterations: 20000
+  use_mini_batch: false
+  ce_loss_coeff: 0.001
+  kl_loss_coeff: 1
+  mse_loss_coeff: 0
+  ema_decay: 0.999
+  optimizer:
+    lr: 0.0004
+    betas:
+    - 0.95
+    - 0.99
+eval:
+  freq: 1000
+  num_samples: 32
+  num_trajectories: 1
+  num_translations: 1

checkpoints /tokenizer_amazon.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52d20f72ce3968e1af9b0f0fbea672c7915db6db5fa1ef639d32ee95c5d53a22
+size 309914

checkpoints /vqgan_celeba_f8_1024.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59e325381f46e6b43a5e4ded5d651928d52bc8910ae8a4114be4c871e827fd56
+size 936814106