模組化重構LLaVA，替換元件只需新增1-2個檔案，開源TinyLLaVA Factory來了
机器之心發表於2024-05-27
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
TinyLLaVA 專案由清華大學電子系多媒體訊號與智慧資訊處理實驗室 (MSIIP) 吳及教授團隊和北京航空航天大學人工智慧學院黃雷老師團隊聯袂打造。清華大學 MSIIP 實驗室長期致力於智慧醫療、自然語言處理與知識發現、多模態等研究領域。北航團隊長期致力於深度學習、多模態、計算機視覺等研究領域。近日，清華和北航聯合推出了 TinyLLaVA Factory, 一款支援定製、訓練、評估多模態大模型的程式碼庫，程式碼和模型全部開源。該程式碼庫以軟體工程的工廠模式作為設計理念，模組化地重構了 LLaVA 程式碼庫，注重程式碼的可讀性、功能的擴充套件性、和實驗效果的可復現性。方便研究者和實踐家們更容易地探索多模態大模型的訓練和設計空間。
- Github 專案：https://github.com/TinyLLaVA/TinyLLaVA_Factory
- 論文地址：https://arxiv.org/abs/2405.11788
- Hugging Face 模型地址：https://huggingface.co/tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B or https://huggingface.co/bczhou/TinyLLaVA-3.1B-SigLIP
- 機器之心 SOTA 模型地址：https://sota.jiqizhixin.com/project/tinyllava
LLaVA 作為多模態社群的優質開源專案，備受研究者和開發者的青睞；新入坑多模態大模型的初學者們也習慣以 LLaVA 專案作為起點，學習和訓練多模態大模型。但是 LLaVA 專案的程式碼較為晦澀難懂，一旦不慎更改錯誤，就可能會影響訓練效果，對於新手來說，往往不敢輕易修改其中的細節，給理解和探索多模態大模型的本質細節造成了一定的困難。近日，清華和北航聯合推出 TinyLLaVA Factory，將本來的 LLaVA 程式碼進行模組化重構，專注於簡潔的程式碼實現、新功能的可擴充套件性、以及訓練結果的可復現性，讓你以最小的程式碼量，定製並訓練屬於自己的多模態大模型，同時減少程式碼錯誤率！相同的模型配置、訓練資料和訓練策略條件下，使用 TinyLLaVA Factory 可訓練出比用 LLaVA 程式碼效能略勝一籌的模型。為了讓使用者更容易地讀懂程式碼和使用模型，TinyLLaVA Factory 專案還配備了程式碼文件和 Demo 網站。其總體架構圖如下。
資料預處理部分，TinyLLaVA Factory 摒棄了 LLaVA 程式碼中燒腦的圖片處理和 Prompt 處理過程，提供了標準的、可擴充套件的圖片和文字預處理過程，清晰明瞭。其中，圖片預處理可自定義 Processor，也可使用一些官方視覺編碼器的 Processor，如 CLIP ViT 和 SigCLIP ViT 自帶的 Image Processor。對於文字預處理，定義了基類 Template，提供了基本的、共用的函式，如新增 System Message (Prompt)、Tokenize、和生成標籤 Ground Truth 的函式，使用者可透過繼承基類就可輕鬆擴充套件至不同 LLM 的 Chat Template。
模型部分，TinyLLaVA Factory 很自然地將多模態大模型模組化成 3 個元件 —— 大語言模型元件、視覺編碼器元件、中間的聯結器元件。每個元件由一個工廠物件控制，負責新模型的註冊和替換，使使用者能夠更容易地替換其中任何一個元件，而不會牽連到其他部分。
TinyLLaVA Factory 為每個元件提供了當前主流的模型，如下表所示。
訓練器仍然仿照 LLaVA，採取 Hugging Face 自帶的 Trainer，整合了 Gradient Accumulation，Wandb 做日誌記錄等特性，同樣支援 DeepSpeed ZeRO2/ZeRO3 並行訓練。對於評估部分，TinyLLaVA Factory 提供了 SQA/GQA/TextVQA/VQAv2/POPE/MME/MM-Vet/MMMU 8 個 Benchmark 的評估。接下來，劃重點！TinyLLaVA Factory Github 專案還手把手教你定製自己的多模態大模型。只需簡單地新增 1-2 個檔案，就可以輕鬆替換 LLM 元件、視覺編碼器元件、聯結器元件。拿替換 LLM 模型舉例。據使用過 LLaVA 程式碼庫的同學反應，LLaVA 程式碼想替換非 Llama 系列的語言模型容易出錯。而 TinyLLaVA Factory 可以方便地替換語言模型，只需新增 2 個 py 檔案，一個是 Chat Template 檔案，一個是模型檔案。替換視覺編碼器時，也只需新增 1 個 py 檔案，繼承視覺編碼器的基類即可。
TinyLLaVA Factory 還支援對訓練策略進行定製，對使用者來說只需在配置檔案中進行修改，就能在 pretraining 和 finetuning 階段對 3 個模組元件（LLM / 視覺編碼器 / 聯結器）實現凍住 / 全量微調 / 部分微調 /lora 微調的任意組合。堪稱小白易上手式的教程！
早在今年 2 月，TinyLLaVA 專案就敏銳地捕捉到了 3B 以下 LLM 在多模態大模型中的潛力，利用市面主流的小規模 LLM，訓練了一系列多模態大模型，引數量在 0.89B-3.1B。實驗結果表明經過高質量的資料選擇和更加細緻的訓練策略，利用小規模 LLM 同樣可以實現和大模型相近甚至更加優越的任務表現。（細節詳見技術報告 https://arxiv.org/abs/2402.14289）