DocFormer: End-to-End Transformer for Document Understanding

In plain words

A scanned form or a receipt is hard for software because its meaning depends on three things at once: the words, where they sit on the page, and how the page looks (boxes, lines, fonts). DocFormer is a single neural network that reads all three together. Every word gets its text, its position and the part of the page image it sits on. Instead of merging these signals only at the end, the network lets them influence each other inside every layer.

Training has two stages. First the model learns from a large set of unlabeled pages: it fills in hidden words, redraws the page image from its internal state, and judges whether a piece of text belongs to the page. Then these training heads are removed and the whole model is fine-tuned on a smaller labeled set for one task, such as labeling form fields or classifying a document.

TL;DR

DocFormer is an encoder-only transformer for visual document understanding. It takes OCR words, their boxes and the page image, and its multi-modal self-attention adds spatial terms whose query and key weights are shared between the text and visual branches §3.1 eq. 6 eq. 7. It is pre-trained on a 5M-page subset of IIT-CDIP with three tasks, then fine-tuned end to end §3 §3.2. Author-reported: the base model reaches 83.34 F1 on FUNSD tab. 1 and 96.17% accuracy on RVL-CDIP tab. 2. No official code was found, so an implementation has to be written from this card.

Problem

Text-only encoders ignore layout and appearance. Text-plus-spatial models such as LayoutLM and BROS ignore the pixels §2. Joint multi-modal transformers concatenate image and text tokens into one sequence, which makes self-attention hard because of cross-modality feature correlation: a word has to be tied to an arbitrary span of pixels §1 fig. 2. Two-stream designs fuse only at the end §3 fig. 2. Strong prior document models also depend on a pre-trained object detector or custom OCR §1.

The method assumes OCR text with word-level bounding boxes for every page §3. The authors state that pre-training matters most in low-to-medium data regimes (FUNSD has 149 training documents; CORD has 800) §4 tab. 6.

Method

Overview: each page becomes three inputs, a visual grid, a word-piece text sequence and per-word spatial features. Each encoder layer runs two attention maps, one for the visual branch and one for the text branch, and adds them. Visual and spatial features are passed to every layer, not only the first §3.1 §3.

1. Visual features §3.1. Pass the page image through a ResNet50 whose weights are trained from scratch as part of the model §1 §3.1. Take the layer-4 map of shape (c=2048, h/32, w/32), apply a 1×1 convolution to reduce c to d=768, flatten the spatial grid, and apply a linear map to N=512 tokens. The paper writes this as V̄ = linear(conv1×1(f_cnn(θ, v))). The input resolution is not stated.

2. Text features §3.1. Tokenize the OCR text with a word-piece tokenizer. The sequence is [CLS] followed by up to 511 tokens; tokens beyond that are dropped. Pad short pages with [PAD] and exclude padding from attention. The embedding layer W_t is initialized from LayoutLMv1 weights, following LayoutLMv2 §3.1.

3. Spatial features §3.1 eq. 1 eq. 2. For each word, use the top-left and bottom-right coordinates (x1, x3, y1, y3), width w and height h. Add relative features A_rel: differences between each word's corner and centre coordinates and those of the next word, and an absolute 1D position embedding P_abs. Build two separate spatial embeddings, one for vision (V̄_s, eq. 1) and one for text (T̄_s, eq. 2). Each is a sum of separate linear maps, and all are trainable.

4. Multi-modal self-attention §3.1 eq. 3 eq. 4 eq. 5 eq. 6 eq. 7. Start from standard scaled dot-product attention (eq. 3 is the weighted sum over values; eq. 4 is the scaled score). For each branch, the score has four terms: the content term, a 1D relative query term and a 1D relative key term that use a learned embedding a_ij of the token offset, and a spatial term. The offset embedding is clipped to a span of 8 tokens on each side [supp. §6.1]. The spatial term uses query and key matrices W_s^Q and W_s^K that are shared by the visual and text branches, so the two branches learn cross-modal spatial interactions §3.1 [supp. §6.5].

5. Combine and carry forward §3.1. The visual branch output V̂_l and the text branch output T̂_l are summed: M_l = V̂_l + T̂_l. The final M_12 feeds the downstream linear layers. The paper says features are multi-modal from layer 2 onward (see Pitfalls for what this means for the inputs).

6. Pre-training data §3. A subset of 5 million pages from IIT-CDIP. OCR is extracted from each page; the OCR engine is not named §3.

7. Pre-training tasks and loss §3.2 fig. 3.

   • Multi-modal masked language modeling (MM-MLM). Mask words as in BERT and predict them from the fused features with cross-entropy. Unlike LayoutLMv2 and LayoutLM, the image regions under masked words are not masked, so vision can help the text reconstruction §3.2.
   • Learn to reconstruct (LTR). A shallow decoder predicts the input page image from the fused features; the loss is smooth-L1. The decoder architecture is not stated §3.2.
   • Text describes image (TDI). Pool the fused features with a linear layer to a binary output. In each batch, 80% of pairs match and 20% are mismatched, and the loss is binary cross-entropy. The LTR loss is ignored for mismatched pairs §3.2.
   • Total loss: L_pt = λ·L_MM-MLM + β·L_LTR + γ·L_TDI with λ=5, β=1, γ=5. Train for 5 epochs, then remove the three heads §3.2.

8. Downstream fine-tuning §3.2 [supp. §6.4]. Add one linear head sized to the task (the supplement describes a pooling block for classification; see Pitfalls). Fine-tune all components, including the visual branch §3.2 [supp. §6.4].

# Multi-modal self-attention, one layer, one head  [eq. 3] [eq. 6] [eq. 7]
# Input:  text, visual: (N, d); spatial_text, spatial_visual: (N, d)
# Output: fused: (N, d)

def multimodal_self_attention(text, visual, spatial_text, spatial_visual):
    # Visual scores: content, relative offsets and spatial terms  [eq. 6]
    # Relative offsets are clipped to 8 tokens each side  [supp. §6.1]
    visual_scores = (
        content_scores(visual)
        + relative_query_scores(visual)
        + relative_key_scores(visual)
        + spatial_scores(spatial_visual)   # W_s^Q, W_s^K shared  [§3.1]
    )
    # Text scores: same form with text inputs  [eq. 7]
    text_scores = (
        content_scores(text)
        + relative_query_scores(text)
        + relative_key_scores(text)
        + spatial_scores(spatial_text)     # same W_s^Q, W_s^K  [§3.1]
    )
    # Each branch takes softmax over tokens and mixes its own values  [eq. 3]
    # The scale for the logits is not stated in eq. 6 or 7  [eq. 4]
    visual_context = softmax(visual_scores) @ visual_value(visual)
    text_context = softmax(text_scores) @ text_value(text)
    # Output is the sum of the two branch outputs  [§3.1]
    return visual_context + text_context

# Pre-training objective  [§3.2]
# Input:  batch of pages: OCR words, boxes, page image, is_matched flag
# Output: scalar loss

def pretraining_loss(batch):
    # Mask words; keep their image regions visible  [§3.2]
    masked_text, word_targets = mask_words(batch.text)
    fused = docformer(masked_text, batch.image, batch.spatial)
    # Predict hidden words from fused features  [§3.2]
    loss_mm_mlm = cross_entropy(word_head(fused), word_targets)
    # Redraw the page image; smooth-L1  [§3.2]
    # Decoder architecture is not stated
    loss_ltr = smooth_l1(image_decoder(fused), batch.image)
    # Pool fused features and classify matched vs mismatched pairs  [§3.2]
    # Pooling choice is not stated
    loss_tdi = binary_cross_entropy(
        pair_head(pool(fused)), batch.is_matched)
    # LTR is ignored for mismatched pairs  [§3.2]
    loss_ltr = loss_ltr * batch.is_matched
    # Weights lambda=5, beta=1, gamma=5  [§3.2]
    return 5 * loss_mm_mlm + 1 * loss_ltr + 5 * loss_tdi

Equations, as written in the paper (cited by number; the paper's own notation):

\overline{V_{s}} = W_{v}^{x}(x_1,x_3,w,A_{rel}^x) + W_{v}^{y}(y_1,y_3,h,A_{rel}^y) + P_{v}^{abs} eq. 1

\overline{T_{s}} = W_{t}^{x}(x_1,x_3,w,A_{rel}^x) + W_{t}^{y}(y_1,y_3,h,A_{rel}^y) + P_{t}^{abs} eq. 2

\overline{M}_{i}^{l}=\sum_{j=1}^{L} \frac{\exp(\alpha_{ij})}{\sum_{j'=1}^{n}\exp(\alpha_{ij'})}\left(x_{j}^{l} W^{V,l}\right) eq. 3

\alpha_{ij}=\frac{1}{\sqrt{d}}\left(x_{i}^{l} W^{Q,l}\right)\left(x_{j}^{l} W^{K,l}\right)^{T} eq. 4

\alpha_{ij}=\left(x_{i} W^{Q}\right)\cdot\left(x_{j} W^{K}\right)^{T} eq. 5, the simplified single-layer form used before the modification

\alpha_{ij}^{v}= (x_{i}^{v} W^{Q}_{v})(x_{j}^{v} W^{K}_{v})^{T} + (x_{i}^{v} W^{Q}_{v} a_{ij}) + (x_{j}^{v} W^{K}_{v} a_{ij}) + (\overline{V_{s}}W_s^{Q})(\overline{V_{s}}W_s^{K}) eq. 6

\alpha_{ij}^{t}= (x_{i} W^{Q}_{t})(x_{j} W^{K}_{t})^{T} + (x_{i} W^{Q}_{t} a_{ij}) + (x_{j} W^{K}_{t} a_{ij}) + (\overline{T_{s}}W_s^{Q})(\overline{T_{s}}W_s^{K}) eq. 7

The pre-training loss is written inline, not as a numbered equation: L_pt = λ L_MM-MLM + β L_LTR + γ L_TDI §3.2.

Figures are described in words and cited; none are reproduced. fig. 1 shows a sample page with key–value pairs and other text. fig. 2 compares four multi-modal encoder designs (joint, two-stream, single-stream, and the paper's discrete design). fig. 3 is the pre-training overview. fig. 4 plots FUNSD F1 against the number of pre-training pages, with the base model at 83.34 after 5M pages. fig. 5 shows FUNSD predictions for the full model and a text-plus-spatial baseline.

Hyperparameters

Results

Author-reported. Each downstream model is fine-tuned on the training split and scored on the test or validation split named in the table; the authors state no dataset-specific tuning and use official annotations unless noted §4. Ablations use 1M pre-training pages for 5 epochs, not 5M §4.5.

FUNSD, entity-level P/R/F1, test set tab. 1:

The DocFormer (T+S) rows are pre-trained with image, text and spatial inputs and fine-tuned without the image §4.1.

RVL-CDIP, test accuracy, 16 classes tab. 2:

† encoder + decoder. * estimated by the authors. Text and OCR for RVL-CDIP come from Textract §4.2.

CORD, entity-level P/R/F1, test set tab. 3:

Kleister-NDA, F1 on the validation split (no test ground truth), tab. 4:

Ablations (pre-training on 1M pages for 5 epochs, so these are not the headline numbers):

• Shared vs separate spatial weights tab. 5: shared, 183M params, FUNSD 76.9, CORD 93.36; separate, 198M params, FUNSD 75.58, CORD 92.51.
• Pre-training from scratch vs with pre-training tab. 6 (100 epochs without pre-training): FUNSD 83.34 vs 4.18; CORD 96.33 vs 0.54; RVL-CDIP 96.17 vs 93.95.
• Deeper head (fc → ReLU → LayerNorm → fc) vs linear head tab. 7: FUNSD 83.34 vs 82.93; CORD 96.33 vs 96.87; RVL-CDIP 96.17 vs 96.85.
• Pre-training tasks tab. 8, FUNSD / CORD F1: MLM only 72.40 / 90.58; + MM-MLM 73.91 / 90.98; + LTR 74.68 / 92.61; + TDI (full) 76.90 / 93.36.
• Architecture tab. 9, FUNSD / CORD F1: text-only BERT-base 61.56 / 89.23; + spatial 73.01 / 92.28; + multi-modal self-attention 76.90 / 93.36.

Paper vs code

No official code.

What was searched: the paper, its abstract and the arXiv metadata have no code link or project page. The paper says "Code in supple" §3.1, but the supplement contains only pseudocode [supp. §6.3]. A GitHub repository search for "docformer" returned 10 repositories, none owned by an author. shabie/docformer (MIT) is a third-party implementation; uakarsh/docformerv2 is a later version by another owner. Neither is used.

Pitfalls

• Spatial features. The text describes the relative features as "the Euclidean distance" from each corner to the corresponding corner of the next box §3.1, but the definition that follows is signed coordinate differences A_{k+1} − A_k for x and y over corners 1–4 and the centre §3.1. Implement the signed differences as written in the equation and confirm the choice against the authors.
• Which box corners. The box is given as eight coordinates (x1…y4) but only x1, x3, y1, y3 enter the spatial embedding. Corner order is assumed to be top-left and bottom-right §3.1. Coordinate normalization is not stated.
• Separate vs shared spatial inputs. Eqs. 1–2 build separate spatial embeddings V̄_s and T̄_s, and eqs. 6–7 apply the shared W_s^Q and W_s^K to each one. The supplement's pseudocode passes a single spatial_feat to both branches [supp. §6.3], which drops the separate inputs. Follow the equations.
• Relative offset terms. The shape of a_ij is not stated. Eqs. 6–7 use the same a_ij in both branches, but the pseudocode keeps separate rel_pos_embed1 and rel_pos_embed2 tables [supp. §6.3]. The behaviour for offsets beyond ±8 is not stated; the paper says only that the attention is "clipped".
• Which inputs feed each layer. The paper says the output of layer l is M_l = V̂_l + T̂_l, and that for l > 1 the inputs x in eq. 7 are multi-modal §3.1. It does not say whether the visual branch at layer l takes M_(l-1) or its own visual stream. The pseudocode keeps the two streams separate by name [supp. §6.3]. Settle this with the authors before implementing; the two readings give different models.
• Scaling and heads. Eq. 4 scales by 1/√d, with d the hidden size. The pseudocode scales each head by √d_k [supp. §6.3]. Eqs. 3–7 are single-head, while the base model has 12 heads §4. Eq. 3 sums over j = 1…L and normalizes over j' = 1…n, which mixes two lengths.
• Pseudocode is not runnable as written. The relative-key term is computed with index order "bhlr" but added to scores with order "hblt" [supp. §6.3]. The pseudocode also omits dropout, layer norm, residuals and the feed-forward block, which the paper says are omitted "for brevity" [supp. §6.3].
• Input size and truncation. The linear map from h/32 × w/32 to N=512 fixes the grid size, so the input resolution must be constant. The resolution is not stated. Text beyond 511 word pieces is dropped §3.1, so long pages lose content, and the paper does not measure the effect.
• Pre-training heads. The LTR decoder architecture, the TDI pooling and the downstream loss functions are not stated. The masking ratio is given only as "the same as BERT" §3.2.
• Classification head. Section 4.4 says downstream models use a single linear layer §4. The supplement's classification figure shows a pooling block of linear, ReLU, linear [supp. §6.4]. Section 4.4's deeper-head ablation (fc, ReLU, LayerNorm, fc) helps CORD and RVL-CDIP and hurts FUNSD tab. 7.
• Fine-tuning epochs. The supplement gives 100 for FUNSD and Kleister-NDA, 200 for CORD and 30 for RVL-CDIP [supp. §6.1]. The "Effect of pre-training" table says 100 epochs for all three tab. 6. The main text also refers to "Table 1 in supplemental" for the implementation details, which is the table in [supp. §6.1].
• Headline numbers vs ablations. The headline numbers use 5M pre-training pages. The ablations use 1M pages and 5 epochs, so the ablation rows (76.90 FUNSD, 93.36 CORD) are not comparable with the headline (83.34, 96.33).
• Claims about RVL-CDIP. The text says DocFormer beats all base and large variants on RVL-CDIP §4.2. That holds for the base model (96.17). The large model (95.50) is below LayoutLMv2-large (95.65) tab. 2.
• Pre-training data count. Section 2 says LayoutLM (v1) was pre-trained on 11M pages §2. Section 4.1 says the previous state of the art needed more than twice DocFormer's pre-training documents, "(11M)", to reach 82.76, which is the LayoutLMv2-base score §4.1 tab. 1. The two passages attach the 11M figure to different models.
• Parameter counts. DocFormer-large (T+S) is listed at 536M in tab. 1 and at 502M in tab. 3. Use the value that matches the table you cite.
• Kleister-NDA. Results are on the validation split, with ground truth made by an unspecified OCR heuristic §4.3. The caption's "2.5x the learning capacity" is not defined. Do not compare these numbers with test-set results.
• Cost. Pre-training takes 17 hours per epoch on A100 40GB, with 5 epochs over 5M pages [supp. §6.1]. The GPU count is not stated. The supplement's complexity table gives the multi-modal attention as 2·O(n²d), about two self-attention passes, with the spatial and relative terms left out [supp. §6.2].
• Numbering. Section, equation, table and figure numbers were counted from the LaTeX source, because no compiled PDF was in the download. Section 4.5 (ablations) and the supplement's section numbers (§6.x) follow LaTeX counting and may differ from the typeset PDF; check them against the PDF before relying on them.

Lineage

• Builds on: LayoutLM (text embedding initialization and the text-plus-spatial design) §2 §3.1; LayoutLMv2 (visual features as separate tokens §2, and the masking of image regions under masked words that DocFormer drops §3.2); BERT (masked language modeling) §3.2; the scaled dot-product attention of Vaswani et al. §3.1; ResNet50 §3.1.
• Compared against: BERT, RoBERTa, UniLMv2, LayoutLMv1, LayoutLMv2, BROS, TILT, SPADE and LAMBERT tab. 1–tab. 4.
• Superseded by: not known from the wiki's content.