Jiachen Lei1,5 ·
Julius Berner2 ·
Jiongxiao Wang3 ·
Zhongzhu Chen4
Zhongjie Ba1 ·
Kui Ren1 ·
Jun Zhu5,6 ·
Anima Anandkumar2
1 Zhejiang University 2Caltech 3 UW-Madison 4Amazon 5Shengshu Tech 6Tsinghua University
- Upload Pre-trained Weights
pip install -r requirements.txt
- the ImageNet dataset should be downloaded and organized as follows:
imagenet directory- train/
- n01440764
- *.JPEG
- ...
- n01440764
- val/
- n01440764
- *.JPEG
- ...
- n01440764
- train/
- the CIFAR10 dataset should be downloaded and organized following the official way provided by pytorch. (We read CIFAR10 images using torchvision)
# RCM-B: train on 4 8xGPU machines
# run the following command on each machine while modifying the machine rank from 0 to 3
accelerate launch --num_processes 32 --num_machines 4 --machine_rank $1 train.py --config=configs/imagenet_pretrain_rcmb.py --workdir=[path to save logs/ckpts] --config.dataset.data_dir="path to train/"
# RCM-S: rain on 2 8xGPU machines
# run the following command on each machine while modifying the machine rank from 0 to 1
accelerate launch --num_processes 16 --num_machines 2 --machine_rank $1 train.py --config=configs/imagenet_pretrain_rcms.py --workdir=[path to save logs/ckpts] --config.dataset.data_dir="path to train/"# RCM-B: train on 1 8xGPU machines
# run the following command on each machine while modifying the machine rank from 0 to 3
accelerate launch --num_processes 8 train.py --config=configs/cifar10_pretrain.py --workdir=[path to save logs/ckpts] --config.dataset.data_dir="path to cifar10 folder"Noticeably, we add noise to normalized data. Therefore, to align with
# Fine-tune RCM-B at sigma=0.25
accelerate launch --num_processes 8 --mixed_precision no rcm_tune.py \
--config=configs/imagenet_finetune.py \
--task=consistency \
--eta=0.5 \
--lbd=10.0 \
--noise_aug=0.5 \
--path "path to pre-trained .net ckpt" \
--config.train.epochs=150 \
--config.train.layernorm=True \
--config.optimizer.lr_layer_decay=0.65 \
--config.optimizer.name="adamw" \
--config.optimizer.param.lr=2e-4 \
--config.optimizer.param.weight_decay=0.0 \
--config.dataset.batch_size=1024 \
--config.dataset.augmentation_type="autoaug" \
--config.dataset.mixup=0.8 \
--config.dataset.cutmix=1.0 \
--config.dataset.label_smoothing=0.1 \
--config.dataset.mixup_switch_prob=0.5 \
--config.nnet.hidden_dim=4096 \
--config.dataset.data_dir="path to directory that contains train/ and val/ folders" \
--name="name of directory that saves the experiment results, including logs and ckpts"
# Fine-tune RCM-B at sigma=0.5
accelerate launch --num_processes 8 --mixed_precision no rcm_tune.py \
--config=configs/imagenet_finetune.py \
--task=consistency \
--eta=0.5 \
--lbd=20.0 \
--noise_aug=1.0 \
--path "path to pre-trained .net ckpt" \
--config.train.epochs=50 \
--config.train.layernorm=True \
--config.optimizer.lr_layer_decay=0.65 \
--config.optimizer.name="adamw" \
--config.optimizer.param.lr=2e-4 \
--config.optimizer.param.weight_decay=0.0 \
--config.dataset.batch_size=2048 \
--config.dataset.augmentation_type="strong" \
--config.dataset.mixup=0.8 \
--config.dataset.cutmix=1.0 \
--config.dataset.mixup_switch_prob=0.5 \
--config.dataset.label_smoothing=0.0 \
--config.nnet.hidden_dim=4096 \
--config.dataset.data_dir="path to directory that contains train/ and val/ folders" \
--name="name of directory that saves the experiment results, including logs and ckpts"
# Fine-tune RCM-B at sigma=1.0
accelerate launch --num_processes 8 --mixed_precision no rcm_tune.py \
--config=configs/imagenet_finetune.py \
--task=consistency \
--eta=0.5 \
--lbd=20.0 \
--noise_aug=2.0 \
--path "path to pre-trained .net ckpt" \
--config.train.epochs=150 \
--config.train.layernorm=True \
--config.optimizer.lr_layer_decay=0.65 \
--config.optimizer.name="adamw" \
--config.optimizer.param.lr=1e-4 \
--config.optimizer.param.weight_decay=0.0 \
--config.dataset.batch_size=256 \
--config.dataset.augmentation_type="weak" \
--config.dataset.mixup=0.8 \
--config.dataset.cutmix=1.0 \
--config.dataset.mixup_switch_prob=0.5 \
--config.dataset.label_smoothing=0.0 \
--config.nnet.hidden_dim=4096 \
--config.dataset.data_dir="path to directory that contains train/ and val/ folders" \
--name="name of directory that saves the experiment results, including logs and ckpts"# Fine-tune RCM-B at sigma=0.25
accelerate launch --num_processes 8 --mixed_precision no rcm_tune.py \
--config=configs/cifar10_finetune.py \
--task=consistency \
--eta=0.5 \
--lbd=10.0 \
--noise_aug=0.5 \
--path "path to pre-trained .net ckpt" \
--config.train.epochs=100 \
--config.train.layernorm=False \
--config.optimizer.lr_layer_decay=0.75 \
--config.optimizer.name="adamw" \
--config.optimizer.param.lr=2e-4 \
--config.optimizer.param.weight_decay=0.0 \
--config.dataset.batch_size=1024 \
--config.dataset.augmentation_type="weak" \
--config.dataset.mixup=0.0 \
--config.dataset.cutmix=0.0 \
--config.dataset.label_smoothing=0.0 \
--config.dataset.data_dir="path to training images" \
--name="name of directory that saves the experiment results, including logs and ckpts"
# Fine-tune RCM-B at sigma=0.5
accelerate launch --num_processes 8 --mixed_precision no rcm_tune.py \
--config=configs/cifar10_finetune.py \
--task=consistency \
--eta=0.5 \
--lbd=20.0 \
--noise_aug=1.0 \
--path "path to pre-trained .net ckpt" \
--config.train.epochs=100 \
--config.train.layernorm=False \
--config.optimizer.lr_layer_decay=0.75 \
--config.optimizer.name="adamw" \
--config.optimizer.param.lr=2e-4 \
--config.optimizer.param.weight_decay=0.0 \
--config.dataset.batch_size=1024 \
--config.dataset.augmentation_type="weak" \
--config.dataset.mixup=0.0 \
--config.dataset.cutmix=0.0 \
--config.dataset.label_smoothing=0.0 \
--config.dataset.data_dir="path to training images" \
--name="name of directory that saves the experiment results, including logs and ckpts"
# Fine-tune RCM-B at sigma=1.0
accelerate launch --num_processes 8 --mixed_precision no rcm_tune.py \
--config=configs/cifar10_finetune.py \
--task=consistency \
--eta=0.5 \
--lbd=20.0 \
--noise_aug=2.0 \
--path "path to pre-trained .net ckpt" \
--config.train.epochs=100 \
--config.train.layernorm=False \
--config.optimizer.lr_layer_decay=0.65 \
--config.optimizer.name="adamw" \
--config.optimizer.param.lr=2e-4 \
--config.optimizer.param.weight_decay=0.0 \
--config.dataset.batch_size=1024 \
--config.dataset.augmentation_type="weak" \
--config.dataset.mixup=0.0 \
--config.dataset.cutmix=0.0 \
--config.dataset.label_smoothing=0.0 \
--config.dataset.data_dir="path to training images" \
--name="name of directory that saves the experiment results, including logs and ckpts"# At noise level [0.25, 0.5, 1.0],, certify RCM-B, pre-trained on Imagenet for 600k steps, with 100k smoothing noise
python certify.py --config=configs/imagenet_finetune.py --n=100000 --step=600000 --path="PATH TO .net file" --name="name of directory for saving the certification results" --config.train.layernorm=True --config.dataset.data_dir="path to the directory that contains train/ and val/ folders"
# At noise level [0.25, 0.5, 1.0], certify RCM-B, pre-trained on CIFAR10 for 300k steps, with 100k smoothing noise
python certify.py --config=configs/cifar10_finetune.py --n=100000 --step=300000 --path="PATH TO .net file" --name="name of directory for saving the certification results" --config.train.layernorm=False --config.dataset.data_dir="path to the dataset directory"@misc{lei2025robustrepresentationconsistencymodel,
title={Robust Representation Consistency Model via Contrastive Denoising},
author={Jiachen Lei and Julius Berner and Jiongxiao Wang and Zhongzhu Chen and Zhongjia Ba and Kui Ren and Jun Zhu and Anima Anandkumar},
year={2025},
eprint={2501.13094},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2501.13094},
}