# How Doc2Graph Performs Key-Value Pair Detection Using Edge Classification

> Doc2Graph detects key-value pairs by framing extraction as binary edge classification, distinguishing key-value relationships from unrelated pairs.

- Repository: [Andrea Gemelli/doc2graph](https://github.com/andreagemelli/doc2graph)
- Tags: deep-dive
- Published: 2026-02-24

---

**Doc2Graph detects key-value pairs by framing the extraction task as a binary edge classification problem, where edges connecting text blocks are classified as either key-value relationships (class 1) or unrelated pairs (class 0).**

The andreagemelli/doc2graph repository implements a graph neural network approach to document understanding where **key-value pair detection using edge classification** replaces traditional sequence labeling. By constructing a DGL graph from document layout and text, the system applies Graph Convolutional Networks (GCN) to predict which edges connect semantic keys to their corresponding values, enabling robust extraction from structured documents.

## Graph Construction: Building Document Graphs

The pipeline begins with the **GraphBuilder** class in [`doc2graph/data/graph_builder.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/data/graph_builder.py), which transforms raw document images or PDFs into graph representations suitable for neural processing.

Nodes correspond to text blocks (words or lines) and are enriched with **multimodal features** combining visual appearance, textual content, and spatial layout. The builder aggregates these features between lines 19-27, creating feature vectors that capture both semantic and positional information.

Edges represent potential relationships between text blocks. Depending on the `edge_type` configuration, the builder generates either fully-connected graphs or applies a k-NN heuristic to limit connections to spatially proximate nodes (lines 57-78). This edge generation strategy ensures that the model evaluates only geometrically plausible key-value candidates while maintaining computational efficiency.

## Edge Classification Architecture

The core of Doc2Graph's key-value extraction lies in the **EdgeClassifier** network defined in [`doc2graph/models/graphs.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/models/graphs.py) (lines 66-100). This architecture processes the constructed graphs through several stages:

1. **Node Feature Projection** – Input features are projected into a latent space compatible with message passing.
2. **GCN Message Passing** – Multiple GCN layers propagate information across the graph, allowing nodes to aggregate context from neighboring text blocks.
3. **Edge Prediction** – The **MLPPredictor** consumes the refined node embeddings to output edge logits, performing binary classification between key-value links and non-relationships.

For end-to-end training scenarios, the repository provides the **E2E** variant featuring the `MLPPredictor_E2E` class (lines 236-244). This enhanced predictor incorporates **node-wise class probabilities** (`cls_u`, `cls_v`) and **polar features** that encode directional relationships between text blocks. By concatenating these additional signals with standard node embeddings, the E2E variant achieves more accurate edge classification by leveraging entity type information alongside structural cues.

## Inference Pipeline: From Edge Logits to Key-Value Pairs

During inference, the trained model processes input graphs through the pipeline defined in [`doc2graph/inference.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/inference.py). The system applies `softmax` to the raw edge logits `e` to obtain class probabilities, then identifies edges where the predicted class equals 1 (lines 54-60).

For each detected key-value edge, the inference logic extracts the **source node** as the key and the **target node** as the value (lines 68-73). This directional interpretation ensures that semantic relationships maintain proper orientation—keys precede values in the logical flow. The extracted pairs are serialized to JSON format and optionally visualized as bounding box overlays on the original document images.

## Implementation Examples

The following examples demonstrate practical usage of Doc2Graph's edge classification system for key-value extraction.

*Running inference on a folder of document images:*

```python
from doc2graph.inference import inference
from pathlib import Path

# Choose a pretrained checkpoint (e.g., "EDGE-funsd-epoch10.pt")

weights = ["EDGE-funsd-epoch10.pt"]

# Path(s) to the input images (can be a list or a glob pattern)

paths = list(Path("data/images").glob("*.png"))

# Perform inference – the function returns PNG visualisations and JSON files with

# detected key-value pairs in `doc2graph/INFERENCE/`.

inference(weights, paths, device=0)   # device=-1 uses CPU

```

*Instantiating the edge-classification model directly for custom workflows:*

```python
from doc2graph.models.graphs import SetModel
from doc2graph.utils import get_config

# Load model configuration (edge-type = "EDGE")

model_name = "EDGE"
cfg = get_config("preprocessing")   # contains feature settings, etc.

chunks = [...]                      # list of feature chunk sizes from FeatureBuilder

device = "cpu"

set_model = SetModel(name=model_name, device=device)
model = set_model.get_model(
    node_classes=5,            # number of node classes (depends on dataset)

    edge_classes=2,            # binary edge classification (key-value vs other)

    chunks=chunks,
    verbatim=False,
)
model.eval()

```

*Manually extracting key-value pairs from graph predictions:*

```python
import torch.nn.functional as F

# `graph` is a DGL graph, `features` contains node features under "feat"

node_emb, edge_logits = model(graph.to(device), graph.ndata["feat"].to(device))

# Edge predictions (class 1 = key-value)

_, edge_pred = torch.max(F.softmax(edge_logits, dim=1), dim=1)
kvp_edge_ids = (edge_pred == 1).nonzero(as_tuple=True)[0]

u, v = graph.edges()
key_texts   = [features["texts"][u[i]] for i in kvp_edge_ids]
value_texts = [features["texts"][v[i]] for i in kvp_edge_ids]

kvp_pairs = [{"key": k, "value": v} for k, v in zip(key_texts, value_texts)]
print(kvp_pairs)

```

## Summary

- **GraphBuilder** ([`doc2graph/data/graph_builder.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/data/graph_builder.py)) constructs DGL graphs from documents using multimodal node features (visual, textual, layout) and configurable edge generation strategies (fully-connected or k-NN).
- The **EdgeClassifier** ([`doc2graph/models/graphs.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/models/graphs.py)) implements a GCN-based architecture with MLP edge predictors to perform binary classification on graph edges.
- During inference ([`doc2graph/inference.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/inference.py)), edges predicted as class 1 indicate **key-value relationships**, with source nodes mapping to keys and target nodes to values.
- The E2E variant enhances classification accuracy by incorporating node class probabilities and polar directional features into the edge prediction mechanism.

## Frequently Asked Questions

### What distinguishes the E2E variant from the standard EdgeClassifier?

The **E2E** (end-to-end) variant extends the base architecture by feeding node classification probabilities (`cls_u`, `cls_v`) and polar geometric features into the `MLPPredictor_E2E`. As implemented in [`doc2graph/models/graphs.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/models/graphs.py) lines 236-244, this allows the edge classifier to leverage entity type information and directional cues between text blocks, improving key-value detection accuracy compared to the standard **EdgeClassifier** which relies solely on node embeddings.

### How does Doc2Graph determine which text blocks to connect with edges?

The system uses the **GraphBuilder** class in [`doc2graph/data/graph_builder.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/data/graph_builder.py) (lines 57-78) to generate edges based on the `edge_type` configuration. It supports two strategies: **fully-connected** graphs where every node connects to every other node, or **k-NN** heuristics that limit connections to the k-nearest spatial neighbors. This configuration balances between capturing all possible relationships and maintaining computational tractability.

### What features represent individual text blocks as graph nodes?

Nodes receive **multimodal embeddings** combining three feature categories: visual features encoding appearance and rendering, textual features from word embeddings or OCR outputs, and layout features capturing bounding box coordinates and spatial positions. The `GraphBuilder` aggregates these features (lines 19-27) to create comprehensive node representations that enable the GCN to distinguish between different semantic and structural document elements.

### How are key-value pairs formatted after edge classification?

During inference ([`doc2graph/inference.py`](https://github.com/andreagemelli/doc2graph/blob/main/doc2graph/inference.py) lines 68-73), the system identifies edges classified as class 1 (key-value links), extracts the source node's text as the **key** and the target node's text as the **value**, and structures them as JSON objects with `{"key": "...", "value": "..."}` format. These pairs are saved to disk alongside optional visualizations showing the detected relationships overlaid on the original document image.