Krea 2: Báo cáo kỹ thuật về mô hình hình ảnh 12B trọng số mở

 (krea.ai)
1 điểm bởi GN⁺ 4 giờ trước | 1 bình luận | Chia sẻ qua WhatsApp
  • Krea 2 là một mô hình nền tảng tạo ảnh coi trọng khám phá sáng tạo hơn là một polished default duy nhất, công khai trọng số mô hình và suy luận theo permissive license
  • Quy trình huấn luyện diễn ra theo chuỗi pretraining → midtraining → SFT → preference optimization → RL, trong đó tuyển chọn dữ liệu, chú thích, mở rộng prompt và tham chiếu phong cách từng bước tinh chỉnh phân phối đầu ra
  • Kiến trúc dựa trên họ DiT đơn giản, kết hợp GQA, gated sigmoid attention, SwiGLU, Qwen 3 VL, Qwen Image VAE, FLUX 2 VAE, v.v. để cân bằng độ ổn định và hiệu quả
  • Krea 2 lọt top 10 trên Artificial Analysis leaderboard ở hạng mục text-to-image và đứng thứ 2 trong số các mô hình từ independent labs
  • Để huấn luyện quy mô lớn, nhóm đã xây dựng hệ thống dựa trên PyTorch, FSDP2, song song hóa tensor, Kubernetes, Virtual Kubelet, Weka và PostgreSQL; ở bước tiếp theo đang xem xét MoE, sparse attention, native 2K–4K, NVFP4 và Muon scaling

Mô hình nền tảng hình ảnh hướng tới khám phá sáng tạo

  • Krea 2 là một dòng mô hình nền tảng tạo ảnh nhằm đạt sự đa dạng thẩm mỹ rộng và khả năng kiểm soát sáng tạo của người dùng
  • Tài liệu công khai được cung cấp tại Release page, Hugging Face weights/license, GitHub code/license, Krea Image tool
  • Trọng số mô hình và suy luận được công khai theo permissive license
  • Krea cho rằng các mô hình hình ảnh dựa trên diffusion và flow-matching đã tiến bộ tới mức tạo ảnh độ phân giải cao, photorealism, cấu trúc ổn định, render chữ dày đặc, kiến thức thế giới rộng và tuân thủ prompt chi tiết, nhưng nhiều hệ thống lại hội tụ vào một mỹ học mặc định hẹp
  • Thay vì tối ưu hóa một polished default duy nhất, Krea 2 hướng tới một phương tiện tạo sinh cho phép khám phá nhiều phong cách, bầu không khí, bố cục và hướng thị giác khác nhau
  • Mô hình lọt top 10 trên Artificial Analysis leaderboard ở hạng mục text-to-image và đứng thứ 2 trong số các mô hình từ independent labs

Tuyển chọn dữ liệu và chiến lược chú thích

  • Nhóm Krea đã xây dựng hạ tầng dữ liệu quy mô lớn và framework huấn luyện phân tán từ đầu để tạo pretraining dataset có broad world knowledge và style coverage
  • Họ cho rằng một “data mix tốt” không chỉ cần ảnh chất lượng cao mà còn cần tính đa dạng và domain coverage rộng
  • Việc lọc dựa trên aesthetic-score và image-quality-assessment có thể tạo ra implicit bias
    • motion blur hoặc softness có thể là lựa chọn nghệ thuật có chủ đích, nhưng có thể bị chấm thấp như blurry image
    • Nếu caption mô tả chính xác hình ảnh, ngay cả những hình ảnh không mong muốn cũng có thể hữu ích cho huấn luyện downstream
  • Trong pretraining dataset, họ lọc duplicated samples, các concepts bị over-represented, mẫu mà VLM không nắm bắt được yếu tố quan trọng, mẫu gây ra undesired biases và artifacts, mẫu có high visual complexity khó mô hình hóa ổn định ở low resolution, và AI-generated samples
  • Krea 2 pretraining mix không sử dụng AI-generated images
    • synthetic data và distillation có thể là shortcut để đạt được capability
    • Họ cho rằng ngay cả một lượng nhỏ AI-generated images cũng đưa bias vào phân phối đầu ra của mô hình và trên thực tế ấn định upper bound của model quality
    • Để lọc chúng, họ thiết kế in-house classifiers
  • Caption được xây dựng theo phương pháp multi-stage
    • Chạy OCR model trên target image để trích xuất visible text
    • Cung cấp kết quả OCR và metadata cho captioning model để tạo enriched caption bao gồm extracted text và world knowledge
    • Dùng một LLM rẻ hơn để tái cấu trúc context-rich long-form caption thành nhiều độ dài và định dạng khác nhau, giúp mô hình tiếp xúc với nhiều prompt style
  • long prompts cung cấp dense supervision, tạo ra hội tụ nhanh hơn và training loss thấp hơn; đồng thời họ vẫn duy trì việc cho mô hình tiếp xúc với short/medium prompt để phục vụ sử dụng downstream

Dữ liệu huấn luyện theo độ phân giải và midtraining

  • Dữ liệu pretraining trải qua các giai đoạn độ phân giải 256px, 512px, 1024px
    • Phân bổ phần lớn FLOPs cho giai đoạn độ phân giải thấp để học core capability một cách hiệu quả
    • Sau đó tăng dần độ phân giải để cung cấp capability tạo sinh high-fidelity
    • Pretraining ở độ phân giải thấp học basic text-image alignment và structure
  • Dataset độ phân giải thấp có quy mô billions of images nên phụ thuộc nhiều vào các bộ lọc chi phí thấp dựa trên CPU
    • Loại bỏ ảnh không phù hợp bằng các bộ lọc broken-file, resolution, aspect-ratio
    • Loại bỏ ảnh có extreme textures và noise patterns bằng Laplacian filters
    • Giảm flat-color backgrounds và border artifacts bằng RGB entropy, white/black pixel ratios, custom heuristics, in-house classifiers
  • In-house classifier được xây dựng bằng cách dùng large VLM tạo system prompt cho filtering task, sinh pseudo-labeled dataset, rồi huấn luyện classifier dựa trên small DINOv3 hoặc SigLIP-2
    • Ở giai đoạn độ phân giải thấp, filtering model cần GPU compute được giữ dưới 1B parameters để đảm bảo hiệu quả
  • Low-resolution deduplication chủ yếu dùng các hash-based methods kết hợp md5, phash, colorhash
    • phash 8x8 cơ bản không xét đến color nên false-positive rate cao
    • Để deduplication mạnh mẽ hơn, kết hợp phash 12x12 với colorhash
  • Khi training resolution tăng lên, các image-quality và aesthetic filters được đưa vào
    • Quality score chỉ được dùng để loại bỏ ảnh có chất lượng rất poor quality, không dùng cho oversampling dựa trên score
    • Loại trừ các ảnh khó biểu đạt text và content một cách có ý nghĩa ở low resolution bằng OCR-based image-complexity score và text density
  • Huấn luyện sparse autoencoder trên SigLIP-2 embeddings để tạo SAE-based tagging system, và dùng để lọc clear visual artifacts mà không cần explicit classifier
  • Khác với pretraining, midtraining chọn rõ các image sources cung cấp stylistic coverage tốt và high-quality images trong những visual domain cụ thể
    • Pretraining là quy trình bottom-up bắt đầu từ general pool
    • Midtraining là quy trình tuyển chọn top-down, chọn domains và sources trước
    • Đây là bước nối mượt giữa general pretraining distribution và high-quality SFT distribution
  • Bổ sung world-knowledge coverage bằng semantic clustering và retrieval-based strategies
    • Thực hiện hierarchical k-means clustering bằng FAISS
    • VLM kiểm tra ảnh gần cluster centroid để đặt tên cho cluster và flag khi cần
    • Các flagged clusters được human review để loại bỏ cluster low quality hoặc problematic
    • Trong các leaf cluster còn lại, thực hiện semantic deduplication bằng SigLIP similarity
  • Để có named entity coverage, dùng Danker chạy PageRank trên English Wikipedia và giữ top 90% articles theo rank
    • Loại bỏ unrepresentable subjects bằng Wikidata metadata
    • Thực hiện full-text search trên caption của toàn bộ dataset cho khoảng 5 million concepts còn lại
    • Khi sampling, ưu tiên ảnh có caption nhắc đến rare concepts

Lựa chọn kiến trúc và ablation

  • Krea 2 đã phát triển một kiến trúc diffusion transformer (DiT) đơn giản nhưng hiệu năng tốt thông qua ablation
  • Ablation kiến trúc được đánh giá theo bốn hạng mục: stability, performance, efficiency, simplicity
    • stability xem xét việc giảm loss/gradient spike và độ ổn định khi huấn luyện
    • performance xem xét tốc độ hội tụ và khả năng duy trì ở độ phân giải cao cũng như horizon dài
    • efficiency xem liệu có thể giảm parameter count, FLOPs, memory, communication mà không làm tổn hại quality hay không
    • simplicity kiểm tra liệu có thể đơn giản hóa model mà không làm ảnh hưởng các hạng mục khác hay không
  • Nhiều quyết định kiến trúc chịu ảnh hưởng từ xu hướng áp dụng trong LLM space, và nhóm cho rằng các kernel và optimization của LLM ecosystem cũng có thể được tận dụng trong diffusion model
  • Các lựa chọn chính của kiến trúc cuối cùng như sau
    • Attention sử dụng GQA with gated sigmoid attention
    • MLP được đổi từ GeLU MLP sang các SwiGLU layer với 4x expansion factor
    • Residual giữ standard residual
    • Text encoder sử dụng Qwen 3 VL
    • Modulation được đổi từ per-block MLP modulation sang light modulation with bias
    • Autoencoder sử dụng Qwen Image VAE và FLUX 2 VAE
    • Block design sử dụng single stream transformer block
    • Norm sử dụng zero-center RMSNorm và QKNorm
    • Positional encoding giữ 3D Axial RoPE
  • GQA cải thiện computational efficiency trong khi chỉ gây minimal degradation
    • MLA cho thấy gain nhẹ so với GQA, nhưng không được chọn do additional computational overhead
    • gated sigmoid attention không tạo performance gain lớn, nhưng cho thấy dynamics ổn định hơn trên loss và gradient-norm curves
  • Giữa single-stream, dual-stream và hybrid-stream design không có khác biệt performance lớn; hybrid-stream nhỉnh hơn một chút nhưng single-stream blocks được dùng vì tính đơn giản
  • Per-block MLP modulation của MMDiT có thể chiếm 20–30% total parameter count, nên Krea 2 thay thế nó bằng per-block tunable bias term
  • Trong thí nghiệm timestep conditioning, ở 256px, 4–16 timestep tokens đủ để thay AdaLN, nhưng ở 512px và 1024px thì kém hơn AdaLN baseline
  • Positional encoding cuối cùng là 3D axial RoPE, trong đó head dimensions được phân bổ cho frame, height, width
    • RoPE indices của text tokens được đặt bằng zero
    • partial RoPE cho kết quả zero-shot inference tốt khi scale từ 256px lên 512px, nhưng sau high-resolution training, hiệu năng cuối cùng thấp hơn baseline
  • Autoencoder bắt đầu với FLUX.1-dev autoencoder làm baseline và so sánh với Qwen Image VAE, DC-AE, FLUX 2 VAE, internal autoencoder
    • DC-AE được xem là đặt ra hard upper limit cho khả năng phân giải fine detail do reconstruction error
    • Qwen Image VAE và FLUX 2 VAE duy trì reconstruction quality xuất sắc đồng thời latent space cho convergence nhanh hơn nhiều
    • Các early model dùng Qwen Image autoencoder, còn các larger model áp dụng FLUX 2 VAE
  • Text encoder so sánh T5-XXL, T5Gemma, umT5, Qwen 2.5 VL, Qwen 3 VL, và sử dụng Qwen 3 VL làm text encoder cuối cùng
    • VLM cung cấp richer input space bao gồm text và image, cùng khả năng multilingual generalization mạnh hơn
    • Thay vì chỉ dùng last layer của VLM feature, nhóm đưa vào shallow attention layer để aggregate hidden features từ toàn bộ layers
    • Thêm lightweight bidirectional transformer layers trên token axis để giảm autoregressive bias

Pipeline huấn luyện, tối ưu hóa ưu tiên, RL

  • training pipeline là cấu trúc nhiều giai đoạn lấy cảm hứng từ modern LLM training pipeline
  • pretraining thiết lập các basic capabilities như text-image alignment, text rendering, stylistic coverage, structural consistency
    • final model được huấn luyện với standard rectified-flow loss và v-parameterization
    • ở first epoch của 256px stage, dùng iREPA để tăng tốc đáng kể early stage convergence rồi sau đó loại bỏ
    • ở các 256px và 512px stages, quan sát thấy training speed gain 15–20% so với bf16 baseline nhờ 8-bit training
    • từ 1024px đến final RL stage, dùng standard bf16 training
  • trong high-resolution pretraining, việc thích nghi resolution-dependent timeshift schedule là quan trọng
    • cả training và inference đều dùng shifted logit-normal sampling schedule
    • khi resolution tăng, shift được tăng dần
    • sweep chỉ áp dụng cho training shift, còn inference shift schedule được giữ constant
  • trong pretraining, dùng warmup-stable-decay learning-rate schedule và áp dụng PMA
    • PMA đạt performance comparable với EMA đồng thời tránh significant memory overhead của EMA
  • optimizer dùng AdamW làm primary optimizer xuyên suốt pipeline
    • Muon hội tụ nhanh hơn AdamW ở initial steps, nhưng ở longer horizons cho performance thấp hơn và có stability issues
    • khi loại first and last linear layers của MMDiT khỏi Muon parameters và thêm Nesterov momentum, nó nhất quán vượt AdamW baseline ở cả low/high resolution
    • ở pretraining run mới nhất, do hạn chế thời gian nên không áp dụng Muon, và dự định sẽ áp dụng trong pretraining cycle tiếp theo
  • ở SFT stage, curate một small dedicated set of highly aesthetic images
    • mục tiêu là bias model nhiều hơn theo các aesthetically desirable directions
    • đặc biệt hữu ích để xử lý high-saturation và texture issues thường gặp ở earlier checkpoints
    • sau khi huấn luyện domain-specific SFT checkpoints, tạo generalist SFT checkpoint bằng model merging
  • preference optimization là bước đầu tiên của post-training stack và gồm two-stage pipeline
    • giai đoạn 1 thực hiện initial refinement bằng large-scale synthetic preference-pair generation pipeline
    • majority của preference pairs chứa ít nhất một on-policy sample
    • giai đoạn 2 là calibration stage chỉ dùng human annotations
    • human annotations được thu thập bởi nhân sự nội bộ quen thuộc với strengths, weaknesses, quirks của model
  • trong PO, policy divergence xuất hiện như một common phenomenon
    • các phương pháp kiểu DPO khuyến khích tăng margin giữa preferred sample likelihood và dispreferred sample likelihood
    • trong nhiều preference-dataset mixtures, quan sát thấy mô hình đạt objective bằng cách giảm generation likelihood của cả hai sample, nhưng với rate khác nhau
    • divergence làm mô hình rời xa general pretraining distribution và biểu hiện thành high-frequency artifacts ở giai đoạn sau của training
    • để giảm thiểu điều này, thiết kế một DPO variant gọi là STPO
  • RL là final stage của training pipeline
    • dùng multi-reward GRPO-style method
    • reward models gồm general aesthetic model, prompt-following reward, text-rendering reward, artifact and structure reward
    • general aesthetic model thu được bằng cách finetune open-source VLM với preference data thu thập ở PO stage
  • prompt-specific rubric reward phân rã prompt thành verifiable requirements và đánh giá generated image có đáp ứng chúng hay không
    • giúp thỏa mãn các fine-grained prompt constraints thay vì quy prompt following về generic image quality
  • để giảm structural artifacts, huấn luyện dedicated artifact reward model
    • các lỗi như extra fingers, malformed limbs, distorted text rất rõ với con người nhưng general-purpose VLM judges thường bỏ sót
  • toàn bộ RL stage được huấn luyện không có CFG
    • cải thiện nhanh conditional model distribution, khiến no-CFG samples ở đầu training gần với guided samples hơn nhiều
    • tại inference time, vẫn có thể bật CFG như một control knob bổ sung
  • sau RL stage có optional timestep-distillation stage
    • đã xem xét DMD, DMD2, Decoupled DMD, piFlow, APT nhưng chọn Trajectory Distribution Matching(TDM)
    • TDM áp dụng DMD trên toàn bộ timesteps để thực hiện distribution matching ở trajectory level

Mở rộng prompt và tham chiếu phong cách

  • Khi huấn luyện, mô hình sử dụng các caption phong phú mô tả chi tiết thị giác dày đặc của ảnh, nhưng đầu vào thực tế của người dùng thường ngắn, mơ hồ và có thói quen diễn đạt đa dạng
  • Prompt expander chuyển đổi user prompt đơn giản hoặc chưa đủ thông tin thành định hướng thị giác phong phú hơn mà không ghi đè ý định của người dùng
    • Được huấn luyện bằng pipeline SFT 2 giai đoạn và RL trên một LLM mã nguồn mở hiện có
    • Mục tiêu không chỉ gồm cải thiện chất lượng ảnh mà còn bao gồm creative variation và controllable exploration
  • Dữ liệu SFT được tạo bằng cách sinh synthetic “user captions” từ long captions
    • Synthetic user captions là các prompt ngắn, mang tính hội thoại và bán chỉ dẫn, cố ý lược bỏ nhiều visual details của target caption
    • Tạo paired data dạng underspecified user prompt → expanded model-friendly caption
    • Cũng tạo synthetic thinking traces để bảo toàn reasoning ability
  • Cũng áp dụng một lượng nhỏ targeted distribution shaping
    • Oversample visually rich and artistic imagery
    • Thêm lightweight photographic-medium bias cho các prompt cần được mở rộng thành photorealistic descriptions
    • Mục tiêu không phải là ép buộc house style, mà là bao gồm cả expressive art-directed imagery lẫn straightforward photorealistic requests
  • RL cho prompt expander nhằm thoát khỏi việc bắt chước target caption, tạo ra các expansion cải thiện image quality trong khi vẫn bảo toàn user intent
    • Được huấn luyện bằng GDPO và multi-reward objective
    • Image-level rewards đo quality và preference của resulting generations
    • Prompt-level verifiable rewards kiểm tra expansion có trung thành với original request hay không
    • Safety và constraint checks được dùng làm gates cho overall reward
  • Một trong các failure mode của prompt expander là diversity collapse
    • Khi image rewards chiếm ưu thế, mô hình có thể học một single safe high-reward house style
    • Thêm DINOv3 embedding diversity score trên prompt groups để thưởng cho intra-group visual diversity cùng với quality và alignment
    • Để duy trì variation, cần giữ diversity reward active trong suốt quá trình training
  • Style-reference system được xây dựng trên base model
    • Người dùng có thể tạo ảnh bằng text đồng thời dùng một hoặc nhiều reference images để guide output style
    • Mục tiêu thiết kế gồm smooth semantic mixing của multiple styles, continuous control cường độ của từng style reference, và adherence ở mức state-of-the-art với complex styles
    • Một failure mode phổ biến là content và subject matter của style image bị leakage sang final image
    • Nhóm đã thiết kế một self-supervised technique để huấn luyện style-reference module, sau đó align thêm outputs bằng bước preference-optimization

Hạ tầng và vận hành huấn luyện phân tán

  • Framework huấn luyện phân tán của Krea được xây dựng từ đầu dựa trên PyTorch, chủ yếu sử dụng các tính năng native của torch được hỗ trợ bởi abstraction DTensor và dự án torchtitan
    • Hầu hết các lần chạy tiền huấn luyện và hậu huấn luyện sử dụng kết hợp FSDP2 với song song tensor kiểu Megatron-LM
    • Trong các cấu hình có kích thước TP lớn hơn 2, bật async-TP bằng cờ torch.compile để đạt mức tăng tốc vừa phải so với TP naive
    • Các tham số autoencoder được nhân bản trên tất cả thiết bị, còn text encoder và backbone MMDiT chính mới được sharding
    • Sử dụng NVLinkSharp cho kết nối nội bộ node và InfiniBand cho kết nối giữa các node
  • Để tăng hiệu quả huấn luyện, họ dùng một mô hình hơi rộng hơn với hidden dimension lớn hơn
    • Khi hidden size tăng, cường độ tính toán của từng layer cao hơn, giúp FSDP2 prefetching dễ che giấu độ trễ hơn
    • Giảm số lượng layer làm giảm số phép toán all-gather và reduce-scatter
    • Thay đổi này giúp giảm đáng kể các lỗi liên quan đến NCCL trong toàn bộ quá trình tiền huấn luyện
    • Kích thước phép nhân ma trận lớn hơn giúp bù lại overhead quantization/dequantization của huấn luyện 8-bit
  • Trọng tâm của chiến lược tối ưu hóa là torch.compile
    • Với attention, mặc định sử dụng các kernel cuDNN mới nhất và khi cần thì dùng FlexAttention hoặc FlashAttention 3
    • Ở độ phân giải thấp, sử dụng selective activation checkpointing
    • Ở độ phân giải cao, activation bắt đầu chiếm phần lớn bộ nhớ nên sử dụng full activation checkpointing
  • Định dạng cơ bản cho nạp dữ liệu là Parquet
    • Mỗi row lưu tham chiếu ảnh, kích thước crop/resize, caption và các metadata khác
    • Với các lần chạy quy mô lớn, row được shuffle và packing trước để nạp các batch ảnh có cùng aspect ratio
    • Nhờ packing, latent có thể được mã hóa bằng một lần pass autoencoder duy nhất
  • Trong huấn luyện phân tán quy mô lớn, một GPU đơn lẻ bị lỗi hoặc một straggler cũng có thể làm dừng toàn bộ lần chạy
    • Ở quy mô của Krea, cách tối ưu MTBF và MTTR bằng checkpointing nhanh, thường xuyên và cải thiện startup time là đủ
  • Nghiên cứu được chạy trên một cụm Kubernetes duy nhất dùng chung GPU với production inference
    • Hệ thống được thiết kế để có thể chiếm toàn bộ GPU pool khi nghiên cứu cần
    • Khi tất cả GPU trong cụm được cấp cho các lần chạy huấn luyện, inference workload của Krea tự động migrate sang nơi khác
    • Hệ thống xử lý traffic failover để duy trì responsiveness của production ngay cả khi không còn GPU cục bộ
  • Kueue là thành phần cốt lõi của workload scheduling
    • Kueue cung cấp hệ thống ưu tiên 2 tầng, kết hợp Workload priority và Kubernetes Pod priority
    • Cho phép gang-scheduling cần thiết cho huấn luyện multi-node
    • Các primitive queueing “borrowing”, “lending”, “reclamation” giúp tối đa hóa utilization
  • Khi tất cả GPU được cấp cho nghiên cứu, Virtual Kubelet được dùng cho các thành phần scale inference ở nơi khác
    • Khi pod được schedule vào virtual Kubernetes node, mã của Krea chuyển đổi pod specification sang dạng tương thích với target provider
    • Nếu xảy ra failure phía provider, hệ thống reconcile trạng thái hai bên
    • Recovery được ủy thác cho Kubernetes, còn hệ thống phát hiện failure và truyền nó sang Kubernetes
  • Observability là lĩnh vực họ học được nhiều nhất trong tiền huấn luyện quy mô lớn
    • Không có các subsystem metric liên quan đến GPU, PCIe, NVLink và InfiniBand thì không thể huấn luyện ở quy mô này
    • Metric được thu thập bằng sự kết hợp giữa DCGM và custom DaemonSet
    • Khi GPU vượt 75–78°C, throttling bắt đầu và throughput tổng thể giảm, làm tăng training instability
    • DCGM_FI_PROF_PIPE_TENSOR_ACTIVE là indicator được ưu tiên để xác định liệu quá trình huấn luyện có diễn ra như kỳ vọng hay không
    • InfiniBand metric là thiết yếu để chẩn đoán fabric instability, link flapping, packet error, congestion, symbol error và throughput disparity
  • Scaling số lượng GPU rất khó
    • Các run dưới 128 GPU rất ổn định và thường chạy nhiều ngày không gặp vấn đề
    • Khi tăng số GPU, các run bắt đầu crash thường xuyên hơn nhiều
    • Ở quy mô rất lớn, không hoàn tất được run nào vượt quá 24 giờ
    • Nhiều crash không có nguyên nhân rõ ràng và xuất hiện như NCCL timeout trong khi mọi metric đều trông healthy
  • Một trong những sai lầm lớn ban đầu là chọn Ceph, sau đó họ chuyển sang Weka
    • Các vấn đề liên quan đến filesystem và downtime giảm mạnh, performance cũng cải thiện với mức tương tự
    • Weka là yếu tố then chốt cho phép checkpointing aggressive trong huấn luyện Krea 2
    • Checkpoint hoàn tất chỉ trong khoảng 30 giây, nên thời gian mất cho checkpointing là rất ít

Kho dữ liệu và hàng đợi tác vụ

  • Để thu thập và tuyển chọn dữ liệu cho K2, nhóm đã xây dựng một custom warehousing and queueing system xoay quanh cluster máy chủ PostgreSQL
  • Mỗi Krea tablet server được gọi là “krablet”
    • Mỗi krablet gồm một Postgres instance chứa một data shard và một deployment server “funnel” dùng để batch/queue bất đồng bộ các mutation nhằm giảm lock contention
  • Mọi thao tác read đều được proxy thông qua deployment server “RPC” quy mô lớn
    • RPC server thay thế traditional connection pooler như PgBouncer
    • Mỗi RPC server duy trì connection pool tới mọi shard của database
  • Hệ thống krablet đã scale tới 208TB chỉ riêng metadata và có thể xử lý hàng chục nghìn transaction UPSERT bị tranh chấp mỗi giây
    • Cung cấp single source of truth cho toàn bộ research data
    • Cho phép stream-processing layer trở nên đồng nhất với data layer
  • Workflow job-processing thông thường dùng Postgres table như một hàng đợi
    • OCR worker tìm và xử lý các row có contains_text IS NULL
    • embed worker xử lý các row có embedding_path IS NULL và contains_text = FALSE
    • Claim row bằng FOR UPDATE SKIP LOCKED và cập nhật các column kiểu last_tried_at
  • Mô hình queue có hành vi retry khác với Kafka hay Ray
    • Khi failure, không drop row hoặc gửi vào dead-letter queue
    • Nhờ atomic update last_tried_at, các row xử lý thất bại cũng được retry ở cuối queue
    • Đồng thời tránh head-of-line blocking
  • Có thể điều chỉnh số lượng worker một cách động
    • Processing job được deploy bằng Kubernetes và có thể scale up/down tùy ý mà không cần data resharding
    • Job có thể chạy với 1 worker hoặc 1000 worker
    • Bằng Prometheus scaling metric, từng phần của pipeline có thể autoscale theo lượng available work
  • Để thuận tiện cho nhà nghiên cứu, nhóm cung cấp một hệ thống tên là “pluck”
    • Cung cấp global map API phù hợp để dùng trong notebook
    • t.map trả về một handle mà người dùng có thể attach vào để xem live progress
    • UDF được pickle bằng cloudpickle và chạy trên remote worker
  • Cho thế hệ nghiên cứu tiếp theo, nhóm đang xây dựng một successor system vẫn giữ krablet và queue semantics của FOR UPDATE SKIP LOCKED, nhưng lưu data trên LSM tree nằm trên object storage

Định hướng tương lai

  • Với Krea 2, nhóm đã chọn architecture và optimizer tương đối thận trọng, ưu tiên tính ổn định và iteration speed
  • Trong pretraining cycle tiếp theo, nhóm muốn áp dụng modern LLM transformer design cho diffusion transformer
    • Các hướng đang xem xét gồm MoE, scale native resolution 2K–4K thông qua sparse attention, NVFP4 pretraining và Muon scaling
    • Model hiện tại còn undertrained, và nhóm cho rằng training lâu hơn sẽ có ích
  • Hiện tại, Krea 2 training pipeline kết thúc bằng multi-reward RL stage
    • Krea đã dùng expert nội bộ để xác minh rằng OPD và MOPD là các distillation method hiệu quả cho diffusion model
    • Nhóm hy vọng sớm chia sẻ thêm nhiều kết quả
  • Production diffusion model đòi hỏi một cấu hình phức tạp gồm nhiều model phụ thuộc lẫn nhau
    • Latent diffusion model serving thường cần autoencoder, diffusion transformer, text encoder và prompt-expansion model
    • Tùy stack, có thể có thêm các module như style-reference model hoặc upscaler
    • Việc duy trì nhiều component vừa phải được huấn luyện độc lập vừa có phụ thuộc lẫn nhau khiến việc phối hợp trong research team trở nên khó khăn
  • Trong pretraining cycle tiếp theo, Krea dự định đơn giản hóa architecture và hợp nhất nhiều component dưới một model duy nhất
  • Krea 2 chủ yếu tập trung vào image generation cho creative exploration, và trong tương lai muốn mở rộng capability sang robust editing, image reference và native 2K/4K generation
  • Nhóm cho rằng chỉ prompting bằng ngôn ngữ tự nhiên truyền thống không còn đủ để hỗ trợ toàn bộ phạm vi request của người dùng
    • Trong prompt của người dùng, nhóm quan sát thấy nhiều prompting style đa dạng như natural language, tag, detailed JSON, bounding box, instruction, visual guideline, Markdown, v.v.
    • Prompt expansion có thể giải quyết một phần, nhưng nhóm cho rằng việc model hiểu các prompt này một cách native cũng nên trở thành core capability

Bài viết liên quan

1 bình luận

 
GN⁺ 4 giờ trước
Ý kiến trên Hacker News

  • Họ công bố trọng số của mô hình text-to-image mới nhất, đồng thời đăng kèm một bài viết đi khá sâu vào quá trình huấn luyện
    Bài này còn đưa vào khá nhiều phần thường không được viết chi tiết, như hạ tầng dữ liệu và huấn luyện thực tế, nên có vẻ sẽ có nội dung đáng quan tâm với những ai theo dõi mảng này

    • Đây là một báo cáo kỹ thuật đồ sộ về mô hình tạo ảnh open-weight, rất ấn tượng
      Với tư cách là người đã theo dõi lĩnh vực này một thời gian, việc đọc về các thí nghiệm và nỗ lực đằng sau sản phẩm cuối cùng thực sự rất thú vị, và nếu họ công bố thêm một phần công cụ fine-tuning để cộng đồng có thể thử nghiệm thì có lẽ còn có thể đẩy xa hơn tiềm năng của mô hình
    • Tôi tò mò Krea xử lý nội dung như khiêu dâm hoặc gore thế nào
      Việc các mô hình lớn có xu hướng loại bỏ rất mạnh những loại nội dung này với lý do an toàn, ngay cả trong trường hợp hợp pháp, từ lâu đã khiến tôi thấy bức bối
    • Nhìn vào Ideogram4, Flux2, Qwen-Image, ZiT và Krea thì có thể thấy phía open-weight đang có nhiều chuyển động tích cực hơn
      Flux.1 Krea gốc đã có mặt trên trang benchmark GenAI Showdown của tôi từ tháng 7 năm ngoái, mà trong lĩnh vực này điều đó đã thấy như từ rất lâu rồi. Tôi cũng muốn test tử tế mô hình mới này

  • Tôi là Diego Rodriguez, đồng sáng lập kiêm CTO của Krea. Lần này chúng tôi công bố trọng số cùng một báo cáo kỹ thuật khá đầy đặn theo mặt bằng hiện tại của ngành
    Báo cáo đi chi tiết vào tuyển chọn dữ liệu/captioning, kiến trúc mô hình, hậu huấn luyện, pipeline học tăng cường, mở rộng prompt, tham chiếu phong cách và hạ tầng
    Trọng số được chia làm hai loại: Krea 2 Turbo là mô hình tăng tốc suy luận bằng cách chưng cất guidance và timestep, còn Krea 2 RAW là mô hình hướng đến việc hack và fine-tuning
    Tôi cho rằng cộng đồng LLM mở làm rất tốt ở chỗ công bố mô hình với nhiều kích cỡ và ở nhiều giai đoạn khác nhau của pipeline huấn luyện. Lần này chúng tôi cũng phát hành cả checkpoint ở giai đoạn huấn luyện trung gian lẫn hậu huấn luyện. Đây là điều hiếm thấy ở mảng hình ảnh/đa phương tiện nên tôi khá tự hào về việc đó
    Theo benchmark text-to-image của Artificial Analysis, chất lượng ảnh ở mức tương đương Nano Banana: https://artificialanalysis.ai/image/leaderboard/text-to-imag...
    Chúng tôi cũng gắn kèm giấy phép khá thoáng cho cá nhân và doanh nghiệp nhỏ
    Giới thiệu bản phát hành OSS: https://www.krea.ai/krea-2-open-source / Mô hình trên Huggingface: https://www.krea.ai/krea-2/huggingface / Kho GitHub: https://www.krea.ai/krea-2/github / Reddit AMA: https://www.reddit.com/r/StableDiffusion/comments/1udnm0a/we... / Báo cáo kỹ thuật: https://www.krea.ai/blog/krea-2-technical-report

  • Kết quả đã có, và đặc biệt là mô hình Turbo thực sự rất ấn tượng nếu tính đến việc nó nhanh như vậy chỉ với 8 bước
    Trong số các mô hình có thể host cục bộ, trước đây chỉ có Ideogram 4 vượt được nó, nhưng bên đó chậm hơn rất nhiều. Khác biệt là tính bằng phút so với tính bằng giây
    Nó vẫn gục trước các “model killer” quen thuộc như ngôi sao chín đỉnh, Count Rugen hay Trái Đất phẳng có quá nhiều người, nhưng nhìn chung đã thể hiện vượt tầm vóc, đạt điểm cao nhất trong nhóm mô hình có thể host cục bộ, và trên toàn bảng thì đứng ngay dưới Ideogram 4 với 6/15 bài test vượt qua
    Link GenAI chỉ so sánh các mô hình có thể host cục bộ: https://genai-showdown.specr.net/?models=fd,hd,kd,qi,f2d,zt,...

    • Đây là lần đầu tôi nghe thấy model killer cho mô hình text-to-image, buồn cười thật
      Khá thú vị khi cách test cuối cùng lại đi đến những hạng mục kỳ quặc mà cực kỳ cụ thể như thế này

  • Việc có thêm nhiều mô hình open-weight hơn là điều tốt, và tôi cũng rất thích những bài viết có chiều sâu
    Cách tiếp cận cố gắng giữ manifold đủ rộng để tạo được nhiều phong cách khác nhau cũng rất hay. Tôi thấy nó tốt hơn là chỉ tinh chỉnh cho khớp đúng vài preset phong cách
    Tuy vậy, các mô hình image-to-image/cấu trúc kiểu agent cao cấp như Nano Banana 2 hay Images 2.0 đã ra rất mạnh, nên giờ nhìn lại có cảm giác như đang “đánh trận cũ”
    Tôi khá nghi ngờ việc nhét chéo Qwen 3 VL gốc vào có thể tiến gần đến mức image-to-image như vậy hay không, trong khi image-to-image vững chắc lại cực kỳ quan trọng cho chỉnh sửa, tinh chỉnh, nhất quán nhân vật và khả năng khái quát hóa của những gì hiện đang dùng cho chuyển phong cách. Phần chuyển phong cách ở đây cũng có vẻ được giải thích chưa đủ
    Đạt đến mức đó chắc không dễ, nhưng tôi nghĩ đây rõ ràng là mặt trận tiếp theo của các mô hình hình ảnh. Có vẻ Ideogram đang xây dần theo hướng đó, nhưng ở phía open-weight thì tôi vẫn chưa thấy nhiều

    • Tôi hiểu sự hoài nghi đó, nhưng nội bộ chúng tôi dùng mô hình này còn nhiều hơn Nano Banana trong nhiều trường hợp như moodboard. Việc nó rẻ hơn NBP 4 lần cũng giúp ích
      Workflow kiểu agent tương thích với Krea 2 nên tôi không rõ lắm ý đó là gì. Nếu bạn đang nói về mô hình chỉnh sửa thì cái đó cũng đang được chuẩn bị
      Trên benchmark text-to-image thì nó cũng ở mức tương đương, bạn có thể xem link Artificial Analysis mà tôi đã đăng ở bình luận phía trên
      Cũng không thể huấn luyện lại Nano Banana hay ChatGPT để chúng hiểu thương hiệu của khách hàng, trong khi đó lại chính là điều khách hàng của chúng tôi phàn nàn liên tục. Thêm nữa, vì nó là mã nguồn mở nên việc so sánh 1:1 cũng không đơn giản
    • Mô hình này cũng hỗ trợ image-to-image, nên tôi không rõ vấn đề với Qwen 3 VL là gì
      Nhận xét rằng chuyển phong cách không được giải thích cũng khá mơ hồ. Trên trang có từ “reference” xuất hiện 11 lần, và khi tôi đọc thì thấy phần đó thực ra được đề cập khá nhiều

  • Cảm ơn Krea vì đã cho phép tải xuống trọng số mô hình, nhưng với những điều khoản như thế này trong giấy phép thì không phải mã nguồn mở: https://huggingface.co/krea/Krea-2-Raw/blob/main/LICENSE.pdf
    Chỉ được phép sử dụng thương mại nếu tổng doanh thu hằng năm của toàn công ty trong 12 tháng gần nhất dưới 1 triệu USD; nếu cao hơn thì cần giấy phép enterprise riêng
    Ngoài ra, không được dùng mô hình Krea, các bản phái sinh hoặc đầu ra theo cách vi phạm luật liên quan, hợp đồng hoặc chính sách sử dụng được phép, và khi phân phối phải triển khai bộ lọc nội dung hợp lý để phát hiện, ngăn chặn và giảm thiểu việc tạo ra nội dung bị cấm, có hại hoặc bất hợp pháp
    Cũng phải tuân theo chính sách sử dụng được phép; trang chính sách https://www.krea.ai/krea-2-use-policy còn có các điều khoản cấm lách các biện pháp an toàn, hạn chế sử dụng, bộ lọc nội dung, ghi nguồn và watermark do Krea hoặc bên phân phối triển khai

  • Có vẻ bản chuyển đổi GGUF của Turbo đã có rồi: https://huggingface.co/Abiray/Krea-2-Turbo-GGUF

  • Có một mục khá thú vị trên trang tuyển dụng
    Với ai biết Mellanox kiểu cũ từng là nơi như thế nào thì có thể sẽ hợp gu: https://jobs.ashbyhq.com/krea/ebe94024-eef6-4306-a019-10072a...

  • Là một mô hình tốt, nhưng việc dùng Qwen VAE hơi đáng tiếc một chút

    • Nếu muốn đẩy mạnh tính chân thực thì Krea 2 Large trên website và API được huấn luyện với FLUX 2 VAE
      Sau khi dùng cả hai, tôi thấy Flux VAE thực sự nhỉnh hơn một chút trong việc học texture chân thực, nhưng khác biệt không lớn như tưởng tượng. Qwen VAE cũng nhìn chung rất tốt trong các thử nghiệm có kiểm soát và mạnh ở việc học tạo ra nhiều phong cách đa dạng
    • Cũng có người nói rằng dùng wan2.1 VAE thay thế thì sẽ giải quyết được vấn đề này
      Tôi vẫn chưa có thời gian tự thử

  • Tôi khá mong được dùng thử Krea 2. Tôi dùng Z-Image Turbo hằng ngày, và nó đã thay thế gói đăng ký ảnh stock cho nhu cầu ảnh chân thực lẫn minh họa
    Tôi tò mò chi phí huấn luyện là bao nhiêu

    • Chắc chắn là tốn rất nhiều cà phê
      Khó ước tính chính xác chi phí huấn luyện vì họ dùng một cụm Kubernetes dùng chung, nơi cả suy luận lẫn workload nghiên cứu cùng chạy

  • Tôi tò mò mọi người dùng gì để tự host những mô hình như thế này
    Tôi đã thử ollama và open-webui nhưng chúng hoàn toàn không hỗ trợ tạo ảnh

    • Tôi chưa thử riêng mô hình này, nhưng ComfyUI chắc chắn sẽ hỗ trợ, và khi đã quen thì giao diện của nó cũng khá ổn
      Nếu bị khựng thì cứ bắt đầu bằng cách sao chép rồi dán workflow trước
    • Koboldcpp có hỗ trợ tạo ảnh, nhưng hỗ trợ Krea2 thì phải chờ bản phát hành tiếp theo
      https://github.com/LostRuins/koboldcpp