6.82M 輕量級文生圖模型 AuroraIMG-6M 技術架構與評測解析
AuroraIMG-6M 是一個參數量僅 6.82M 的極輕量文生圖模型,涵蓋文字編碼器、去噪器與 VAE,完全由零開始在 CC12M 資料集上訓練完成。本文詳細解析其架構設計、訓練流程、評測數據與硬體適用限制。
現今主流的文字生成圖片模型通常擁有數億甚至數十億參數,且多半仰賴預訓練的大型文字編碼器(例如 CLIP 或 T5)與預訓練 VAE 來建構生成流程。由研究團隊在 AuroraAI-Research 的 Hugging Face 專案頁面 開源的 AuroraIMG-6M,則採取截然不同的實驗方向。
該專案嘗試驗證在極度受限的參數量下,是否能單純依靠小規模資料集從零訓練完整的文生圖系統。整個模型包含文字編碼器、去噪模型與變分自動編碼器(VAE),總參數量僅約 6.82M,能在一般個人電腦甚至 CPU 環境下完成推論。
輕量化文生圖設計:全模組 6.82M 參數從頭訓練
AuroraIMG-6M 的原生生成解析度設定為 64×64 RGB。這項專案的所有模組均未借用任何既有的預訓練權重,全部以 pixparse 整理的 CC12M 資料集 進行自主訓練。
模型的參數配置相當精簡:
- 文字編碼器(Text Encoder):約 1.32M 參數
- 去噪器(Denoiser):約 4.42M 參數
- 變分自動編碼器(VAE):約 1.08M 參數
這三個模組加總僅 6.82M,相較於 Stable Diffusion 1.5 單是 U-Net 就高達 0.86B 參數,AuroraIMG-6M 的體積縮小了超過百倍。
三大核心組件技術細節
為了在極低參數量下維持合理的生成表現,AuroraIMG-6M 針對各個模組進行了針對性的架構精簡與損失函數設計。
1.08M 卷積 VAE
該模型配備的 VAE 負責將 64×64×3 的輸入影像壓縮成 8 個通道、尺寸為 16×16 的潛在空間特徵(Latent)。
架構上採用小型卷積殘差塊(Residual Blocks)搭配群組正規化(GroupNorm)。在訓練損失部分,團隊捨棄了運算負擔較高的對抗損失(Adversarial Loss)與感知損失(LPIPS),改由 L1 損失、基於梯度的邊緣銳利度損失、Gram 矩陣風格損失,以及係數僅 5e-5 的極小 KL 散度項共同組成。在保留驗證集測試中,該 VAE 的重建 PSNR 達到 31.92 dB,FID 為 8.97。
1.32M 聯合訓練文字編碼器
AuroraIMG-6M 沒有採用龐大的預訓練語言模型,而是建立了一個 4 層結構、寬度為 128 且包含 4 個注意力頭的 Pre-LN Transformer。
文字處理層採用在 CC12M 標題資料上訓練的 Byte-Level BPE 分詞器,詞表大小為 4096,最高支援 64 個詞元(Token),並引入旋轉位置編碼(RoPE)。該編碼器在訓練階段與去噪器共同進行聯合訓練,直接學習影像生成過程中最需要的特徵表示。
4.42M 精簡 U-Net 與流匹配演算法
影像生成核心採用小型 U-Net,運作於 16×16、8×8 與 4×4 的潛在維度。
殘差塊使用 AdaLN-Zero 條件調控機制,時間嵌入(Time Embedding)結合隨機傅立葉特徵(Random Fourier Features)與經過平均池化的文字嵌入。演算法採用潛在流匹配(Latent Flow Matching,即 Rectified Flow)技術預測速度向量,推論時搭配 Euler 採樣器運行 50 步,無分類器引導(CFG)係數建議設定為 6.0。
訓練流程與優化策略
模型訓練分為兩階段進行,資料來源為經過嚴格篩選的 CC12M 圖片與文字配對。原始資料經過濾除 HTML、過短標題(少於 3 字)與低解析度圖檔後,保留約 9,882,283 組訓練配對與 50,000 組驗證資料。
第一階段為 VAE 訓練,使用 AdamW 優化器(學習率 5e-4)以 2048 的批次大小訓練 25,000 步。
第二階段為流匹配訓練,共進行 300,000 步(批次大小 2048,約等同 6.14 億樣本或 62 個 Epoch)。權重矩陣與卷積核採用新興的 Muon 優化器(學習率 0.02),其餘嵌入層、正規化層與偏置參數則使用 AdamW 優化器(學習率 5e-4),並配合 10,000 步預熱與 EMA 0.9998 權重平滑。
評測數據與生成能力表現
在 10,000 張 CC12M 保留驗證集的測試中,AuroraIMG-6M 在原生 64×64 解析度下測得的 FID 為 33.01,CLIP 分數為 24.80。作為對照基準,該驗證集真實圖片的 CLIP 分數為 28.25。
在評估文字對齊能力的 GenEval 基準測試中,模型的整體平均準確率為 0.1387:
- 單一主體(Single Object):50.00%
- 顏色辨識(Colors):26.86%
- 雙主體組合(Two Object):3.79%
- 數量計數(Counting):1.56%
- 相對位置(Position):0.75%
- 顏色屬性綁定(Color Attribution):0.25%
在 GenEval 2 測試中,模型測得 Soft-TIFA GM 為 2.21,Soft-TIFA AM 為 21.84。數據顯示出極小型模型在處理單一主要物件時具備一定穩定性,但在多物件邏輯、空間相對關係與複雜語法約束上,與參數量達到數億級別的大型模型仍有明顯差距。
適用情境與技術限制
AuroraIMG-6M 的主要定位在於極端輕量化生成模型的研究驗證,非常適合用於測試 Muon 優化器、精簡 VAE 與流匹配架構的實驗環境,或在運算資源極其有限的邊緣設備上運行。
然而,受限於模型規模與訓練條件,使用時有幾項限制:
- 解析度限制:原生輸出僅 64×64,人臉細節、手部結構與細小文字無法清晰呈現。
- 語言與長度:僅支援小寫英文輸入,建議提示詞保持簡短並聚焦單一主體。
- 概念組合能力有限:對於同時要求多個物體、指定排列位置或計數的提示詞,生成失敗率較高。
- 資料偏誤:模型承襲了網路抓取圖庫的風格與既有偏差,且未施加額外安全性過濾,不適合用於要求嚴格真實性或精確度的應用場景。
常見問題
AuroraIMG-6M 能夠在一般 CPU 上執行嗎?
可以。因為整個模型參數量僅 6.82M,記憶體消耗極低,透過 Hugging Face 的 Diffusers 程式庫載入後,即使在沒有獨立顯示卡的 CPU 環境也能順暢進行推論。
為什麼 AuroraIMG-6M 選擇 64×64 解析度而非更高解析度?
專案的核心目標是探索完全從零訓練極小型模型的可行性。將解析度限制在 64×64,能大幅降低潛在空間維度與計算需求,使 1.08M 參數的 VAE 與 4.42M 的去噪器能在有限資源內收斂。
模型在複雜提示詞上的生成表現如何?
根據 GenEval 的測試數據,模型在處理單一物體的簡單描述時準確率約為 50%,但在面對計數、多物體關聯與空間位置時準確率顯著下降。建議使用類似資料庫標註風格的簡短描述,例如單純描繪「綠草地上的狗」或「日落的海面」。
結語
AuroraIMG-6M 展示了利用現代優化演算法與流匹配技術,將完整文生圖流程壓縮至 6.82M 參數的可能性。該模型現已在 Apache 2.0 授權下開放原始碼與模型權重,為輕量化生成式模型與演算法研究提供了一個結構精簡且具備重現價值的實驗範例。