klingairesearch/imba-loss

IMBA Loss is a research project from Kuaishou's Kling team that tackles a common problem in AI image and video generation: models often struggle to correctly combine multiple concepts in a single image. For example, if you ask a model to generate "a red car and a blue bicycle," it might get one object right but blur or distort the other. IMBA Loss is a new training technique that helps generation models balance different concepts during training, so the final model is better at composing complex scenes with multiple elements correctly represented. At a technical level, IMBA Loss replaces the standard training objective (called diffusion loss) with a "self-equalization" loss function. The key idea is that during training, the loss dynamically adjusts based on the data distribution being processed, so no single concept dominates the learning process. This helps the model give balanced attention to all the concepts it encounters rather than over-indexing on frequent or prominent ones. The authors show this approach works well both when training a model from scratch and when fine-tuning an existing model, and it benefits both image and video generation tasks. This tool is primarily aimed at AI researchers and engineers who train or fine-tune diffusion-based generation models. If you're building a custom image generation model and finding that it produces visually inconsistent results when prompts contain multiple objects, attributes, or styles, IMBA Loss could help improve those outcomes. The team behind it works on Kuaishou's Kling video generation product, so the technique has been validated in a large-scale production setting. One practical limitation worth noting: the actual pre-trained weights and the code trained on their proprietary 31-million-image dataset are not open-sourced due to copyright restrictions. Instead, the repository provides a reference implementation built on top of x-flux, an open-source FLUX model training framework. All the IMBA Loss modifications are contained in a single training script, making it relatively straightforward to experiment with if you're already familiar with that ecosystem.