Pytorch implementation of TAS-LoRA

This is the pytorch implementation of the paper TAS-LoRA: Transformer Architecture Search with Mixture-of-LoRA Experts

For detailed information, please checkout the project site [website]

Requirements

torch==1.7.1+cu110
torchvision==0.2.1
timm==0.3.2
numpy==1.23.1
protobuf==3.20.0
easydict
PyYAML
pillow
tensorboard

Install dependencies:

pip install torch==1.7.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html && pip install numpy==1.23.1 seaborn torchvision==0.2.1 timm==0.3.2 protobuf==3.20.0

Dataset

We use ImageNet (ILSVRC2012). Set the --data-path argument to your ImageNet root directory.

Getting Started

1. Training

Train the TAS-LoRA supernet with frozen pretrained weights + LoRA experts + Router:

bash run.sh

2. Evolutionary Search

After training, search for optimal sub-network architectures under different parameter budgets:

bash search.sh

3. Evaluation

Evaluate specific sub-network architectures:

bash eval.sh

You can also provide custom candidate architectures via a JSON file:

python -m torch.distributed.launch --nproc_per_node=1 --use_env evaluate_subnets.py \
    ... \
    --candidates ./candidates.json

Project Structure

.
├── supernet_train.py          # Main training script
├── supernet_engine.py         # Training and evaluation engine
├── evolution.py               # Evolutionary architecture search
├── evaluate_subnets.py        # Sub-network evaluation
├── model/
│   ├── supernet_transformer.py   # Vision Transformer supernet
│   ├── lora_utils.py             # Shared LoRA utilities
│   ├── utils.py                  # Model utilities
│   └── module/
│       ├── MoeLORA_layer_for_linear.py  # MoE-LoRA for linear layers
│       ├── MoeLORA_layer_for_qkv.py     # MoE-LoRA for QKV projections
│       ├── router.py                     # Architecture-aware router
│       ├── Linear_super.py              # Supernet linear layer
│       ├── multihead_super.py           # Supernet attention
│       ├── qkv_super.py                 # Supernet QKV projection
│       ├── embedding_super.py           # Supernet patch embedding
│       └── layernorm_super.py           # Supernet layer norm
├── lib/                       # Utility libraries
├── timm/                      # PyTorch Image Models (timm 0.3.2)
├── experiments/               # YAML configuration files
│   ├── supernet/supernet-T.yaml
│   └── subnet/AutoFormer-T.yaml
├── run.sh                     # Training script
├── search.sh                  # Search script
└── eval.sh                    # Evaluation script

Acknowledgement

The codebase of this repository is largely borrowed from AutoFormer.

Citation

@inproceedings{jeon2026tas,
  title={TAS-LoRA: Transformer Architecture Search with Mixture-of-LoRA Experts},
  author={Jeon, Jeimin and Lee, Hyunju and Ham, Bumsub},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  year={2026}
}