MSRoPE-BiL Embedding

• The paper presents a novel embedding mechanism, MSRoPE-BiL, that enables single-pass, multi-layer RGBA image generation by explicitly encoding spatial and layer distinctions.
• It employs a unique mathematical formulation that extends standard RoPE by incorporating an explicit layer index, ensuring translation invariance and efficient cross-layer attention.
• Experimental results demonstrate significant improvements in concurrent multi-image inference, layer-conditioned generation, and a marked reduction in matting error.

MSRoPE-BiL (Multi-axis, bi-directional Rotary Positional Embedding with a Layer axis) is a positional embedding mechanism designed to enable transformer architectures, such as Diffusion Transformers (DiTs), to natively process multiple input and output RGBA image layers in parallel. By introducing a third, bi-directionally extendable "layer" axis into Rotary Positional Encoding (RoPE), MSRoPE-BiL empowers a single transformer model to distinctly and efficiently represent spatial and cross-layer information within a unified sequence, which is critical for flexible sequence-to-sequence RGBA image generation and editing (Yu et al., 25 Nov 2025).

1. Motivation and Conceptual Overview

Traditional RoPE mechanisms address token sequence or 2D spatial patch embedding, but are insufficient for tasks requiring explicit awareness of both spatial location and layer membership, such as multi-image or RGBA layer-aware generative modeling. In sequence-to-sequence RGBA image generation, it is imperative to distinguish:

• The image (layer) from which a token originates,
• Whether the token corresponds to an input (conditioning) or target (generative) layer,
• The 2D spatial coordinate within that layer.

MSRoPE-BiL achieves this by conceptualizing each token as a point $(x, y, z)$, where $x$ and $y$ denote spatial coordinates and $z$ the "image index" or layer, which is assigned nonnegative indices for inputs and negative indices for output layers. This explicit augmentation avoids the ambiguity and collapse in attention that would arise from treating all patches from different images or layers as indistinguishable along the sequence dimension (Yu et al., 25 Nov 2025).

2. Mathematical Formulation

A hidden token $h \in \mathbb{R}^d$ in an attention head is tagged with coordinates $p = (x, y, z_\mathrm{raw})$. The domain is: $$x, y \in \{0, \dots, H-1\}, \quad z_\mathrm{raw} \in \{-m, \dots, -1\} \cup \{0, \dots, n-1\} \cup \{n, \dots\}$$ where negative $z_\mathrm{raw}$ enumerate target layers, nonnegative identify input layers, and higher values accommodate VLM text tokens.

To permit a single frequency table, an offset $S_\mathrm{offset} \geq m$ (typically $S_\mathrm{offset} = m$) is added: $$z_\mathrm{impl} = z_\mathrm{raw} + S_\mathrm{offset}$$, making all indices nonnegative. The rotary transform $R$ is then computed across this 3D position:

• Each attention head's embedding is split into $d/2$ complex pairs.
• For pair $i$, the angular frequency is $\omega_i = 1/10000^{2i/d}$.
• The rotation angle: $\theta_i(p) = \omega_i (x + y + z_\mathrm{impl})$.
• The rotated pair: $$R(h_{2i:2i+1}, p) = \mathrm{Rot}(\theta_i(p)) \begin{pmatrix}h_{2i} \ h_{2i+1}\end{pmatrix}$$ with $\mathrm{Rot}(\theta)$ the standard $2 \times 2$ rotation matrix.

This construction yields translation invariance: shifting all $z_\mathrm{raw}$ by a constant does not affect attention, as attention scores depend solely on relative $(x, y, z)$ differences. Furthermore, assigning non-overlapping $z$ ranges to inputs and targets ensures they remain distinct within a single sequence pass (Yu et al., 25 Nov 2025).

3. Extension Beyond Standard RoPE

Standard 1D RoPE, as used in GPT-type models, encodes only a single sequence index, while "2D RoPE" in image transformers encodes $(x, y)$ through compound or successive rotations. MSRoPE-BiL generalizes this by incorporating a third, explicit $z$ axis for layer identification:

• Standard 1D RoPE: $\theta = \omega_i \cdot (\text{token index})$
• 2D RoPE: $\theta_x = \omega_i x$, $\theta_y = \omega_i y$ (summed or composed)
• MSRoPE-BiL: $\theta = \omega_i (x + y + z_\mathrm{impl})$

Efficiently, this reduces to a single scalar rotation per pair. The design ensures that (i) input and output layers never share $z$, preserving their separability, and (ii) all attention operations are translation-invariant and indexed into a unified embedding structure (Yu et al., 25 Nov 2025).

4. Integration into Diffusion Transformer Architectures

Within each block of the DiT backbone, the standard RoPE mechanism applied to $Q$ (query) and $K$ (key) projections is replaced with the MSRoPE-BiL transform. Tokens constituting the long sequence are each labeled with $(x_j, y_j, z^{\text{raw}}_j)$. The pipeline proceeds:

• Compute $$z^{\text{impl}}_j = z^{\text{raw}}_j + S_{\text{offset}}$$.
• For all head-pairs $i$, compute $$\theta^j_i = \omega_i(x_j + y_j + z^{\text{impl}}_j)$$.
• Apply 2D complex rotation to $Q_j$ and $K_j$ for each attention head’s pairs.

This process embeds all RGBA patch and text tokens into a single, undifferentiated sequence, permitting the transformer to reason jointly about all layers, spatial locations, and modalities, without reliance on hacks such as multiple forward passes or zero-padding (Yu et al., 25 Nov 2025).

5. Algorithmic Implementation

The MSRoPE-BiL pseudocode for each attention layer is as follows:

function MSRoPE_BiL(LatentTokens): # LatentTokens is list of (h, x, y, z_raw) z_offset = m # number of target layers for each token index j: (h, x, y, z_raw) = LatentTokens[j] z_impl = z_raw + z_offset # Project to Q and K q = Q_proj(h) # shape [n_heads, d_head] k = K_proj(h) # Apply rotary on each head for head in 1..n_heads: for pair i in 0..(d_head/2 -1): omega = 1.0 / (10000 ** (2*i / d_head)) theta = omega * (x + y + z_impl) c = cos(theta); s = sin(theta) q_head[2*i ] = q_head[2*i ]*c - q_head[2*i+1]*s q_head[2*i+1] = q_head[2*i ]*s + q_head[2*i+1]*c k_head[2*i ] = k_head[2*i ]*c - k_head[2*i+1]*s k_head[2*i+1] = k_head[2*i ]*s + k_head[2*i+1]*c Q[j], K[j] = q, k return Q, K

Every patch or text token is mapped with this process, preserving full information about its spatial and layer provenance (Yu et al., 25 Nov 2025).

6. Empirical Performance, Advantages, and Ablation

MSRoPE-BiL is fundamental to concurrent multi-image inference in unified multi-task RGBA models. With traditional 2D RoPE, $n+m$ forward passes or intricate sequence padding would be needed to handle all input-output images. In contrast, MSRoPE-BiL supports a single forward pass per diffusion step, with $O(1)$ wall-clock time complexity with respect to the number of layers. Experimental results from "OmniAlpha: A Sequence-to-Sequence Framework for Unified Multi-Task RGBA Generation" demonstrate:

• Concurrent multi-image inference: Single-pass processing of arbitrarily many layers, unattainable with standard RoPE.
• Layer-conditioned generation: Over 90% win rate in human preferences for layer-conditioned completion, with the gain attributed primarily to effective cross-layer attention enabled by MSRoPE-BiL.
• Significant reduction in matting error: In mask-free matting tasks on AIM-500, an 84.8% reduction in SAD over the strongest specialized baseline is observed. Removing MSRoPE-BiL increases the SAD nearly threefold (Yu et al., 25 Nov 2025).

A plausible implication is that MSRoPE-BiL is not merely a minimal extension but a critical enhancement for tasks where semantic, spatial, and layer information must be jointly modeled within a generative transformer system.

7. Context, Significance, and Future Directions

MSRoPE-BiL demonstrates that bi-directional, multi-axis rotary embeddings can facilitate fully unified, multi-task generative models for complex data modalities such as RGBA image stacks. It preserves the translation invariance and efficiency of RoPE while providing an extensible mechanism for multiplexed attention across multiple image and text channels.

Its adoption in OmniAlpha (Yu et al., 25 Nov 2025) provided empirical evidence that unified frameworks can outperform specialized ones, substantiating the utility of shared, contextualized representation learning across input-output layer boundaries. This suggests broader applicability in domains beyond RGBA, wherever joint reasoning about multiple modalities or axes of provenance is required. A plausible implication is expansion to more general multi-axis generative and sequence modeling tasks involving diverse data layers or modalities.