Multimodal Cross Quantizer (MXQ-VAE)

• The paper introduces an unconditional image–text generation framework by learning a shared quantized latent space for semantically aligned multimodal data.
• It employs a two-stage training process combining multimodal Transformers, VQ-VAE style quantization, and input masking to enhance cross-modal correlations.
• Empirical evaluations on synthetic and real-world datasets demonstrate state-of-the-art performance in semantic consistency and generative alignment.

The Multimodal Cross Quantizer Variational Autoencoder (MXQ-VAE) is a vector quantization framework for the unconditional joint generation of image–text pairs. The model introduces a quantized joint representation space for images and text by combining VQ-VAE-style discrete bottlenecks, multimodal Transformer self-attention, and an input masking strategy. Through a two-stage training and sampling procedure, MXQ-VAE serves as an end-to-end solution for generating semantically aligned multimodal data without conditioning. Quantitative experiments on synthetic and real-world datasets demonstrate that MXQ-VAE achieves state-of-the-art semantic consistency in image–text generative tasks (Lee et al., 2022).

1. Architecture and Quantization Framework

MXQ-VAE consists of two main stages. In Stage 1, the model learns a discrete joint latent space for image–text pairs. The architecture decomposes into the following components:

• Image Encoder: A lightweight 2D CNN maps input image $I \in \mathbb{R}^{H \times W \times 3}$ to feature maps $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$, with $h = H/f$, $w = W/f$, typically $f=8$, $d=128$.
• Text Encoder: A 1D CNN projects the input token sequence $T = \{t_1, ..., t_N\}$ to $z_{\text{txt}} \in \mathbb{R}^{n \times d}$, $n = N/f$.
• MXQ Module: The concatenated sequence of $h \cdot w + n$ $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$0-dimensional vectors is passed through a two-layer Transformer encoder (multi-head self-attention), producing $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$1, $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$2.
• Shared Codebook: Vector quantization is performed via a codebook $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$3, with $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$4, $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$5. Each $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$6 is quantized by hard nearest-neighbor lookup:

$$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$7

• Decoding: After quantization, the unified code sequence is split by two linear heads into reconstructed image codes $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$8 and text codes $$z_{\text{img}} \in \mathbb{R}^{h \times w \times d}$$9, which are decoded by mirror CNNs to pixels and tokens, respectively.

The vector quantization objectives combine VQ-VAE reconstruction, codebook, and commitment losses:

$h = H/f$0

with $h = H/f$1, $h = H/f$2.

2. Multimodal Input Masking and Joint Latent Correlation

During Stage 1 training, input pairs $h = H/f$3 are subjected to independent random masking with probability $h = H/f$4 in both modalities:

• Images: $h = H/f$5 of non-overlapping patches are replaced with zeros.
• Text: $h = H/f$6 of tokens are replaced by a [MASK] token.

This joint masking compels the Transformer to reconstruct the missing regions in each modality using not only the observed regions within that modality but also information from the other modality. This strategy directly strengthens cross-modal semantic correlation in the quantized code space. Ablation studies indicate that removing the input masking mechanism ("w/o IM") causes a $h = H/f$7–$h = H/f$8 point drop in semantic correlation and generation alignment metrics.

3. Two-Stage Unconditional Image–Text Pair Generation

Following the learning of the joint discrete latent space, Stage 2 models the unconditional generative process:

• Code Sequence Extraction: Each training pair is encoded into a code index sequence $h = H/f$9, $w = W/f$0.
• Autoregressive Modeling: A GPT-2-style Transformer (8 layers, 8 heads, hidden size 512) is trained autoregressively to maximize $w = W/f$1.
• Unconditional Sampling: At inference, code indices are sampled via top-$w = W/f$2 sampling ($w = W/f$3) to generate synthetic unified code sequences $w = W/f$4.
• Joint Decoding: The decoder reconstructs both image $w = W/f$5 and caption $w = W/f$6 from the latent codes, with no conditioning on real data.

4. Empirical Evaluation and Performance

MXQ-VAE was evaluated on synthetic and natural image–text datasets:

• Caption MNIST: 600k synthetic $w = W/f$7 images, captions with fixed vocabulary. Used for precise cross-modal consistency measurement.
• Oxford-102 Flowers & CUB-200 Birds: Real images ($w = W/f$8), 10 captions per image, 64-token fixed-length truncation.
• COCO: 120k images, 5 captions each.
• Degree Variants: Caption corruption (color/digit swaps) for controlled alignment testing.

Key findings include:

Dataset/Metric	Baseline	MXQ-VAE (ours)	Performance Gain
Caption MNIST Gen Align	0.945	0.992	+0.047
Flowers Prec. / BERT Sim	0.407 / 0.937	0.428 / 0.941	+0.021 / +0.004
CUB Prec. / BERT Sim	0.422 / 0.941	0.478 / 0.948	+0.056 / +0.007
COCO CLIP P@1/5/10	0.083/0.290/0.425	0.106/0.323/0.491	+0.023/+0.033/+0.066

On synthetically controlled settings (Degree variants on Caption MNIST Quad3), MXQ-VAE achieves near-ideal "stair-step" accuracy: $w = W/f$9, showing semantically precise joint representation, in contrast to ablations and baselines.

5. Ablation Studies and Architectural Insights

Ablation studies elucidate the necessity of key components:

• Input Masking (IM): Essential for cross-modal correlation. Ablating IM reduces alignment by 15–20 points.
• Text-Compression (TC): Replacing the 1D CNN with a full Transformer text encoder (w/o TC) results in almost no cross-modal binding and total alignment collapse.
• Codebook Sharing Only: A codebook without attention ("Only-Sharing C") does not induce sufficient multimodal interaction.
• Hyperparameters: Best convergence and code utilization are found at $f=8$0, $f=8$1. Larger values slow convergence and under-utilize the codebook.
• Transformer Depth: One stack is insufficient; two saturate performance; four yield minimal gains and double computational cost.
• Representation Visualization: t-SNE plots show well-separated clusters in the joint latent space. Cross-attention maps indicate exact word–patch correspondence, even without explicit cross-modal supervision.

6. Training Regimen and Hyperparameter Configuration

The MXQ-VAE and its associated sampling Transformer are trained with the following settings:

• Codebook size: $f=8$2 vectors, embedding dimension $f=8$3.
• Downsampling: $f=8$4, so $f=8$5 images $f=8$6 $f=8$7 codes; 64-token text $f=8$8 8 codes.
• Masking during Stage 1: $f=8$9 per patch/token.
• Loss weights: $d=128$0; $d=128$1.
• Optimizers: AdamW ($d=128$2, $d=128$3), cosine-decay learning rate schedule; learning rate $d=128$4; batch size 800.
• Autoregressive transformer (Stage 2): Same optimization; top-$d=128$5 for sampling.

7. Significance and Implications

MXQ-VAE provides the first unconditional, end-to-end model for simultaneous image–text pair generation by leveraging a shared quantized code space, multimodal attention, and a strong masking strategy. This architecture yields state-of-the-art semantic consistency on both synthetic and real multimodal datasets—demonstrated quantitatively and by qualitative visualizations (Lee et al., 2022). A plausible implication is that quantized joint representations, when learned with cross-modal masking and joint attention, enable not only accurate reconstruction but also robust unconditional sampling in multimodal domains. The public availability of code and data provides a foundation for replication and further research into unconditional generative modeling with strong semantic alignment.