歡迎回到 PRX 系列的第四部分。前三部分涵蓋了模型架構、訓練設計和 24 小時快速衝刺,而這次我們將揭開所有這些背後默默支撐的關鍵:資料。在塑造 PRX 品質的所有環節中,資料管線的建構雖然不是最光鮮亮麗的部分,卻是必須做對的重要一環。

我們將分享所做的工作、會如何改進,以及一些我們緩慢學習到的經驗。簡而言之,我們從公共和內部資料集混合組裝訓練資料,使用視覺語言模型 (VLM) 重新標註圖像,並將結果轉化為 PRX 訓練所用的可串流語料庫。

我們的目標是為預訓練組裝一個龐大且多樣化的資料集。在這個階段,模型正在學習世界的樣貌:視覺概念、物體和場景、事物如何構成和打光,以及圖像所能包含的廣泛內容。這是一個關於覆蓋範圍和多樣性的問題,而非單張圖像的完美度。

一個廣泛且具代表性的語料庫能比一個較小、較精美的語料庫,教導模型更多關於視覺世界的結構,即使許多單張圖像只是普通的快照或輕微壓縮。在這個階段過度篩選美學實際上會造成傷害,它會縮小分佈範圍,並讓模型失去日後無法彌補的概念和構成多樣性。讓生成內容看起來精美是另一個後期的考量,我們將其留給在小型、嚴格策劃的資料集上進行微調和偏好對齊。預訓練是為了廣度;微調是為了品味。

我們從公共和內部資料集的混合中組裝預訓練資料。這個階段的優先考量是廣度、多樣性,以及站在現有策劃的基礎上,而不是重新自行策劃。如果一個資料來源已經經過品質過濾、去重,並過濾了不適內容和個人資訊,我們就會利用這些工作,而不是大規模地重複。

資料來源的形式各異:有些直接附帶圖像資料本身,有些則以中繼資料加上基準標註的形式提供,我們將其統一格式。我們採取了務實的方法:與其從頭開始建立一個語料庫,不如利用現有的資料集和我們自己的工具快速組裝一個。事後看來,這不一定是能建立的絕對最佳資料集,但對於預訓練一個 7B 模型來說,它是一個堅實且輕量級的起點。

根據我們的經驗,對於預訓練來說,最重要的是使用能準確描述圖像中所有內容的長標註。我們在第二部分中直接看到了這一點,從短標註切換到長標註顯著改善了樣本品質。如果標註是忠實的,我們就不必擔心圖像中偶爾出現的截圖、廣告、標誌或文字片段,因為這些內容也會在標註中被描述。

這樣一來,模型就會將它們學習為有條件、可控制的屬性,而不是無條件地複製它們。精確的標註將「噪音」轉化為你可以提示生成或提示消除的東西。這正是我們之後進行的過濾刻意輕微的原因。我們只移除真正無法使用的內容,而不是所有不完美的內容。

我們長期以來一直使用 Mosaic Streaming 和 Mosaic Data Shards (MDS) 作為分散式訓練的資料集格式。結合 Mosaic Composer,我們發現它是一個維護成本低、靈活且性能良好的分散式訓練框架。此外,MDS 資料集可以輕鬆有效地混合和打亂,並且還允許直接從 S3 或 GCS 等物件儲存進行分散式訓練。

然而,MDS 資料集非常僵化。添加一個欄位或為給定篩選條件建立子集,基本上意味著必須掃描並重寫整個資料集。這就是為什麼我們使用 Lance 進行這類特徵工程和資料策劃。Lance 是一種欄位式資料格式,具有廉價的述詞下推、純量索引和向量搜尋功能,是建立和探索數十億行資料集的正確工具。

這兩種格式在本文和 PRX 資料管線中相互配合:Lance 用於建構,MDS 用於串流。

對於之前的訓練運行,我們使用 T5Gemma 作為文字編碼器,並預先計算文字潛在向量,將它們以位元組形式儲存在 MDS 中。這次,在將文字編碼器切換到 Qwen3-VL 後,我們決定在訓練期間即時計算文字潛在向量。在訓練迴圈中運行文字編碼器會犧牲吞吐量,但犧牲多少取決於模型。

對於小型去噪器來說,這可能很顯著,但在 PRX 的 7B 規模下,文字編碼器的計算量相對於去噪器來說可以忽略不計,我們測量到的吞吐量成本約為 3-4%(在 30 天的運行中約多出 1 天)。作為回報,我們獲得了兩點好處。跳過預計算使我們的 MDS 分片小得多,小到足以將完整的預訓練資料集儲存在我們 SLURM 叢集的 SSD 支援共享檔案系統上,而不是從物件儲存透過網路串流。

它還讓我們可以自由地在以後更改文字編碼器,而無需重寫數 TB 的儲存潛在向量,這正是我們轉向 Qwen3-VL 時所做的切換。

我們將所有圖像編碼為品質 92 的 JPEG 格式,而不是 PNG 等無損格式。我們不只是假設品質 92 是安全的,我們還進行了測量。真實世界的圖像通常已經經過多次 JPEG 壓縮,因此真正的問題是再進行一次重新編碼是否會造成損害。

在對高解析度(1-2 MP)和較低解析度(0.25-0.5 MP)的真實圖像進行重複解碼/編碼循環後,第一次以品質 92 重新編碼基本上是無法察覺的。每個進一步的循環幾乎沒有增加任何東西,因為 JPEG 很快就會收斂到穩定狀態,即使經過 10 個循環,圖像仍然保持在無法察覺的範圍內,而 PNG 會大 3-10 倍,卻沒有任何感知上的增益。

由於大多數來源圖像已經是 JPEG 壓縮的,以 PNG 無損儲存幾乎沒有好處,因此我們將所有內容轉換為高品質(品質 92)的 JPEG 格式。我們還檢查了最重要的一點:在 JPEG 上訓練是否會改變模型產生的內容。我們的語料庫大部分都是 JPEG,因此 PNG 唯一能提供的優勢是不引入新的偽影。

我們特意從高解析度來源進行比較,預期即使是原始 JPEG 圖像在高解析度下也幾乎沒有偽影。我們在相同的圖像上訓練了兩個相同的 1024px PRX 模型,一個儲存為 PNG,另一個儲存為品質 92 的 JPEG。兩者在我們的指標上訓練得同樣好,它們的生成內容實際上無法區分,包括當我們使用一種已知技術,透過匹配量化表來估計 JPEG 品質時,評估輸出看起來有多像 JPEG 壓縮。

檢測率是指在生成內容中發現任何量化結構的比例。平均和中位數估計 JPEG 品質僅針對那些被標記的圖像計算。這兩個模型實際上無法區分:兩者中大約只有十分之一的生成內容顯示出任何可檢測的量化結構,而且差異小到我們確信訓練圖像格式對輸出品質的影響微乎其微。

因此,高品質 JPEG 儲存並沒有在來源已有的基礎上增加任何可測量的東西。對於大規模的文字轉圖像訓練來說,這已經足夠好了。對於用於訓練其他 Photoroom 模型(例如我們的客製化 AI Shadows 模型)的偽影敏感資料,我們仍然依賴 PNG。

您無法互動式地探索包含數億甚至數十億行的 Parquet 表格。因此,我們將資料儲存在 Lance 中,在那裡我們可以進行索引、查詢和瀏覽。我們使用 Ray Data 進行攝取,並行讀取來源表格,並在叢集上寫入 Lance 表格的許多片段。

一個值得分享的教訓是關於片段化。一個 Lance 表格被分割成片段,而一些操作的擴展性取決於片段的總數,而不是行數:掃描會打開每個片段的檔案,並且一些中繼資料操作是 O(片段數)。因此,一個被分割成太多小片段的表格,無論其大小如何,查詢速度都會很慢。

Lance 的性能指南是將該計數保持在較低水平(經驗法則是即使對於十億行的表格,也應保持在數百個片段的數量級),並定期運行壓縮以將小片段合併為較大的片段。

我們是透過緩慢的方式學到這一點的。我們最初的攝取目標是每個片段只有 100,000 行,這導致了數千個微小片段,甚至使簡單的篩選和全文查詢都變得緩慢。在壓縮到每個片段大約一百萬行之後,對於我們數億行的表格來說,這意味著大約一千個片段,查詢速度變得很快,我們就不再擔心了。

正確的目標取決於行寬度、查詢模式以及資料的寫入方式,事後看來,更少、更大的片段可能是一個更好的預設值。Lance 支援分散式壓縮(例如透過 Lance-Ray 整合),因此事後調整很容易,但最好一開始就不要過度片段化。

我們選擇自行重新標註每張圖像,而不是依賴某些資料集碰巧附帶的標註。原因是為了保持一致性:在混合的資料來源中,標註的長度、風格和品質差異很大,我們希望在整個語料庫中採用單一統一的標準(關於標註器,請參閱第三節)。

然而,資料集已有的中繼資料仍然很有價值,只是用於不同的目的。預先存在的標註,以及有時隨之而來的向量嵌入(例如 CLIP 嵌入),使得資料集可以在 Lance 中立即進行探索:對標註進行全文搜尋,以及對嵌入欄位進行最近鄰搜尋,讓我們可以在運行自己的標註之前,導航資料並快速判斷其品質和所需的過濾。因此,我們兩者都保留,當它們位於單獨的表格中時將它們連接起來,並依賴它們進行我們接下來描述的瀏覽。

透過可查詢的資料,我們可以快速分析其特性。例如,解析度分佈告訴我們在哪裡設定截止點。我們最小的訓練桶是 512² 像素,我們將升級到桶的上限設定為 4/3(約 33%),因此有效截止點是 384² ≈ 147k 像素,加上長寬比在 [0.5, 2.0] 之間。低於此的所有內容我們都會丟棄。

Lance 表格開箱即用支援文字欄位的全文搜尋以及向量嵌入欄位的最近鄰相似度搜尋,兩者都可以透過建立索引來加速。我們對標註和嵌入欄位正是這樣做的,然後建立了一個小型使用者介面來瀏覽資料集:對標註進行即時文字搜尋,以及透過向量搜尋導航視覺上相似圖像的叢集。以這種方式瀏覽資料使我們能夠定性評估其多樣性和品質,並透過觀察及早發現問題。