Thiết kế huấn luyện cho các mô hình Text-to-Image: Bài học từ các thử nghiệm loại bỏ

Chào mừng trở lại! Đây là phần thứ hai trong loạt bài của chúng tôi về việc huấn luyện các mô hình text-to-image hiệu quả từ đầu.

Trong bài viết đầu tiên của loạt bài này, chúng tôi đã giới thiệu mục tiêu của mình: huấn luyện một mô hình nền tảng text-to-image có tính cạnh tranh hoàn toàn từ đầu, công khai và ở quy mô lớn. Chúng tôi tập trung chủ yếu vào các lựa chọn kiến trúc và giải thích các quyết định thiết kế cốt lõi đằng sau mô hình PRX của chúng tôi. Chúng tôi cũng đã phát hành một phiên bản nhỏ (1,2 tỷ tham số) ban đầu của mô hình dưới dạng xem trước những gì chúng tôi đang xây dựng (hãy thử nó nếu bạn chưa làm 😉).

Trong bài viết này, chúng tôi chuyển trọng tâm từ kiến trúc sang huấn luyện. Mục tiêu là ghi lại những gì thực sự tạo nên sự khác biệt đối với chúng tôi khi cố gắng làm cho các mô hình huấn luyện nhanh hơn, hội tụ đáng tin cậy hơn và học các biểu diễn tốt hơn. Lĩnh vực này đang phát triển nhanh chóng và danh sách các “mẹo huấn luyện” ngày càng dài, vì vậy thay vì cố gắng khảo sát một cách toàn diện, chúng tôi đã cấu trúc bài viết này như một nhật ký thử nghiệm: chúng tôi tái tạo (hoặc điều chỉnh) một loạt ý tưởng gần đây, triển khai chúng trong một thiết lập nhất quán và báo cáo cách chúng ảnh hưởng đến quá trình tối ưu hóa và hội tụ trong thực tế. Cuối cùng, chúng tôi không chỉ báo cáo các kỹ thuật này một cách riêng lẻ; chúng tôi cũng khám phá xem những kỹ thuật nào vẫn hữu ích khi kết hợp với nhau.

Trong bài viết tiếp theo, chúng tôi sẽ xuất bản công thức huấn luyện đầy đủ dưới dạng mã, bao gồm cả các thử nghiệm trong bài viết này. Chúng tôi cũng sẽ chạy và báo cáo về một “speedrun” công khai nơi chúng tôi kết hợp các phần tốt nhất vào một cấu hình duy nhất và kiểm tra căng thẳng nó từ đầu đến cuối. Bài tập này sẽ vừa đóng vai trò là một bài kiểm tra căng thẳng cho quy trình huấn luyện hiện tại của chúng tôi, vừa là một minh chứng cụ thể về mức độ xa có thể đi xa với thiết kế huấn luyện cẩn thận trong các ràng buộc chặt chẽ. Nếu bạn chưa làm vậy, chúng tôi mời bạn tham gia Discord của chúng tôi để tiếp tục thảo luận. Một phần đáng kể của dự án này đã được định hình bởi các trao đổi với các thành viên cộng đồng, và chúng tôi đặt giá trị cao vào phản hồi bên ngoài, các thử nghiệm loại bỏ và các diễn giải thay thế về kết quả.

Cơ sở

Trước khi giới thiệu bất kỳ kỹ thuật nào tăng hiệu quả huấn luyện, trước tiên chúng tôi thiết lập một lần chạy tham chiếu rõ ràng. Cơ sở này cố tình đơn giản. Nó sử dụng các thành phần tiêu chuẩn, tránh các mục tiêu phụ trợ và không dựa vào các lối tắt kiến trúc hoặc thủ thuật để tiết kiệm tài nguyên tính toán. Vai trò của nó là phục vụ như một điểm so sánh ổn định cho tất cả các thử nghiệm tiếp theo.

Cụ thể, đây là một thiết lập huấn luyện Flow Matching thuần túy (như được giới thiệu trong Phần 1) mà không có mục tiêu bổ sung và không có các thủ thuật tăng tốc kiến trúc. Chúng tôi sẽ sử dụng mô hình PRX-1.2B nhỏ mà chúng tôi đã trình bày trong bài viết đầu tiên của loạt bài này (kiến trúc luồng đơn với sự chú ý toàn cục cho các token hình ảnh và token văn bản) làm cơ sở và huấn luyện nó trong không gian tiềm ẩn của Flux VAE, giữ cấu hình cố định trên tất cả các so sánh trừ khi có quy định khác.

Cấu hình huấn luyện cơ sở như sau:

Cài đặt	Giá trị
Bước	100k
Tập dữ liệu	Bittensor_subnet_19_06_04_24 1 triệu ảnh tổng hợp do MidJourneyV6 tạo ra
Độ phân giải	256×256
Batch toàn cục	256
Trình tối ưu hóa	AdamW
lr	1e-4
weight_decay	0.0
eps	1e-15
betas	(0.9, 0.95)
Bộ mã hóa văn bản	GemmaT5
Mã hóa vị trí	Rotary (RoPE)
Mặt nạ chú ý	Mặt nạ đệm
EMA	Vô hiệu hóa

Cấu hình cơ sở này cung cấp một điểm neo minh bạch và có thể tái lập. Nó cho phép chúng tôi quy kết các cải tiến và suy thoái được quan sát cho các can thiệp huấn luyện cụ thể, thay vì cho sự thay đổi siêu tham số hoặc thay đổi thiết lập ẩn. Trong phần còn lại của bài viết này, mọi kỹ thuật đều được đánh giá so với tham chiếu này với một câu hỏi chỉ đạo duy nhất trong tâm trí:

Liệu sự sửa đổi này có cải thiện sự hội tụ hoặc hiệu quả huấn luyện so với cơ sở hay không?

Các chỉ số đánh giá hiệu suất

Để giữ bài viết này thực tế, chúng tôi dựa vào một bộ chỉ số nhỏ để theo dõi các điểm kiểm tra theo thời gian. Không có chỉ số nào trong số đó là một đại diện hoàn hảo cho chất lượng hình ảnh cảm nhận được, nhưng cùng nhau chúng cung cấp một bảng điểm thực tế khi chúng tôi lặp lại.

• Fréchet Inception Distance (FID): (Heusel et al., 2017) Đo lường mức độ gần của phân phối các hình ảnh được tạo ra và hình ảnh thực tế, sử dụng thống kê đặc trưng của Inception-v3 (trung bình và hiệp phương sai). Các giá trị thấp hơn thường tương quan với chất lượng mẫu cao hơn.
• CLIP Maximum Mean Discrepancy (CMMD): (Jayasumana et al., 2024) Đo khoảng cách giữa phân phối hình ảnh thực tế và hình ảnh được tạo ra bằng cách sử dụng các nhúng hình ảnh CLIP và Maximum Mean Discrepancy (MMD). Không giống như FID, CMMD không giả định phân phối đặc trưng là Gaussian và có thể hiệu quả hơn về mẫu; trong thực tế, nó thường theo dõi chất lượng cảm nhận tốt hơn FID, mặc dù nó vẫn là một đại diện không hoàn hảo.
• DINOv2 Maximum Mean Discrepancy (DinoMMD): Cùng một khoảng cách MMD như CMMD, nhưng được tính toán trên các nhúng hình ảnh DINOv2 (Oquab et al. 2023) thay vì CLIP. Điều này cung cấp một cái nhìn bổ sung về sự dịch chuyển phân phối dưới một khung nhìn tự giám sát.
• Thông lượng mạng: Số lượng mẫu trung bình được xử lý mỗi giây (mẫu/giây), như một thước đo hiệu quả huấn luyện từ đầu đến cuối.

Với bảng điểm được xác định, giờ đây chúng ta có thể đi sâu vào các phương pháp mà chúng tôi đã khám phá, được nhóm thành bốn nhóm: Căn chỉnh biểu diễn, Mục tiêu huấn luyện, Định tuyến và làm thưa token, và Dữ liệu.

Căn chỉnh biểu diễn

Mô hình khuếch tán và mô hình luồng thường được huấn luyện với một mục tiêu duy nhất: dự đoán một mục tiêu giống nhiễu (hoặc trường vector) từ một đầu vào bị lỗi. Ngay từ đầu quá trình huấn luyện, mục tiêu duy nhất đó đang thực hiện hai nhiệm vụ cùng một lúc: nó phải xây dựng một biểu diễn nội bộ hữu ích và học cách khử nhiễu dựa trên đó. Căn chỉnh biểu diễn làm cho điều này trở nên rõ ràng bằng cách giữ lại mục tiêu khử nhiễu và thêm một tổn thất phụ trợ giám sát trực tiếp các đặc trưng trung gian bằng một bộ mã hóa thị giác mạnh mẽ, đóng băng. Điều này có xu hướng tăng tốc quá trình học sớm và đưa các đặc trưng của mô hình lại gần hơn với các bộ mã hóa tự giám sát hiện đại. Do đó, bạn thường cần ít tính toán hơn để đạt được cùng một chất lượng.

Một cách hữu ích để xem xét nó là phân tách bộ khử nhiễu thành một bộ mã hóa ngầm tạo ra các trạng thái ẩn trung gian và một bộ giải mã ánh xạ các trạng thái đó đến mục tiêu khử nhiễu. Lập luận là học biểu diễn là yếu tố hạn chế: các bộ chuyển đổi khuếch tán và luồng học các đặc trưng phân biệt, nhưng chúng tụt hậu so với các bộ mã hóa thị giác nền tảng khi quá trình huấn luyện bị giới hạn về tính toán. Do đó, việc mượn một không gian biểu diễn mạnh mẽ có thể làm cho vấn đề khử nhiễu trở nên dễ dàng hơn.

REPA (Yu et al., 2024)

REPA bổ sung một thuật ngữ khớp biểu diễn lên trên mục tiêu gốc flow-matching. Cho x0∼pdatax_0 \sim p_{\text{data}}x0​∼pdata​ là một mẫu sạch và x1∼ppriorx_1 \sim p_{\text{prior}}x1​∼pprior​ là mẫu nhiễu. Mô hình được huấn luyện trên một trạng thái nội suy xtx_txt​ (t∈[0,1]t \in [0,1]t∈[0,1]) và dự đoán một trường vector vθ(xt,t)v_\theta(x_t, t)vθ​(xt​,t). Trong REPA, một bộ mã hóa thị giác được đào tạo trước fff xử lý mẫu sạch x0x_0x0​ để tạo ra các nhúng vá y0=f(x0)∈RN×Dy_0 = f(x_0) \in \mathbb{R}^{N \times D}y0​=f(x0​)∈RN×D, trong đó NNN là số lượng token vá và DDD là chiều nhúng của giáo viên. Song song đó, bộ khử nhiễu xử lý xtx_txt​ và tạo ra các token ẩn trung gian hth_tht​ (một token cho mỗi vá). Một đầu chiếu nhỏ hϕh_\phihϕ​ ánh xạ các token ẩn sinh viên này vào không gian nhúng của giáo viên, và một tổn thất phụ trợ tối đa hóa sự tương tự vá giữa các token giáo viên và sinh viên tương ứng: LREPA(θ,ϕ)=−Ex0,x1,t[1N∑n=1Nsim(y0,[n], hϕ(ht,[n]))]

\mathcal{L}{\text{REPA}}(\theta,\phi) = -\mathbb{E}{x_0,x_1,t}\Big[\frac{1}{N}\sum_{n=1}^{N} \text{sim}\big(y_{0,[n]},, h_\phi(h_{t,[n]})\big)\Big]

LREPA​(θ,ϕ)=−Ex0​,x1​,t​[N1​n=1∑N​sim(y0,[n]​,hϕ​(ht,[n]​))].

Link bài viết gốc

• Tags:
• Ai
• 03 February 2026
• Huggingface.co