UniMERNet Architecture Overview

• UniMERNet is a universal architecture for mathematical expression recognition, integrating a detail-aware Swin Transformer encoder with an mBART decoder.
• It utilizes a Length-Aware Module to predict output sequence lengths, enhancing decoding stability and preventing over/under-generation.
• Trained on the UniMER-1M dataset, the model achieves state-of-the-art performance across printed, handwritten, and noisy math expressions.

UniMERNet is a universal neural network architecture for mathematical expression recognition (MER) in real-world complex scenes. Developed to leverage the large-scale UniMER-1M dataset, UniMERNet targets robust, accurate, and generalizable conversion of visual formula images into structured LaTeX token sequences by integrating a detail-aware Swin Transformer encoder, a length-aware conditioning module, and an mBART-based Transformer decoder. By incorporating both token sequence structure and high-resolution visual cues, UniMERNet achieves state-of-the-art results across noisy, handwritten, and printed mathematical expressions (Wang et al., 2024).

1. Design Principles and Motivation

MER presents unique challenges due to the dense two-dimensional visual structure, symbol ambiguity, variable sequence lengths, and the prevalence of occlusions, noise, and layout variation in real-world data. Previous models frequently struggled with over/under-generation in the sequence output, loss of spatial context, insufficient robustness against visual distortions, and poor performance on out-of-domain data.

UniMERNet addresses deficiencies in prior approaches by:

• Utilizing a multi-scale, detail-aware visual encoder (Swin Transformer) to capture precise local context and discriminative global cues.
• Explicitly predicting output sequence length via a dedicated Length-Aware Module (LAM), informing the decoding process and stabilizing output termination.
• Employing a pretrained mBART decoder with cross-attention, combining benefits of state-of-the-art language modeling and multimodal conditioning for improved syntactic accuracy and expressiveness.

2. Architectural Components

UniMERNet processes an input RGB image $I \in \mathbb{R}^{3 \times H_0 \times W_0}$ through the following pipeline:

2.1 Image Augmentation

The input undergoes strong on-the-fly augmentations including random dilation/erosion, synthetic "weather" noise, and other domain-specific perturbations. This step exposes the encoder to the diversity present in the UniMER dataset, increasing robustness to naturally occurring artifacts.

2.2 Swin Transformer Encoder

The encoder is based on Swin-Tiny (Liu et al., ICCV 2021), consisting of:

• Conv Patch-Partition: A $4 \times 4$ convolution (stride 4), converting the image into $96 \times (H_0/4) \times (W_0/4)$ features.
• Hierarchical Stages: Four sequential stages, each containing Swin Transformer blocks with increasing channel widths $[96,192,384,768]$ and interleaved 2x2 patch-merging downsamplings. The number of blocks per stage is $[2,2,18,2]$.
• Windowed Self-Attention: In each block, features are partitioned into $M \times M$ (typically $7 \times 7$) local windows for efficient multi-head self-attention. Shifted windows in half the blocks ensure cross-window information propagation.
• Normalization and Feedforward Layers: LayerNorm applied before attention and MLP sublayers, each MLP being two-layer, hidden dimension $4C_l$, using GELU activations.

Post-encoding, a feature map of $768 \times (H_0/32) \times (W_0/32)$ is produced, then flattened into $T = (H_0/32) \cdot (W_0/32)$ visual tokens $Z \in \mathbb{R}^{B \times T \times 768}$, where $B$ is the batch size.

2.3 Length-Aware Module (LAM)

LAM estimates the target LaTeX token sequence length as follows:

• Applies a single self-attention (hidden=768, heads=12) layer to $Z$.
• Aggregates token information by global average pooling, yielding $B \times 768$ vectors.
• Projects via an MLP head ($768 \rightarrow 1$) to predict scalar sequence length $\hat\ell$ for each input.
• The predicted $\hat\ell$ is linearly projected to a dense vector $e_{\text{len}} \in \mathbb{R}^{B \times 768}$, used as a global conditioning cue for the decoder.

2.4 mBART Decoder

An mBART-base decoder (12 layers, 768-dim, 12 attention heads) sequentially generates structured LaTeX symbol tokens:

• At decoding step $t$, input is the sum $$e_t = e_{\text{tok}} + e_{\text{pos}} + e_{\text{len}}$$.
• Each decoder layer performs:
  • Masked self-attention over previously generated tokens.
  • Cross-attention to encoder outputs $Z$.
  • Feedforward transformation.

Cross-attention is computed as: $$A = \mathrm{Softmax}\left(\frac{QK^\top}{\sqrt{768}}\right) V$$ where $Q$ is decoder-projected hidden state (via LayerNorm and $W^{Q}$), $K$ and $V$ are projections of encoder outputs.

The output logit distribution is computed via a final linear transform and softmax over the vocabulary.

3. Layer-by-Layer Specification

A summary of UniMERNet’s information flow is presented in the following table:

Layer	Input Shape	Output Shape
Input Image	$3 \times H_0 \times W_0$	$3 \times H_0 \times W_0$
Augmentation	$3 \times H_0 \times W_0$	$3 \times H_0 \times W_0$
Conv Patch Partition	$3 \times H_0 \times W_0$	$96 \times H_0/4 \times W_0/4$
Swin Stage 1 (2 blocks)	$96 \times H_0/4 \times W_0/4$	same size
Patch-Merge 1	$96 \times H_0/4 \times W_0/4$	$192 \times H_0/8 \times W_0/8$
Swin Stage 2 (2 blocks)	$192 \times H_0/8 \times W_0/8$	same size
Patch-Merge 2	$192 \times H_0/8 \times W_0/8$	$384 \times H_0/16 \times W_0/16$
Swin Stage 3 (18 blocks)	$384 \times H_0/16 \times W_0/16$	same size
Patch-Merge 3	$384 \times H_0/16 \times W_0/16$	$768 \times H_0/32 \times W_0/32$
Swin Stage 4 (2 blocks)	$768 \times H_0/32 \times W_0/32$	same size
Flatten	$768 \times H_0/32 \times W_0/32$	$T \times 768$
Length-Aware Module (LAM)	$B \times T \times 768$	$\hat\ell$ (scalar), $e_{\text{len}}$ (768)
mBART Decoder	$e_{\text{tok}} + e_{\text{pos}} + e_{\text{len}}$	Token logits (size $|\mathcal{V}|$)

Default parameters: batch size 64, Adam optimizer with 0.05 weight decay, 180k total training iterations. Learning rate is linearly warmed up to 1e-5 over 10k iterations, then cosine decayed from 1e-4 to 1e-8. Decoder sequence capped at 1024 tokens.

4. Loss Functions and Training Dynamics

UniMERNet is trained under composite supervision:

• Language Modeling Loss: Standard cross-entropy between predicted and ground-truth token distributions,

$\ell_{\text{lm}} = -\sum_{o=1}^{C} y_o \log p_o$

where $C$ is vocabulary size, $y_o$ is the ground truth indicator, and $p_o$ is predicted probability.

• Length Loss: Smooth-$L_1$ between predicted and true sequence lengths,

$$\ell_{\text{len}}(\hat\ell,\ell) = \begin{cases} 0.5(\hat\ell - \ell)^2, &\text{if }|\hat\ell-\ell| < 1 \ |\hat\ell-\ell| - 0.5, &\text{otherwise} \end{cases}$$

The total objective is $L = \ell_{\text{lm}} + 0.5 \ell_{\text{len}}$.

By conditioning decoding on length estimation, UniMERNet improves decoder stability, preventing "drift" (over/under-generation). Empirically, this results in 10–15% faster decoding and greater reliability.

5. Computational Complexity and Empirical Performance

For a 224×224 input, the encoder computes approximately 4.5 GFLOPs. The per-token decoder cost is $O(|\mathcal{V}| \cdot 768 + 768^2 \cdot 12)$, with total generation time notably reduced by sequence length conditioning. The LAM module introduces only ∼50 MFLOPs per image, yielding minimal overhead.

On the UniMER-1M and UniMER-Test benchmarks, UniMERNet improves both BLEU and edit-distance metrics compared to prior work, demonstrating superior performance across printed, handwritten, complex, and noisy subsets. The architecture retains end-to-end trainability and generalizes effectively to in-the-wild mathematical expression instances.

6. Significance in Mathematical Expression Recognition

UniMERNet sets a new technical baseline for large-scale, robust mathematical expression recognition. The combination of a detail-aware visual encoder, explicit length modeling, and cross-attentive language decoding systematically resolves longstanding issues such as sequence drift and sensitivity to visual artifacts.

A plausible implication is that the explicit integration of predicted sequence length into the decoding process may benefit other sequence-to-sequence vision–language applications characterized by variable output lengths and strict syntactic requirements. Additionally, UniMERNet’s robustness to visual perturbations, demonstrated via strong augmentations and tested on diverse datasets, suggests suitability for deployment in practical, high-noise environments such as classroom capture and document analysis pipelines (Wang et al., 2024).

All UniMERNet data, code, and trained models are available at https://github.com/opendatalab/UniMERNet.