Hao Shao, Shengju Qian, Han Xiao, Guanglu Song, Zhuofan Zong, Letian Wang, Yu Liu, Hongsheng Li

This repository contains code for the paper Visual CoT: Unleashing Chain-of-Thought Reasoning in the Multi-Modal Language Model and it was built based on LLaVA

The work proposes a multi-turn processing pipeline for the multi-modal language model that dynamically focuses on visual inputs and provides interpretable thoughts. We also collect and introduce the Visual CoT dataset comprising 373k question-answer pairs, annotated with intermediate bounding boxes highlighting key regions essential for answering the questions. Importantly, the released benchmark is capable of evaluating MLLMs in scenarios requiring specific local region identification.

• Install
• Model Zoo
• Train
• Evaluation
• Visualization

1. Clone this repository and navigate to Visual-CoT folder

git clone https://github.com/deepcs233/Visual-CoT.git
cd Visual-CoT

2. Install Package

conda create -n viscot python=3.10 -y
conda activate viscot
pip install --upgrade pip  # enable PEP 660 support
pip install -e .

3. Install additional packages for training cases

pip install -e ".[train]"
pip install flash-attn --no-build-isolation

Please refer to https://github.com/haotian-liu/LLaVA/tree/main?tab=readme-ov-file#demo

The model weights below are merged weights. You do not need to apply delta. The usage of VisCoT checkpoints should comply with the base LLM's model license.

Our training steps are largely consistent with LLaVA; for further details, please refer to the LLaVA documentation/issues.

VisCoT training consists of two stages: (1) feature alignment stage: use our 558K subset of the LAION-CC-SBU dataset to connect a frozen pretrained vision encoder to a frozen LLM; (2) visual instruction tuning stage: use 665K dataset including multimodal instruction-following data and academic VQA tasks from LLaVA-1.5, 1.4M dataset with positional annotations from Shikra, and 373K visual CoT dataset from ours, to teach the model to follow multimodal instructions and obtain the visual CoT ability.

VisCoT is trained on 8 A100 GPUs with 80GB memory. To train on fewer GPUs, you can reduce the per_device_train_batch_size and increase the gradient_accumulation_steps accordingly. Always keep the global batch size the same: per_device_train_batch_size x gradient_accumulation_steps x num_gpus.

We use a similar set of hyperparameters as Vicuna in finetuning. Both hyperparameters used in pretraining and finetuning are provided below.

1. Pretraining

2. Finetuning

Our base model Vicuna v1.5, which is an instruction-tuned chatbot, will be downloaded automatically when you run our provided training scripts. No action is needed.

In our work, we directly use the project's weight from LLaVA-1.5. If you do not need to train it by yourself, projector weights can be downloaded here: https://github.com/haotian-liu/LLaVA/blob/main/docs/MODEL_ZOO.md#projector-weights

Please download the 558K subset of the LAION-CC-SBU dataset with BLIP captions used in LLaVA here.

Pretrain takes around 5.5 hours for VisCoT-13B on 8x A100 (80G), due to the increased resolution to 336px. It takes around 3.5 hours for VisCoT-7B.

Training script with DeepSpeed ZeRO-2: pretrain.sh.

• --mm_projector_type mlp2x_gelu: the two-layer MLP vision-language connector.
• --vision_tower openai/clip-vit-large-patch14-336: CLIP ViT-L/14 336px.

1. Prepare data

Please download the annotation of the final mixture our instruction tuning data viscot_mixed_2m.json to ./playground/data. We also provide our 363k visual CoT datasetviscot_363k.json for building your own dataset. Please download the images from constituting datasets, and some of them may need to register/complete the form first.

• COCO: train2017
• GQA: images
• OCR-VQA: download script, we save all files as .jpg
• TextVQA: train_val_images
• VisualGenome: part1, part2
• Visual7W: repo
• Flickr30k: homepage
• DocVQA: homepage
• InfographicsVQA: homepage
• Open images: download script, we only use 0-5 splits
• VSR: images
• DUDE: images
• SROIE: homepage
• CUB: images

After downloading all of them, organize the data as follows in ./playground/data,

├── coco
│   └── train2017
│   └── train2014
├── gqa
│   └── images
├── ocr_vqa
│   └── images
├── textvqa
│   └── train_images
└── vg
│   ├── VG_100K
│   └── VG_100K_2
└── v7w
│   └── images
└── flickr30k
│   └── images
└── cot
│   └── flickr30k
│   └── docvqa
│   └── gqa
│   └── infographicsvqa
│   └── openimages
│   └── textvqa
│   └── vsr
│   └── dude
│   └── sroie
│   └── cub

2. Start training!

We have prepared LLaVA's pretrained projectors in our repo (checkpoints/llava_7b_mm_projector.bin and checkpoints/llava_13b_mm_projector.bin). It is not recommended to use legacy projectors, as they may be trained with a different version of the codebase, and if any option is off, the model will not function/train as we expected.

Visual instruction tuning takes around 60 hours for VisCoT-7b-224 on 8x A100 (80G).

Training script with DeepSpeed ZeRO-3: finetune.sh.

If you are interested in finetuning LLaVA model to your own task/data, please check out Finetune_Custom_Data.md。

Some options to note:

• --mm_projector_type mlp2x_gelu: the two-layer MLP vision-language connector.
• --vision_tower openai/clip-vit-large-patch14-336: CLIP ViT-L/14 336px.
• --ft_vision_tower True: finetune the vision encoder with the same learning rate as the backbone.
• --vision_tower_lr 2e-6: use a specific vision encder learning rate.

1. Single-GPU inference, VisCoT-7b-336 can be changed to other model names saved in the ./checkpoints/

bash scripts/v1_5/eval/cot_benchmark.sh VisCoT-7b-336

2. Obtain the score using ChatGPT-3.5, the API KEY need to be set in llava/eval/eval_cot_score.py

bash scripts/v1_5/eval/cot_score.sh VisCoT-7b-336

3. Stat the overall score

python tools/cot_get_result.py VisCoT-7b-336

4. Stat the detection accuracy of visual CoT bounding boxes (optional)

python tools/cot_detection_get_result.py VisCoT-7b-336

1. Single-GPU inference, VisCoT-7b-336 can be changed to other model names saved in the ./checkpoints/

bash scripts/v1_5/eval/refcoco.sh VisCoT-7b-336

2. Stat the overall accuracy

python tools/refcoco_get_result.py VisCoT-7b-336

Please refer to LLaVA's scripts.

This implementation is based on code from several repositories.

• LLaVA
• Shikra
• Vicuna

If you find our repo, dataset or paper useful, please cite us as

@misc{shao2024visual,
      title={Visual CoT: Unleashing Chain-of-Thought Reasoning in Multi-Modal Language Models}, 
      author={Hao Shao and Shengju Qian and Han Xiao and Guanglu Song and Zhuofan Zong and Letian Wang and Yu Liu and Hongsheng Li},
      year={2024},
      eprint={2403.16999},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

All code within this repository is under Apache License 2.0.