[ECCV 2026] PixelREPA: Representation Alignment for Just Image Transformers is not Easier than You Think

Jaeyo Shin^*   ·   Jiwook Kim^*   ·   Hyunjung Shim
KAIST AI  
^*Equal Contribution  

This repository is a PyTorch official code of the paper Representation Alignment for Just Image Transformers is not Easier than You Think.

News & Updates

• [2026/06/18] ✈️ Our paper is accepted to ECCV 2026!
• [2026/03/17] 📝 Our paper is now available! The paper is released here.
• [2026/03/16] 🧑‍💻 Our code is released!

PixelREPA

PixelREPA aligns JiT intermediate features to the feature space of an external semantic encoder using our Masked Transformer Adapter (MTA), instead of an MLP. This alternative design accelerates training speed of JiT. The detailed design is described as a following figure:

1. Environment setting and Dataset Preparation

Code Download

git clone https://github.com/kaist-cvml/PixelREPA.git
cd PixelREPA

Dataset Download

Download ImageNet-10k.

Conda Environment

conda env create -f environment.yaml
conda activate pixelrepa

If you get undefined symbol: iJIT_NotifyEvent when importing torch, simply

pip uninstall torch
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124

Check this issue for more details.

2. Training

We trained on 8 NVIDIA H200 GPUs.

Here is the training script on ImageNet 256x256 for 600 epochs:

1. PixelREPA-B/16

  torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_pixelREPA.py \
    --model JiT-B/16 \
    --mta MTA-B/16 \
    --proj_dropout 0.0 \
    --P_mean -0.8 --P_std 0.8 \
    --img_size 256 --noise_scale 1.0 \
    --batch_size 128 --blr 5e-5 \
    --epochs 600 --warmup_epochs 5 \
    --gen_bsz 128 --num_images 50000 --cfg 3.1 --interval_min 0.1 --interval_max 1.0 \
    --output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
    --data_path ${IMAGENET_PATH} --online_eval --eval_freq 10

2. PixelREPA-L/16

  torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_pixelREPA.py \
    --model JiT-L/16 \
    --mta MTA-L/16 \
    --proj_dropout 0.0 \
    --P_mean -0.8 --P_std 0.8 \
    --img_size 256 --noise_scale 1.0 \
    --batch_size 128 --blr 5e-5 \
    --epochs 600 --warmup_epochs 5 \
    --gen_bsz 128 --num_images 50000 --cfg 2.8 --interval_min 0.1 --interval_max 1.0 \
    --output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
    --data_path ${IMAGENET_PATH} --online_eval --eval_freq 10

3. PixelREPA-H/16

  torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_pixelREPA.py \
    --model JiT-H/16 \
    --mta MTA-H/16 \
    --proj_dropout 0.2 \
    --P_mean -0.8 --P_std 0.8 \
    --img_size 256 --noise_scale 1.0 \
    --batch_size 128 --blr 5e-5 \
    --epochs 600 --warmup_epochs 5 \
    --gen_bsz 128 --num_images 50000 --cfg 2.4 --interval_min 0.1 --interval_max 1.0 \
    --output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
    --data_path ${IMAGENET_PATH} --online_eval --eval_freq 10

3. Inference

Pre-trained models are available here.

1. PixelREPA-B/16

  torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_pixelREPA.py \
    --model JiT-B/16 \
    --mta MTA-B/16 \
    --img_size 256 --noise_scale 1.0 \
    --gen_bsz 128 --num_images 50000 --cfg 3.1 --interval_min 0.1 --interval_max 1.0 \
    --output_dir ${CKPT_DIR} --resume ${CKPT_DIR} \
    --data_path ${IMAGENET_PATH} --evaluate_gen

2. PixelREPA-L/16

  torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_pixelREPA.py \
    --model JiT-L/16 \
    --mta MTA-L/16 \
    --img_size 256 --noise_scale 1.0 \
    --gen_bsz 128 --num_images 50000 --cfg 2.8 --interval_min 0.1 --interval_max 1.0 \
    --output_dir ${CKPT_DIR} --resume ${CKPT_DIR} \
    --data_path ${IMAGENET_PATH} --evaluate_gen

3. PixelREPA-H/16

  torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_pixelREPA.py \
    --model JiT-H/16 \
    --mta MTA-H/16 \
    --img_size 256 --noise_scale 1.0 \
    --gen_bsz 128 --num_images 50000 --cfg 2.4 --interval_min 0.1 --interval_max 1.0 \
    --output_dir ${CKPT_DIR} --resume ${CKPT_DIR} \
    --data_path ${IMAGENET_PATH} --evaluate_gen

Acknowledgement

Our code is strongly based on JiT. We sincerely appreciate to the following works:

• JiT
• REPA

Citation

@misc{shin2026pixelrepa,
      title={Representation Alignment for Just Image Transformers is not Easier than You Think}, 
      author={Jaeyo Shin and Jiwook Kim and Hyunjung Shim},
      year={2026},
      eprint={2603.14366},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2603.14366}, 
}