VisualRWKV: 基于RWKV的创新视觉语言模型

VisualRWKV: 视觉与语言的创新融合

在人工智能领域,视觉语言模型(Visual Language Models, VLMs)的发展日新月异。然而,将高效的线性循环神经网络(RNNs)架构融入VLMs的尝试却相对较少。本文将为您详细介绍VisualRWKV,这是一个基于RWKV语言模型的创新视觉语言模型,它成功地将线性RNN模型应用于多模态学习任务。

VisualRWKV的诞生与发展

VisualRWKV是由研究人员Haowen Hou、Peigen Zeng、Fei Ma和Fei Richard Yu共同开发的项目。该项目的目标是将RWKV(Receptance Weighted Key Value)语言模型的强大能力扩展到视觉领域,使其能够处理各种视觉任务。

VisualRWKV Logo

自项目启动以来,VisualRWKV经历了多个版本的迭代:

2024年3月25日: VisualRWKV-5.0正式发布
2024年5月11日: VisualRWKV-6.0代码发布
2024年6月25日: VisualRWKV-6.0检查点权重发布

这些更新体现了研究团队对模型持续优化和改进的努力,使VisualRWKV在性能和功能上不断进步。

VisualRWKV的创新架构

VisualRWKV的架构设计体现了视觉和语言模型的巧妙融合。以下是VisualRWKV的架构图:

VisualRWKV Architecture

从架构图中我们可以看到,VisualRWKV主要由以下几个关键组件构成:

视觉编码器: 使用预训练的视觉模型(如CLIP)来提取图像特征。
投影层: 将视觉特征映射到RWKV模型的输入空间。
RWKV主干网络: 处理融合后的视觉-语言信息。
输出层: 生成最终的多模态表示或任务特定输出。

这种设计允许模型有效地处理视觉和语言信息,实现了真正的多模态学习。

VisualRWKV的训练过程

VisualRWKV的训练分为两个主要阶段:预训练和微调。这种分阶段训练策略有助于模型充分学习视觉-语言之间的关联,并适应特定的下游任务。

预训练阶段

在预训练阶段,VisualRWKV主要学习如何将视觉信息与语言模型关联起来。具体步骤如下:

使用预训练数据集(如LLaVA-Pretrain)来训练投影层。
冻结预训练的视觉编码器和RWKV模型,只更新投影层参数。
训练目标是让投影后的视觉特征能够与RWKV模型的语言表示相匹配。

预训练命令示例:

export CUDA_VISIBLE_DEVICES=0,1,2,3
python train.py --load_model /path/to/rwkv/checkpoint \
    --wandb "" --proj_dir path/to/output/ \
    --data_file /path/to/LLaVA-Pretrain/blip_laion_cc_sbu_558k.json \
    --data_type "json" --vocab_size 65536 \
    --ctx_len 1024 --epoch_steps 1000 --epoch_count 9 --epoch_begin 0 --epoch_save 0 \
    --micro_bsz 16 --accumulate_grad_batches 2 --n_layer 24 --n_embd 2048 --pre_ffn 0 \
    --lr_init 1e-3 --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
    --accelerator gpu --devices 4 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 0 \
    --image_folder /path/to/LLaVA-Pretrain/images/ \
    --vision_tower_name /path/to/openai/clip-vit-large-patch14-336 \
    --freeze_rwkv 24 --detail low --grid_size -1 --image_position first \
    --enable_progress_bar True

微调阶段

微调阶段旨在使模型能够执行特定的视觉-语言任务。主要步骤包括:

使用视觉指令数据(如LLaVA数据集)进行微调。
解冻RWKV模型的部分或全部参数,允许模型进行更全面的学习。
训练目标是让模型能够理解并执行视觉相关的指令。

微调命令示例:

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
python train.py --model_path path/to/pretrained-visualrwkv \
    --wandb "" --proj_dir out/rwkv1b5-v060_mix665k \
    --data_file /path/to/LLaVA-Instruct-150K/shuffled_llava_v1_5_mix665k.json \
    --data_type "json" --vocab_size 65536 \
    --ctx_len 2048 --epoch_steps 1000 --epoch_count 20 --epoch_begin 0 --epoch_save 5 \
    --micro_bsz 8 --accumulate_grad_batches 2 --n_layer 24 --n_embd 2048 --pre_ffn 0 \
    --lr_init 2e-5 --lr_final 2e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 \
    --accelerator gpu --devices 8 --precision bf16 --strategy deepspeed_stage_1 --grad_cp 0 \
    --image_folder /path/to/LLaVA-Instruct-150K/images/ \
    --vision_tower_name /path/to/openai/clip-vit-large-patch14-336 \
    --freeze_rwkv 0 --freeze_proj 0 --detail low --grid_size -1 --image_position middle \
    --enable_progress_bar True