OmniAlpha: Aligning Transparency-Aware Generation via Multi-Task Unified Reinforcement Learning

1. Problem Statement

Current generative AI models, particularly diffusion-based systems, excel at creating high-fidelity RGB images. However, many real-world visual workflows (e.g., visual effects, graphic design, image editing) require transparency-aware generation (RGBA), which involves modeling not just color but also per-pixel opacity (alpha channels) and cross-layer composition.

Existing solutions suffer from two main limitations:

• Fragmentation: Current methods are task-specific silos. Separate pipelines exist for image matting, object removal, layer decomposition, and text-to-image generation. There is no unified model capable of handling diverse RGBA workflows.
• Limitations of Supervised Fine-Tuning (SFT): While scaling SFT across tasks is a logical step toward unification, it relies on localized regression objectives (e.g., pixel-wise L1/L2 loss). These objectives fail to directly optimize critical properties of transparency-aware generation, such as:
  • Cross-layer consistency and physics.
  • Precision at alpha boundaries (e.g., hair, fur).
  • Compositional faithfulness when layers are blended.
  • Consequently, SFT-only models often plateau, struggling to capture fine structural details required for seamless transparency.

2. Methodology: OMNIALPHA

The authors propose OMNIALPHA, a unified multi-task framework that shifts from purely supervised learning to Reinforcement Learning (RL) alignment. The architecture consists of three core components:

A. Unified Architecture

• Alpha-Aware VAE: An end-to-end Variational Autoencoder initialized from a pre-trained RGB VAE using an "opaque initialization" strategy. The encoder zeros out alpha weights initially to focus on RGB, while the decoder is initialized to predict a fully opaque alpha channel ($\alpha=1.0$). This allows efficient adaptation to the 4-channel RGBA space.
• Sequence-to-Sequence Diffusion Transformer (DiT): The denoising backbone processes multiple inputs and outputs simultaneously.
  • Bi-directional Layer Axis: To handle multiple images (inputs and targets) in a single forward pass, the authors extend standard 2D Rotary Positional Embeddings (RoPE) with a bi-directional z-axis.
  • Input latents are assigned non-negative indices ($0$ to $n-1$), target latents are assigned negative indices ($-1$ to $-m$), and text embeddings are assigned distinct positive indices. This allows the model to differentiate layers spatially within the sequence.

B. Training Paradigm

The training follows a two-stage process:

1. Cold-Start (Multi-Task SFT): The model is first trained on a unified dataset covering five task categories using Flow Matching objectives. This establishes a baseline capability across diverse tasks.
2. Post-Training (Group Relative Policy Optimization - GRPO): To overcome SFT limitations, the authors formulate the generation process as a Markov Decision Process (MDP).
   • State: Multimodal conditions, time step, and intermediate latent sequence.
   • Action: The denoising step guided by the predicted velocity field.
   • Reward: A Layer-Aware Reward is computed only at the end of the trajectory (outcome-based).

C. Layer-Aware Reward Design

Instead of task-specific rewards, OMNIALPHA uses four universal objectives defined on the decoded RGBA sequence to guide the RL policy:

1. Layer Fidelity ($r_{layer}$): Measures RGB similarity (SSIM/PSNR) across all valid layers.
2. Composition Fidelity ($r_{comp}$): Evaluates the globally composited image against the ground truth.
3. Background Structure ($r_{bg}$): Penalizes structural deviations in background layers using LPIPS.
4. Foreground Boundary ($r_{fg}$): Enforces precision and recall on alpha boundaries (crucial for hair/fur) using soft boundary masks.

These rewards are dynamically aggregated based on the specific task (e.g., text-to-image skips pixel-aligned rewards as ground truth is not available, while matting tasks heavily weight boundary rewards).

3. Key Contributions

1. First Unified RL Framework for RGBA: OMNIALPHA is the first framework to apply multi-task reinforcement learning to unify transparency-aware generation and manipulation, moving beyond fragmented, task-specific pipelines.
2. Structural Reward Paradigm: The authors introduce structurally-grounded rewards that explicitly optimize cross-layer coherence and alpha-boundary precision, successfully transcending the performance ceiling of standard SFT.
3. Bi-directional Layer Encoding: A novel extension of positional embeddings that enables a single DiT to jointly model multiple input and output RGBA layers in a single forward pass.

4. Experimental Results

The model was evaluated on five categories of tasks: Text-to-Image, Object Removal, Automatic Matting, Referring Matting, and Layer Decomposition.

• Overall Performance: OMNIALPHA consistently outperforms its own SFT baseline and specialized expert models across diverse benchmarks.
• Key Metrics:
  • Layer Decomposition: Achieved a 9.07% relative reduction in RGB L1 compared to the SFT baseline.
  • Automatic Matting: Outperformed conventional tools with a 74% improvement in SAD (Sum of Absolute Differences) and 68% improvement in Grad (Gradient difference).
  • Referring Matting: Achieved state-of-the-art results on the RefMatte-RW100 benchmark (SAD: 14.768), significantly beating specialized models like CLIPMat.
  • Object Removal: Demonstrated strong generalization on out-of-domain datasets (OBER-Wild), achieving ReMOVE scores comparable to the highly specialized ObjectClear model.
• Ablation Studies: Removing any single reward component ($r_{layer}, r_{comp}, r_{bg}, r_{fg}$) resulted in performance degradation, confirming that the full reward suite is necessary for robust multi-task alignment.

5. Significance

OMNIALPHA represents a significant paradigm shift in generative AI for visual editing. By demonstrating that Reinforcement Learning can effectively align diffusion models for complex, multi-layer RGBA tasks, the paper bridges the gap between opaque RGB generation and the nuanced requirements of professional visual workflows.

It suggests that future foundation models for image editing should move away from isolated, task-specific fine-tuning toward unified policies optimized via outcome-based rewards. This approach not only improves the fidelity of alpha channels and layer composition but also enables a single model to handle a wide spectrum of creative tasks, from generating transparent objects to decomposing complex scenes.