PRX Phần 4: Chiến lược dữ liệu của chúng tôi

Quay lại danh sách bài viết

Chào mừng bạn quay trở lại! Đây là Phần 4 trong chuỗi bài về PRX. Các phần từ 1 đến 3 đã đề cập đến kiến trúc mô hình, thiết kế huấn luyện và quá trình chạy thử nghiệm tốc độ trong 24 giờ. Lần này, chúng ta sẽ vén màn phần công việc âm thầm làm nền tảng cho tất cả những điều trên: dữ liệu. Trong tất cả các yếu tố định hình chất lượng của PRX, đường ống dữ liệu (data pipeline) là một trong những phần ít hào nhoáng nhất khi xây dựng nhưng lại là mảnh ghép quan trọng cần phải thực hiện đúng. Dưới đây là những gì chúng tôi đã làm, những gì chúng tôi sẽ làm khác đi và một vài bài học mà chúng tôi chỉ rút ra được sau những trải nghiệm không mấy dễ dàng.

Tóm gọn trong một câu: chúng tôi tập hợp dữ liệu huấn luyện từ sự kết hợp giữa các tập dữ liệu công khai và nội bộ, thực hiện chú thích lại (re-caption) các hình ảnh bằng VLM, và biến kết quả đó thành một kho dữ liệu có thể truyền tải (streamable corpus) mà chúng tôi dùng để huấn luyện PRX.

Ở cấp độ tổng quan, đường ống dữ liệu trông như sau:

Trong phần tiếp theo, chúng ta sẽ đi sâu vào chi tiết.

1. Các nguyên tắc chỉ đạo

Một tập dữ liệu đa dạng cho việc tiền huấn luyện (pre-training)

Mục tiêu là tập hợp một tập dữ liệu lớn, đa dạng cho quá trình tiền huấn luyện. Ở giai đoạn này, mô hình đang học cách thế giới trông như thế nào: các khái niệm thị giác, các vật thể và khung cảnh, cách mọi thứ được bố cục và chiếu sáng, cũng như phạm vi rộng lớn những gì hình ảnh có thể chứa đựng. Đây là vấn đề về độ bao phủ và tính đa dạng, không phải về sự hoàn hảo của từng hình ảnh. Một kho dữ liệu rộng lớn, mang tính đại diện sẽ dạy cho mô hình nhiều điều về cấu trúc của thế giới thị giác hơn là một tập dữ liệu nhỏ và đẹp mắt hơn, ngay cả khi nhiều hình ảnh riêng lẻ chỉ là những bức ảnh chụp thông thường hoặc bị nén nhẹ. Việc lọc quá mức vì mục đích thẩm mỹ ở giai đoạn này thực tế sẽ gây hại, làm thu hẹp phân phối và khiến mô hình mất đi các khái niệm cũng như sự đa dạng về bố cục mà sau này không thể khôi phục được. Làm cho các thế hệ hình ảnh trông bóng bẩy là một mối quan tâm riêng biệt, nằm ở giai đoạn sau, mà chúng tôi để dành cho việc tinh chỉnh (fine-tuning) và căn chỉnh sở thích (preference alignment) trên các tập dữ liệu nhỏ được chọn lọc kỹ lưỡng. Tiền huấn luyện là để lấy độ rộng; tinh chỉnh là để lấy gu thẩm mỹ.

Sự kết hợp các nguồn dữ liệu

Chúng tôi tập hợp dữ liệu tiền huấn luyện từ sự kết hợp giữa các tập dữ liệu công khai và nội bộ. Ưu tiên ở giai đoạn này là độ rộng, tính đa dạng và tận dụng các công tác chọn lọc đã có sẵn thay vì tự làm lại từ đầu. Ở những nơi mà nguồn dữ liệu đã được lọc chất lượng, khử trùng lặp và lọc nội dung NSFW (không phù hợp tại nơi làm việc) cũng như thông tin cá nhân, chúng tôi xây dựng dựa trên công việc đó thay vì lặp lại ở quy mô lớn. Các nguồn dữ liệu đến dưới nhiều hình thức khác nhau: một số đi kèm với dữ liệu hình ảnh, số khác là siêu dữ liệu (metadata) cộng với các chú thích cơ bản mà chúng tôi đưa về một dạng chung. Chúng tôi đã áp dụng cách tiếp cận thực dụng: thay vì xây dựng kho dữ liệu hoàn toàn từ đầu, chúng tôi dựa vào các tập dữ liệu hiện có và công cụ của riêng mình để tập hợp nhanh chóng. Nhìn lại thì đây không nhất thiết là tập dữ liệu tốt nhất có thể xây dựng, nhưng nó là một điểm khởi đầu vững chắc và nhẹ nhàng để tiền huấn luyện một mô hình 7B.

Triết lý về chú thích của chúng tôi

Theo kinh nghiệm của chúng tôi, điều quan trọng nhất đối với tiền huấn luyện là sử dụng các chú thích dài mô tả chính xác mọi thứ trong hình ảnh. Chúng tôi đã thấy điều này trực tiếp trong Phần 2, nơi việc chuyển từ chú thích ngắn sang chú thích dài đã cải thiện đáng kể chất lượng mẫu. Nếu việc chú thích là trung thực, chúng ta không cần lo lắng về các ảnh chụp màn hình, quảng cáo, logo hoặc các đoạn văn bản xuất hiện ngẫu nhiên trong ảnh, vì những thứ đó cũng được mô tả trong chú thích, nhờ vậy mô hình học chúng như những thuộc tính có điều kiện và có thể kiểm soát thay vì tái tạo chúng một cách vô điều kiện. Chú thích chính xác biến "nhiễu" thành thứ mà bạn có thể dùng prompt để gọi ra hoặc loại bỏ. Đây chính là lý do tại sao việc lọc mà chúng tôi thực hiện sau đó được cố tình làm nhẹ nhàng. Chúng tôi loại bỏ những gì thực sự không thể sử dụng được, chứ không phải mọi thứ không hoàn hảo.

Các định dạng dữ liệu

Chúng tôi đã sử dụng Mosaic Streaming và Mosaic Data Shards (MDS) làm định dạng tập dữ liệu cho huấn luyện phân tán một thời gian rồi. Kết hợp với Mosaic Composer, chúng tôi thấy đây là một khung làm việc rất ít cần bảo trì, linh hoạt và hiệu quả cho huấn luyện phân tán. Hơn nữa, các tập dữ liệu MDS có thể được trộn và xáo trộn một cách dễ dàng, hiệu quả, đồng thời cho phép huấn luyện phân tán trực tiếp từ lưu trữ đối tượng (object storage) như S3 hoặc GCS.

Tuy nhiên, các tập dữ liệu MDS rất cứng nhắc. Việc thêm một cột hoặc tạo một tập con cho một bộ lọc nhất định về cơ bản có nghĩa là phải quét và viết lại toàn bộ tập dữ liệu. Đó là lý do tại sao chúng tôi sử dụng Lance cho loại kỹ thuật đặc trưng (feature engineering) và chọn lọc tập dữ liệu này. Lance là một định dạng dữ liệu dạng cột với khả năng đẩy vị từ (predicate pushdown) chi phí thấp, chỉ mục vô hướng (scalar indexes) và tìm kiếm vector, là công cụ phù hợp để xây dựng và khám phá các tập dữ liệu với hàng tỷ hàng.

Hai định dạng này bổ trợ cho nhau trong suốt bài viết này và đường ống dữ liệu PRX: Lance để xây dựng, MDS để truyền tải (stream).

Về các text latents

Đối với các lần huấn luyện trước, chúng tôi đã sử dụng T5Gemma làm bộ mã hóa văn bản (text encoder) và tính toán trước các text latents, lưu trữ chúng trong MDS dưới dạng byte. Lần này, sau khi chuyển bộ mã hóa văn bản sang Qwen3-VL, chúng tôi quyết định tính toán text latents ngay trong quá trình huấn luyện (on the fly). Việc chạy bộ mã hóa văn bản bên trong vòng lặp huấn luyện sẽ tốn băng thông, nhưng mức độ phụ thuộc vào mô hình: đối với một bộ khử nhiễu (denoiser) nhỏ thì có thể đáng kể, trong khi ở quy mô 7B của PRX, khối lượng tính toán của bộ mã hóa văn bản là không đáng kể so với bộ khử nhiễu, và chúng tôi chỉ đo được mức tiêu tốn băng thông khoảng 3–4% (khoảng 1 ngày thêm vào trong một đợt chạy 30 ngày). Đổi lại, chúng tôi nhận được hai thứ. Việc bỏ qua tính toán trước giúp các shard MDS của chúng tôi nhỏ hơn nhiều, đủ nhỏ để lưu trữ toàn bộ tập dữ liệu tiền huấn luyện trên hệ thống tệp chia sẻ hỗ trợ SSD của cụm SLURM thay vì truyền tải qua mạng từ lưu trữ đối tượng. Và nó giúp chúng tôi tự do thay đổi bộ mã hóa văn bản sau này mà không cần viết lại hàng terabyte latents đã lưu trữ, đây chính xác là sự thay đổi chúng tôi đã thực hiện khi chuyển sang Qwen3-VL.

Về mã hóa hình ảnh

Chúng tôi đã mã hóa tất cả hình ảnh dưới dạng JPEG với chất lượng 92, thay vì định dạng không mất dữ liệu như PNG. Chúng tôi không chỉ giả định chất lượng 92 là an toàn, chúng tôi đã đo lường nó. Hình ảnh trong thế giới thực thường đã được nén JPEG vài lần, vì vậy câu hỏi thực sự là liệu việc mã hóa lại thêm một lần nữa có gây hại hay không. Qua các chu kỳ giải mã/mã hóa lặp đi lặp lại trên cả hình ảnh thực có độ phân giải cao (1–2 MP) và độ phân giải thấp hơn (0.25–0.5 MP), lần mã hóa lại đầu tiên ở chất lượng 92 về cơ bản là không thể nhận thấy. Mỗi chu kỳ tiếp theo hầu như không thêm gì cả, vì JPEG nhanh chóng hội tụ về trạng thái ổn định, và ngay cả sau 10 chu kỳ, hình ảnh vẫn nằm trong phạm vi không thể nhận thấy, trong khi PNG sẽ lớn hơn gấp 3–10 lần mà không mang lại lợi ích gì về mặt cảm nhận. Các con số, trung bình trên 100 hình ảnh ở chất lượng 92 và đo lường so với bản gốc (PSNR cao hơn và LPIPS thấp hơn là tốt hơn):

Vì hầu hết các hình ảnh nguồn đã được nén JPEG, việc lưu trữ chúng không mất dữ liệu dưới dạng PNG không mang lại nhiều lợi ích, vì vậy chúng tôi chuyển đổi mọi thứ sang JPEG ở chất lượng cao (chất lượng 92).

Chúng tôi cũng đã kiểm tra điều quan trọng nhất: liệu việc huấn luyện trên các tệp JPEG có làm thay đổi những gì mô hình tạo ra hay không. Kho dữ liệu của chúng tôi dù sao cũng chủ yếu là JPEG, vì vậy lợi thế duy nhất mà PNG có thể mang lại là không tạo ra các hiện tượng giả (artifacts) mới. Chúng tôi cố tình thực hiện so sánh từ các nguồn có độ phân giải cao, với kỳ vọng rằng ngay cả những hình ảnh gốc là JPEG cũng mang rất ít hiện tượng giả ở độ phân giải cao. Chúng tôi đã huấn luyện hai mô hình PRX giống hệt nhau ở 1024px trên cùng một hình ảnh, một được lưu dưới dạng PNG và một dưới dạng JPEG ở chất lượng 92. Cả hai đều được huấn luyện tốt như nhau dựa trên các chỉ số của chúng tôi, và các thế hệ hình ảnh tạo ra gần như không thể phân biệt được, bao gồm cả khi chúng tôi ước tính hình ảnh đầu ra trông bị nén JPEG như thế nào bằng cách sử dụng một kỹ thuật đã biết để ước tính chất lượng JPEG bằng cách khớp các bảng lượng tử hóa (quantization tables):

Tỷ lệ phát hiện là tỷ lệ các thế hệ hình ảnh mà trong đó tìm thấy bất kỳ cấu trúc lượng tử hóa nào. Chất lượng JPEG ước tính trung bình và trung vị chỉ được lấy trên những hình ảnh được gắn cờ đó. Hai mô hình gần như không thể phân biệt được: chỉ khoảng một trong mười thế hệ hình ảnh từ mỗi mô hình cho thấy bất kỳ cấu trúc lượng tử hóa nào có thể phát hiện được, và sự khác biệt đủ nhỏ để chúng tôi tự tin rằng định dạng hình ảnh huấn luyện có tác động tối thiểu đến chất lượng đầu ra. Vì vậy, việc lưu trữ JPEG chất lượng cao không thêm bất kỳ giá trị đo lường nào so với những gì các nguồn dữ liệu đã mang sẵn. Đối với việc huấn luyện chuyển văn bản thành hình ảnh quy mô lớn, điều đó là quá đủ. Đối với dữ liệu nhạy cảm với hiện tượng giả được sử dụng để huấn luyện các mô hình Photoroom khác, chẳng hạn như mô hình AI Shadows tùy chỉnh của chúng tôi, chúng tôi vẫn dựa vào các tệp PNG.

2. Xây dựng tập dữ liệu trong Lance

Bạn không thể khám phá tương tác các bảng Parquet với hàng trăm triệu hoặc thậm chí hàng tỷ hàng. Do đó, chúng tôi lưu trữ dữ liệu trong Lance, nơi chúng tôi có thể lập chỉ mục, truy vấn và duyệt qua nó. Chúng tôi chạy quá trình nhập (ingest) với Ray Data, đọc các bảng nguồn song song và ghi nhiều phân đoạn (fragments) của một bảng Lance trên toàn cụm.

Một bài học đáng chia sẻ là về sự phân mảnh. Một bảng Lance được chia thành các phân đoạn, và một số thao tác sẽ mở rộng quy mô theo tổng số phân đoạn thay vì số lượng hàng: một lần quét sẽ mở các tệp của mọi phân đoạn, và một số thao tác siêu dữ liệu là O(số lượng phân đoạn). Do đó, một bảng bị chia thành quá nhiều phân đoạn nhỏ sẽ truy vấn chậm bất kể kích thước của nó. Hướng dẫn hiệu suất của Lance là giữ số lượng đó ở mức thấp (quy tắc chung là khoảng một trăm phân đoạn ngay cả đối với bảng một tỷ hàng) và chạy nén (compaction) định kỳ để hợp nhất các phân đoạn nhỏ thành các phân đoạn lớn hơn.

Chúng tôi đã học được điều này theo cách chậm chạp. Lần nhập đầu tiên của chúng tôi chỉ nhắm mục tiêu 100.000 hàng mỗi phân đoạn, để lại hàng ngàn phân đoạn nhỏ xíu và làm cho ngay cả các bộ lọc đơn giản và truy vấn toàn văn bản cũng trở nên chậm chạp. Sau khi nén lên khoảng một triệu hàng mỗi phân đoạn, đối với các bảng gồm vài trăm triệu hàng của chúng tôi, điều đó có nghĩa là khoảng một ngàn phân đoạn, các truy vấn đã trở nên nhanh chóng và chúng tôi không còn phải lo lắng về nó nữa. Mục tiêu phù hợp phụ thuộc vào độ rộng hàng, các mẫu truy vấn và cách dữ liệu được ghi, và nhìn lại thì các phân đoạn ít hơn, lớn hơn có thể là một mặc định tốt hơn. Lance hỗ trợ nén phân tán (ví dụ thông qua tích hợp Lance-Ray), vì vậy rất dễ điều chỉnh sau đó, nhưng tốt hơn hết là đừng để bị phân mảnh quá mức ngay từ đầu.

Các chú thích và nhúng (embeddings) hiện có để khám phá

Chúng tôi đã chọn tự chú thích lại mọi hình ảnh thay vì dựa vào các chú thích mà một số tập dữ liệu đi kèm. Lý do là tính nhất quán: trên nhiều nguồn khác nhau, độ dài, phong cách và chất lượng chú thích rất khác nhau, và chúng tôi muốn một tiêu chuẩn thống nhất duy nhất trên toàn bộ kho dữ liệu (xem thêm về bộ chú thích trong Phần 3).

Tuy nhiên, siêu dữ liệu mà một tập dữ liệu đã mang sẵn vẫn có giá trị, chỉ là cho một mục đích khác. Các chú thích có sẵn và các vector nhúng đôi khi đi kèm với chúng (ví dụ: nhúng CLIP) là những thứ giúp một tập dữ liệu có thể khám phá được trong Lance ngay lập tức: tìm kiếm toàn văn bản trên các chú thích và tìm kiếm láng giềng gần nhất (nearest-neighbor search) trên một cột nhúng cho phép chúng tôi điều hướng dữ liệu và đánh giá nhanh chất lượng cũng như việc lọc cần thiết, rất lâu trước khi chúng tôi chạy chú thích của riêng mình trên đó. Vì vậy, chúng tôi giữ cả hai, kết hợp chúng khi chúng nằm trong một bảng riêng biệt, và dựa vào chúng để duyệt dữ liệu như chúng tôi mô tả tiếp theo.

Lập hồ sơ và khám phá dữ liệu

Với dữ liệu có thể truy vấn, chúng tôi có thể lập hồ sơ (profile) nhanh chóng. Ví dụ, phân phối độ phân giải cho chúng tôi biết nơi đặt các điểm cắt. Nhóm huấn luyện nhỏ nhất của chúng tôi là 512² pixel, và chúng tôi giới hạn việc nâng cấp (upscaling) vào một nhóm ở mức 4/3 (khoảng 33%), vì vậy điểm cắt hiệu dụng là 384² ≈ 147k pixel cộng với tỷ lệ khung hình trong [0.5, 2.0]. Mọi thứ dưới mức đó chúng tôi loại bỏ.

Các bảng Lance hỗ trợ tìm kiếm toàn văn bản trên các cột văn bản cũng như tìm kiếm tương đồng láng giềng gần nhất trên các cột vector nhúng ngay khi cài đặt, và cả hai đều có thể được tăng tốc bằng cách xây dựng chỉ mục. Chúng tôi đã làm chính xác điều đó cho các cột chú thích và nhúng, sau đó xây dựng một giao diện người dùng nhỏ để duyệt tập dữ liệu: tìm kiếm văn bản thời gian thực trên các chú thích cũng như điều hướng các cụm hình ảnh tương tự về mặt thị giác thông qua tìm kiếm vector.

Việc duyệt dữ liệu theo cách này cho phép chúng tôi đánh giá tính đa dạng và chất lượng của nó một cách định tính và phát hiện sớm các vấn đề, chỉ bằng cách nhìn vào nó. Chúng tôi đã tìm thấy các chú thích cơ bản không mang nhiều thông tin, một lượng đáng ngạc nhiên các nội dung không phải ảnh chụp (ảnh chụp màn hình, slide, tài liệu, đồ họa thông tin) và một số hình ảnh gần như trùng lặp. Những quan sát này đã định hình trực tiếp phần còn lại của đường ống: chú thích lại mọi thứ (các chú thích dài hơn, chính xác hơn mô tả tốt ngay cả những hình ảnh không hoàn hảo, và đảm bảo rằng chú thích khớp với chính xác hình ảnh chúng tôi huấn luyện), và thêm các bước lọc nhẹ cũng như khử trùng lặp sau đó. Nhìn chung, chúng tôi tin rằng loại công cụ khám phá này là vô giá để có cảm nhận về một tập dữ liệu trước khi cam kết với nó.

3. Chú thích lại mọi thứ bằng VLM

Chúng tôi nhận thấy rằng các chú thích dài, chi tiết là một đòn bẩy rất mạnh đối với chất lượng đầu ra. Trong một bài kiểm chuẩn ban đầu, chúng tôi đã huấn luyện một mô hình khuếch tán (diffusion model) nhỏ hai lần trên cùng một hình ảnh, một lần trên các chú thích chúng tôi tạo bằng Qwen2.5-VL-7B và một lần trên các chú thích ngắn hơn nhiều (được tạo bằng LLaVA-1.5-LLaMA3-8B), và chấm điểm cả hai trong quá trình huấn luyện. Các chú thích dài được tạo bằng Qwen2.5-VL-7B đã giành chiến thắng trên mọi phương diện, với FID, CMMD và DINO-MMD thấp hơn tại mọi điểm kiểm tra (thấp hơn là tốt hơn trên cả ba chỉ số).