UniMERNet: Unified Math Expression Recognition

• UniMERNet is a unified, end-to-end network for mathematical expression recognition that integrates a hierarchical Swin Transformer encoder with an mBART-based decoder.
• It employs a dedicated Length-Aware Module to predict global sequence lengths, enabling accurate decoding of diverse and complex LaTeX expressions.
• Empirical results demonstrate that UniMERNet outperforms previous models on multiple benchmarks while supporting real-time operation.

UniMERNet is a unified, end-to-end network for mathematical expression recognition (MER) in real-world scenarios, introduced alongside the UniMER-1M and UniMER-Test datasets to address the challenges of formula recognition in diverse, unconstrained environments. Designed to effectively handle printed, complex, screen-captured, and handwritten formulas with broad variability in length, noise, and layout, UniMERNet couples a hierarchical Swin Transformer visual encoder with an mBART-based sequence decoder, enhanced by a dedicated Length-Aware Module for global sequence length prediction. This architecture achieves strong generalization across all known MER challenges, outperforming previous models on multiple public and newly introduced benchmarks (Wang et al., 2024).

1. Motivation and Problem Setting

The principal objective of UniMERNet is universal, robust recognition of mathematical expressions from heterogeneous image sources. Conventional datasets and models exhibit strong domain overfitting: IM2LATEX-100K and Pix2tex target only short, clean printed LaTeX; CROHME and HME100K are limited to handwriting. However, real-world scenarios present a mixture of:

• Wide formula-length distributions (from isolated symbols to multi-line, multi-thousand-token expressions)
• Imaging artifacts (e.g., blur, compression, lighting, screen backgrounds)
• Complex spatial layouts reflecting true mathematical syntax

To holistically address this spectrum, the UniMER-1M dataset provides a large-scale, diverse training set across four categories: Simple Printed (SPE), Complex Printed (CPE), Screen-Captures (SCE), and Handwritten Expressions (HWE). UniMER-Test incorporates a balanced evaluation across these modes, establishing comprehensive benchmarks for assessing model generalization and robustness (Wang et al., 2024).

2. Encoder Structure and Visual Representation

UniMERNet employs a hierarchical Swin Transformer backbone as the visual encoder, optimized for multi-scale representation.

• Input images $I \in \mathbb{R}^{3 \times H_0 \times W_0}$ are partitioned into non-overlapping patches (4×4 or 8×8).
• Four sequential Swin stages progressively decrease spatial resolution by a factor of 2 and double the feature dimension at each stage, yielding feature maps $F_1, F_2, F_3, F_4$.
• Local context is captured in each stage via window-based multi-head self-attention:

$$\text{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} + B \right)V$$

where $B$ is a learnable relative position bias.

• Alternating blocks use shifted windows (offset by $(\lfloor M/2 \rfloor, \lfloor M/2 \rfloor)$) to propagate information globally across windows, allowing comprehensive spatial modeling.
• The last encoder output $F_4 \in \mathbb{R}^{H_4 \times W_4 \times D}$ is flattened into a sequence $Z \in \mathbb{R}^{T \times D}$, with $T = H_0 W_0 / s^2$, capturing both fine syntactic details and global layout context.

3. Decoder Design and Length-Aware Module

The decoder of UniMERNet is based on the pre-trained sequence-to-sequence Transformer mBART, adapted for the specifics of long and variable-length LaTeX generation.

Length-Aware Module (LAM):

• From the encoder sequence $Z \in \mathbb{R}^{B \times T \times D}$, self-attention pools global context:

$A = \mathrm{softmax}((ZW_Q)(ZW_K)^T / \sqrt{d})$

$F_1, F_2, F_3, F_4$0

$F_1, F_2, F_3, F_4$1

• An MLP predicts the target LaTeX length $F_1, F_2, F_3, F_4$2.
• The decoder input at each step $F_1, F_2, F_3, F_4$3 is:

$F_1, F_2, F_3, F_4$4

combining previous token, position, and explicit length embedding.

• At each layer, standard Transformer operations and cross-attention over $F_1, F_2, F_3, F_4$5 are used.

Loss Function:

$F_1, F_2, F_3, F_4$6

where $F_1, F_2, F_3, F_4$7 is cross-entropy over next-token prediction, and $F_1, F_2, F_3, F_4$8 is a regression loss for sequence length:

$F_1, F_2, F_3, F_4$9

Default coefficients are $$\text{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} + B \right)V$$0, $$\text{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} + B \right)V$$1.

4. Dataset Scale and Subset Characteristics

The UniMER-1M training set contains 1,061,791 LaTeX–image pairs:

• Simple Printed (SPE): 725,246
• Complex Printed (CPE): 110,332
• Screen-Capture (SCE): ∼4,700
• Handwritten (HWE): 83,338

The test split, UniMER-Test, is equally divided into four subsets (SPE, CPE, SCE, HWE), totaling 23,789 cases. Maximum LaTeX length is 7,037 tokens; average length is 79.48 tokens. This comprehensive span ensures the evaluation of both local symbol decoding and long-range spatial parsing performance across the entire real-world complexity spectrum.

5. Empirical Results and Benchmarking

Printed and Complex Expressions

On UniMER-Test, primary metrics are BLEU (higher is better) for n-gram overlap and normalized edit distance (ED, lower is better):

Method	SPE (BLEU/ED)	CPE (BLEU/ED)	SCE (BLEU/ED)	HWE (BLEU/ED)
Pix2tex	0.873 / 0.088	0.655 / 0.408	0.092 / 0.817	0.012 / 0.920
Texify	0.906 / 0.061	0.690 / 0.230	0.420 / 0.390	0.341 / 0.522
UniMERNet	0.917 / 0.058	0.916 / 0.060	0.616 / 0.229	0.921 / 0.055

UniMERNet achieves substantial improvements on complex (CPE, +0.226 BLEU) and screen-captured (SCE, +0.196) expressions over Pix2tex, and closes the gap on all handwritten variants.

Handwritten Expressions

Measured by ExpRate (exact LaTeX match):

• CROHME2014: prior best 65.9% (CAN-ABM); UniMERNet 67.4%
• CROHME2016: prior best 63.1%; UniMERNet 68.4%
• CROHME2019: prior best 64.5%; UniMERNet 65.4%
• HME100K: prior best 68.1%; UniMERNet 68.0%

Throughput

A single forward pass using Swin+LAM+mBART for a 512×512 image executes in under 50 ms on an NVIDIA A100, supporting real-time operation (>20 fps). This suggests applicability for live document parsing in practical systems.

6. Ablations: Data, Augmentation, and Length Module

Experiments demonstrate:

• Exclusive training on Pix2tex data leads to severe domain overfitting; performance drops severely outside it (e.g., HWE BLEU 0.087).
• Combined datasets (UniMER-1M) enable robust generalization across all subsets.
• Aggressive augmentations (dilation, erosion, fog, snow, shadow) yield +2.5% SCE BLEU for Pix2tex; 0.559 → 0.616 BLEU improvement on UniMER-1M.
• LAM is essential for long-sequence prediction: adding LAM increases CPE BLEU from 0.893 to 0.916 (+2.3%), confirming its role in preventing length errors for both short and long expressions.

7. Architectural Contributions and State-of-the-Art Advancement

UniMERNet integrates several critical advances:

• Hierarchical Swin Transformer mechanism yields high-fidelity extraction of both stroke-level detail and document-level context.
• LAM provides explicit guidance on sequence termination, rarely addressed in prior printed MER systems.
• The single-model, multi-domain paradigm eliminates the need for specialized pipelines and reduces maintenance complexity.
• The scale and diversity of training data, coupled with augmentation, suppresses overfitting and improves robustness.
• Overall, UniMERNet sets new SOTA on all MER tasks—printed, screen-captured, handwritten—establishing a universal, efficient benchmark for the field (Wang et al., 2024).