Invertible Cross-Attention (ICA)

• Invertible Cross-Attention (ICA) is a mechanism that blends Transformer-style expressiveness with bijective mappings, ensuring exact invertibility and tractable Jacobians.
• It employs structured, masked cross-attention and autoregressive partitioning to guarantee a block-upper-triangular Jacobian, enabling maximum-likelihood training in normalizing flows.
• Empirical results on multimodal tasks demonstrate its effectiveness in enhancing cross-modal signal propagation, density modeling, and overall computational efficiency.

Invertible Cross-Attention (ICA) is a mechanism for multimodal data fusion designed to reconcile the representational expressiveness of Transformer-style attention with the bijective, tractable Jacobian properties demanded by normalizing flow (NF) models. The ICA layer provides explicit, interpretable, and exactly invertible mappings between input modalities or tokens while retaining the ability to learn complex, long-range cross-modal relationships, a property that standard affine-coupling layers in flows generally lack. ICA achieves this via a structured, masked cross-attention operation that ensures tractability, invertibility, and suitability for likelihood-based maximum-likelihood learning in the NF setting (Truong et al., 13 Aug 2025).

1. Theoretical Motivation

Normalizing flows, to be suitable for generative modeling, require every transformation layer to be bijective with a tractable computation of the Jacobian determinant. Traditional affine-coupling layers (e.g., as used in RealNVP and Glow) satisfy these requirements but are limited in expressivity, as they only model simple dependencies: one subset of features is affinely transformed conditioned on the other, without capturing complex, cross-feature correlations.

Transformer-style cross-attention modules, in contrast, capture rich, pairwise, and higher-order dependencies over tokens or modalities but lose bijectivity and tractable log-determinants; their global attention leads to dense and non-triangular Jacobians, precluding exact likelihood computation and reversible inference.

ICA bridges this gap by restructuring attention computation through autoregressive masking and specific partitioning of the tokens/modalities. This constrains the attention matrix to be block-upper-triangular, guaranteeing invertibility, a tractable (closed-form) Jacobian determinant, and strong information exchange between modalities (Truong et al., 13 Aug 2025).

2. Mathematical Formulation

Given an input sequence

$$X = [x_1, \ldots, x_N]^T \in \mathbb{R}^{N \times d_{in}},$$

partitioned into disjoint groups $X_1 \in \mathbb{R}^{n \times d_{in}}$ and $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$, ICA transforms the data as follows:

Forward Pass (Coupling via Cross-Attention):

1. $Q = \mathrm{LN}(W_Q X_1)$, $K = \mathrm{LN}(W_K X_1)$, each $\in \mathbb{R}^{n \times d}$.
2. $V = X_2 \in \mathbb{R}^{(N-n) \times d}$.
3. $$A = \mathrm{softmax}((Q K^T)/\sqrt{d} \odot M) \in \mathbb{R}^{n \times (N-n)}$$, where $M$ is a fixed upper-triangular mask enforcing autoregression.
4. $Y_1 = X_1$ (unchanged).
5. $X_1 \in \mathbb{R}^{n \times d_{in}}$0.
6. $X_1 \in \mathbb{R}^{n \times d_{in}}$1.

Inverse Pass:

Given $X_1 \in \mathbb{R}^{n \times d_{in}}$2,

1. Recompute $X_1 \in \mathbb{R}^{n \times d_{in}}$3 using $X_1 \in \mathbb{R}^{n \times d_{in}}$4.
2. $X_1 \in \mathbb{R}^{n \times d_{in}}$5.
3. $X_1 \in \mathbb{R}^{n \times d_{in}}$6.
4. $X_1 \in \mathbb{R}^{n \times d_{in}}$7.

Jacobian Determinant:

The block-triangular structure yields

$X_1 \in \mathbb{R}^{n \times d_{in}}$8

where $X_1 \in \mathbb{R}^{n \times d_{in}}$9 (diagonal entries ensured positive by softmax). Consequently,

$X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$0

3. Specialized Cross-Attention Mechanisms

ICA employs multiple specialized partitioning schemes to tailor cross-attention for different multi-modal interactions:

Mechanism	Partitioning Strategy	Main Objective
MMCA	$X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$1 tokens of modality A, $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$2 tokens of modality B (or reversed)	Directs modality A's (or B's) full information into the other modality
IMCA	Each modality split in halves; $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$3 is one half from each, $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$4 the other halves	Enables cross-modality attention between different quarters
LICA	Input tokens first permuted by a learnable permutation, then split in half	Learns optimal grouping and order for cross-attention coupling

The MMCA mechanism allows for complete transfer from one modality to another. IMCA ensures that each quarter of both modalities can attend into every other, providing combinatorially enhanced fusion. LICA introduces permutation parameters (with LU decomposition to ensure invertibility) to make the splitting scheme itself learnable, maximizing cross-modality interaction flexibility (Truong et al., 13 Aug 2025).

4. Architectural and Implementation Details

All ICA layers use shared weight matrices $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$5, $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$6 for projections and employ LayerNorm before attention computation. Embedding dimensions are set to a fixed $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$7 (e.g., $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$8 in primary experiments), with $X_2 \in \mathbb{R}^{(N-n) \times d_{in}}$9 optionally as a learnable temperature. Each normalizing flow block $Q = \mathrm{LN}(W_Q X_1)$0 stacks eight ICA layers in the order: 2×MMCA, 4×IMCA (covering all partition combinations), and 2×LICA, followed by a standard affine-coupling layer for expressiveness enhancement. Partition sizes (i.e., $Q = \mathrm{LN}(W_Q X_1)$1 vs. $Q = \mathrm{LN}(W_Q X_1)$2) vary according to the mechanism.

To increase expressiveness further, standard NF design elements—such as RealNVP-style coupling and invertible $Q = \mathrm{LN}(W_Q X_1)$3 convolutions as in Glow—are interleaved with ICA blocks. High-dimensional data (e.g., $Q = \mathrm{LN}(W_Q X_1)$4 RGB plus depth) is first compressed via a perceptual autoencoder: the flow models the distribution of the compact latent, substantially reducing computational cost (Truong et al., 13 Aug 2025).

5. Invertibility and Tractability

ICA’s design guarantees full invertibility owing to the block-upper-triangular masking of the attention matrix $Q = \mathrm{LN}(W_Q X_1)$5 and the strictly positive diagonal (from softmax). The invertible mapping allows both analysis ($Q = \mathrm{LN}(W_Q X_1)$6) and synthesis ($Q = \mathrm{LN}(W_Q X_1)$7) without approximation, supporting interpretability, exact latent manipulation, and fully reversible fusion.

The closed-form log-determinant from the Jacobian yields exact log-likelihoods, supporting maximum-likelihood training. This lets models perform not just discriminative tasks but also density modeling and conditional synthesis. The explicit coupling via masked attention retains information exchange surpassing traditional coupling layers (Truong et al., 13 Aug 2025).

6. Comparison with Standard Cross-Attention

Standard Transformer cross-attention applies a global, unmasked softmax, yielding a dense, non-triangular matrix with no guaranteed invertibility or tractable Jacobian. Consequently, such architectures cannot be used for exact likelihood-based generative modeling and rely on discriminative or ELBO-style objectives.

ICA’s advantages include:

• Exact maximum-likelihood learning via tractable Jacobian
• Full invertibility and direct density modeling
• Explicit, transparent token coupling through attention

Trade-offs involve reduced parallelism from autoregressive masking (tokens attend only “forward” within a block) and potentially under-expressed bidirectional dependencies, which alternating partitions and bi-directional scheduling help mitigate. The costly part of computing $Q = \mathrm{LN}(W_Q X_1)$8 in the backward pass is a modest additional computation (Truong et al., 13 Aug 2025).

7. Empirical Performance and Practical Considerations

On benchmark multimodal tasks—including NYUDv2 and SUN RGB-D semantic segmentation, Taskonomy multimodal translation, and MM-IMDB genre classification—the MANGO framework (with stacked ICA layers and coupling layers) achieves state-of-the-art performance. For instance, stacking $Q = \mathrm{LN}(W_Q X_1)$9 ICA blocks (approx. 72 ICA layers plus 12 coupling layers) yields strong metrics (72.9M parameters, 152 GFLOPS, 144ms/inference), outstripping competitors such as GeminiFusion and TokenFusion. Adding depth ($K = \mathrm{LN}(W_K X_1)$0 blocks) further improves mIoU from $K = \mathrm{LN}(W_K X_1)$1 to $K = \mathrm{LN}(W_K X_1)$2 on NYUDv2 with modest cost increase.

Attention visualizations confirm that ICA layers successfully propagate task-relevant inter-modality signals (e.g., depth edges influencing RGB reconstructions), substantiating their role in learning meaningful multimodal correlations (Truong et al., 13 Aug 2025).