├── .gitignore ├── CITATION.cff ├── LICENSE ├── README.md ├── assets └── graphical_abstract.png ├── configs ├── callbacks │ ├── checkpoint_every_n_steps.yaml │ ├── checkpoint_monitor.yaml │ └── learning_rate_monitor.yaml ├── config.yaml ├── data │ ├── ag_news.yaml │ ├── lambada.yaml │ ├── lm1b-gpt2.yaml │ ├── lm1b-streaming.yaml │ ├── lm1b.yaml │ ├── openwebtext-split.yaml │ ├── openwebtext-streaming.yaml │ ├── openwebtext.yaml │ ├── ptb.yaml │ ├── scientific_papers_arxiv.yaml │ ├── scientific_papers_pubmed.yaml │ ├── text8-crop.yaml │ ├── text8.yaml │ ├── wikitext103.yaml │ └── wikitext2.yaml ├── lr_scheduler │ ├── constant_warmup.yaml │ └── cosine_decay_warmup.yaml ├── model │ ├── medium.yaml │ ├── small-ar.yaml │ ├── small.yaml │ ├── tiny-ar.yaml │ ├── tiny-dimamba.yaml │ └── tiny.yaml ├── noise │ ├── ar.yaml │ ├── linear.yaml │ ├── loglinear.yaml │ └── polynomial.yaml └── strategy │ ├── ddp.yaml │ └── fsdp.yaml ├── dataloader.py ├── diffusion.py ├── main.py ├── models ├── __init__.py ├── autoregressive.py ├── dimamba.py ├── dit.py └── ema.py ├── noise_schedule.py ├── requirements.yml ├── scripts ├── ar.sh ├── dfm.sh ├── fb.sh ├── mdlm.sh ├── remdm-cap.sh ├── remdm-conf.sh ├── remdm-loop.sh ├── remdm-rescale.sh └── sedd.sh └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | outputs 2 | watch_folder 3 | __pycache__ -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/CITATION.cff -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/README.md -------------------------------------------------------------------------------- /assets/graphical_abstract.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/assets/graphical_abstract.png -------------------------------------------------------------------------------- /configs/callbacks/checkpoint_every_n_steps.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/callbacks/checkpoint_every_n_steps.yaml -------------------------------------------------------------------------------- /configs/callbacks/checkpoint_monitor.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/callbacks/checkpoint_monitor.yaml -------------------------------------------------------------------------------- /configs/callbacks/learning_rate_monitor.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/callbacks/learning_rate_monitor.yaml -------------------------------------------------------------------------------- /configs/config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/config.yaml -------------------------------------------------------------------------------- /configs/data/ag_news.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/ag_news.yaml -------------------------------------------------------------------------------- /configs/data/lambada.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/lambada.yaml -------------------------------------------------------------------------------- /configs/data/lm1b-gpt2.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/lm1b-gpt2.yaml -------------------------------------------------------------------------------- /configs/data/lm1b-streaming.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/lm1b-streaming.yaml -------------------------------------------------------------------------------- /configs/data/lm1b.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/lm1b.yaml -------------------------------------------------------------------------------- /configs/data/openwebtext-split.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/openwebtext-split.yaml -------------------------------------------------------------------------------- /configs/data/openwebtext-streaming.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/openwebtext-streaming.yaml -------------------------------------------------------------------------------- /configs/data/openwebtext.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/openwebtext.yaml -------------------------------------------------------------------------------- /configs/data/ptb.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/ptb.yaml -------------------------------------------------------------------------------- /configs/data/scientific_papers_arxiv.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/scientific_papers_arxiv.yaml -------------------------------------------------------------------------------- /configs/data/scientific_papers_pubmed.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/scientific_papers_pubmed.yaml -------------------------------------------------------------------------------- /configs/data/text8-crop.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/text8-crop.yaml -------------------------------------------------------------------------------- /configs/data/text8.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/text8.yaml -------------------------------------------------------------------------------- /configs/data/wikitext103.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/wikitext103.yaml -------------------------------------------------------------------------------- /configs/data/wikitext2.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/data/wikitext2.yaml -------------------------------------------------------------------------------- /configs/lr_scheduler/constant_warmup.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/lr_scheduler/constant_warmup.yaml -------------------------------------------------------------------------------- /configs/lr_scheduler/cosine_decay_warmup.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/lr_scheduler/cosine_decay_warmup.yaml -------------------------------------------------------------------------------- /configs/model/medium.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/model/medium.yaml -------------------------------------------------------------------------------- /configs/model/small-ar.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/model/small-ar.yaml -------------------------------------------------------------------------------- /configs/model/small.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/model/small.yaml -------------------------------------------------------------------------------- /configs/model/tiny-ar.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/model/tiny-ar.yaml -------------------------------------------------------------------------------- /configs/model/tiny-dimamba.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/model/tiny-dimamba.yaml -------------------------------------------------------------------------------- /configs/model/tiny.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/model/tiny.yaml -------------------------------------------------------------------------------- /configs/noise/ar.yaml: -------------------------------------------------------------------------------- 1 | type: ar 2 | scale: 6.0 -------------------------------------------------------------------------------- /configs/noise/linear.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/noise/linear.yaml -------------------------------------------------------------------------------- /configs/noise/loglinear.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/noise/loglinear.yaml -------------------------------------------------------------------------------- /configs/noise/polynomial.yaml: -------------------------------------------------------------------------------- 1 | type: polynomial 2 | a: -3 3 | b: 5 4 | c: -4 5 | eps: 1e-3 -------------------------------------------------------------------------------- /configs/strategy/ddp.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/strategy/ddp.yaml -------------------------------------------------------------------------------- /configs/strategy/fsdp.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/configs/strategy/fsdp.yaml -------------------------------------------------------------------------------- /dataloader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/dataloader.py -------------------------------------------------------------------------------- /diffusion.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/diffusion.py -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/main.py -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/models/__init__.py -------------------------------------------------------------------------------- /models/autoregressive.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/models/autoregressive.py -------------------------------------------------------------------------------- /models/dimamba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/models/dimamba.py -------------------------------------------------------------------------------- /models/dit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/models/dit.py -------------------------------------------------------------------------------- /models/ema.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/models/ema.py -------------------------------------------------------------------------------- /noise_schedule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/noise_schedule.py -------------------------------------------------------------------------------- /requirements.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/requirements.yml -------------------------------------------------------------------------------- /scripts/ar.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/ar.sh -------------------------------------------------------------------------------- /scripts/dfm.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/dfm.sh -------------------------------------------------------------------------------- /scripts/fb.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/fb.sh -------------------------------------------------------------------------------- /scripts/mdlm.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/mdlm.sh -------------------------------------------------------------------------------- /scripts/remdm-cap.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/remdm-cap.sh -------------------------------------------------------------------------------- /scripts/remdm-conf.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/remdm-conf.sh -------------------------------------------------------------------------------- /scripts/remdm-loop.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/remdm-loop.sh -------------------------------------------------------------------------------- /scripts/remdm-rescale.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/remdm-rescale.sh -------------------------------------------------------------------------------- /scripts/sedd.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/scripts/sedd.sh -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kuleshov-group/remdm/HEAD/utils.py --------------------------------------------------------------------------------