TASO: Task-Aligned Sparse Optimization for Parameter-Efficient Model Adaptation

Official implementation of the EMNLP 2025 paper: "TASO: Task-Aligned Sparse Optimization for Parameter-Efficient Model Adaptation"

TASO leverages importance information from pretrained model weights to eliminate LoRA redundancy before fine-tuning. It estimates parameter importance on downstream tasks, identifies task-specific core regions, and uses their locations to determine the sparse structure of LoRA modules. With a parameter budget comparable to LoRA rank r=1, TASO consistently outperforms standard LoRA across multiple tasks.

Quick Start

# 1. Clone
git clone https://github.com/123bigmirros/TASO.git
cd TASO

# 2. Create environment
conda env create -f environment.yml
conda activate taso

# 3. Install TASO
pip install -e .

# 4. Prepare data
python scripts/prepare_data.py --dataset arc_challenge

# 5. Train
CUDA_VISIBLE_DEVICES=0 python scripts/train_taso.py \
    --model Qwen/Qwen2.5-3B-Instruct --dataset arc_challenge

Installation

Option A: Conda (Recommended)

conda env create -f environment.yml
conda activate taso
pip install -e .

Option B: pip

pip install -e .

Core dependencies: torch>=2.0, transformers>=4.40, peft>=0.8, datasets>=2.14, accelerate>=0.20

Data Preparation

# Download all datasets (GSM8K, BoolQ, WiC, ARC-Challenge)
python scripts/prepare_data.py

# Or download a specific dataset
python scripts/prepare_data.py --dataset arc_challenge

Training

TASO (Ours)

CUDA_VISIBLE_DEVICES=0 python scripts/train_taso.py \
    --model Qwen/Qwen2.5-3B-Instruct \
    --dataset arc_challenge \
    --rank 1 \
    --density 0.1 \
    --learning-rate 5e-5 \
    --num-train-epochs 10 \
    --batch-size 4

Baselines

# LoRA (r=8)
python scripts/train_lora8.py --model Qwen/Qwen2.5-3B-Instruct --dataset arc_challenge

# LoRA (r=32)
python scripts/train_lora32.py --model Qwen/Qwen2.5-3B-Instruct --dataset arc_challenge

# AdaLoRA
python scripts/train_adalora.py --model Qwen/Qwen2.5-3B-Instruct --dataset arc_challenge

# DoRA
python scripts/train_dora.py --model Qwen/Qwen2.5-3B-Instruct --dataset arc_challenge

# Full Fine-tuning
python scripts/train_full.py --model Qwen/Qwen2.5-3B-Instruct --dataset arc_challenge

Evaluation

python scripts/evaluate.py \
    --model-path saves/arc_challenge_42_Qwen2.5-3B-Instruct \
    --dataset arc_challenge

Algorithm

TASO follows a three-stage process:

1. Importance Computation: Compute gradient-based importance scores |θ × ∂L/∂θ| on a sample of training data
2. Dual-Direction Masking: Globally select the top-k important rows and columns to achieve target density
3. Sparse Training: Train only the selected dimensions with learning rate scaled by √(1/density)

Hyperparameters

Parameter	Default	Description
--rank	1	LoRA rank
--density	0.1	Fraction of parameters to keep (10%)
--learning-rate	5e-5	Base learning rate (auto-scaled by density)
--importance-samples	2000	Samples for importance computation

Project Structure

TASO/
├── src/taso/
│   ├── core/              # Core TASO algorithm
│   │   ├── lora.py        # LowRankLinear (sparse LoRA layer)
│   │   ├── importance.py  # Gradient-based importance scoring
│   │   └── masking.py     # Mask utilities
│   ├── data/              # Data loading and chat templates
│   ├── model/             # Model loading
│   ├── trainer/           # Training pipeline and callbacks
│   └── eval/              # Evaluation
├── scripts/               # Training and evaluation scripts
├── configs/               # YAML configurations
└── environment.yml        # Conda environment

Citation

If you find this work useful, please cite:

@inproceedings{miao-etal-2025-taso,
    title = "{TASO}: Task-Aligned Sparse Optimization for Parameter-Efficient Model Adaptation",
    author = "Miao, Daiye  and
      Liu, Yufang  and
      Wang, Jie  and
      Sun, Changzhi  and
      Zhang, Yunke  and
      Yan, Demei  and
      Dong, Shaokang  and
      Zhang, Qi  and
      Wu, Yuanbin",
    booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2025",
    address = "Suzhou, China",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2025.emnlp-main.1157/",
    doi = "10.18653/v1/2025.emnlp-main.1157",
    pages = "22735--22747",
}

License

Apache License 2.0