Entity-Aware Transformer

Entity-aware transformers employ a dual-encoder design that processes text tokens and entity tokens through separate but interacting pathways:

• Text encoder: Handles standard subword tokenization and contextual word representations using a standard transformer stack
• Entity encoder: Processes entity embeddings derived from a pre-trained knowledge graph embedding space (e.g., TransE, RotatE)
• Cross-encoder attention: A modified self-attention layer where text tokens attend to entity tokens and vice versa, allowing bidirectional knowledge flow

This separation prevents the model from confusing lexical semantics with structured world knowledge. In LUKE (Language Understanding with Knowledge-based Embeddings), the entity-aware self-attention mechanism computes separate attention scores for token-to-token, token-to-entity, entity-to-token, and entity-to-entity interactions.

Rather than treating entities as plain text strings, entity-aware transformers initialize entity representations using pre-trained knowledge graph embeddings that encode relational structure:

• Entity vectors are sourced from large-scale knowledge bases like Wikidata or Wikipedia hyperlink graphs
• Embeddings capture multi-hop relational information (e.g., 'Paris' is connected to 'France' via the capitalOf relation)
• During fine-tuning, these embeddings are updated via gradient descent, allowing the model to refine entity representations for downstream tasks

ERNIE (Enhanced Representation through kNowledge IntEgration) by Baidu takes a complementary approach by masking entity spans during pre-training rather than individual tokens, forcing the model to learn entity-level semantics alongside lexical patterns.

The core innovation lies in reformulating the scaled dot-product attention to accommodate entity tokens as distinct input types:

• Standard attention: Attention(Q, K, V) = softmax(QK^T / √d_k)V
• Entity-aware attention: Expands the key and value matrices to include entity representations alongside token representations
• The model computes four attention sub-matrices: token-to-token, token-to-entity, entity-to-token, and entity-to-entity
• A learned entity type embedding is added to distinguish different entity categories (person, organization, location, etc.)

This design allows the model to resolve ambiguous mentions by consulting structured knowledge. For example, when processing 'Apple released a new product,' the entity-aware attention can disambiguate between the company and the fruit by attending to the entity embedding for Apple Inc.

Entity-aware transformers use specialized pre-training objectives that treat entities as atomic units rather than token sequences:

• Entity masking: Instead of masking random subword tokens, entire entity spans are masked and the model must predict them using both textual context and knowledge graph embeddings
• Entity replacement detection: The model is trained to detect when an entity has been replaced with a random alternative of the same type, forcing it to verify entity-context consistency
• Relation prediction: An auxiliary objective predicts the relationship between two entities mentioned in the same passage using the knowledge graph as a distant supervision signal

These objectives teach the model that entities are semantically coherent units with external grounding, improving performance on entity typing, relation extraction, and question answering benchmarks.

Integrating entity embeddings introduces additional computational overhead that must be managed for practical deployment:

• The number of entity tokens in a sequence scales with the entity density of the input, which can be unpredictable
• Entity-aware attention computes O((n + e)^2) attention scores where n is token count and e is entity count
• Sparse attention patterns are often employed to restrict entity attention to relevant local context windows
• Knowledge graph embedding lookups add latency during inference unless cached or pre-computed

Production implementations like LUKE use a pre-filtering step to identify entity mentions before the transformer forward pass, limiting entity token injection to spans with high NER confidence scores. This reduces the effective entity count to a manageable subset of the full knowledge base.

Entity-aware transformers demonstrate measurable improvements on tasks requiring factual grounding and entity reasoning:

• Named Entity Recognition: +2-3 F1 points over BERT-base on CoNLL-2003 by leveraging entity type embeddings
• Relation Extraction: State-of-the-art results on TACRED and DocRED by attending to entity pair representations
• Entity Linking: Improved disambiguation accuracy on AIDA-CoNLL by incorporating knowledge graph context directly into the encoding
• Question Answering: Enhanced performance on entity-rich questions in datasets like WebQSP and TriviaQA

These gains are most pronounced when the target entities have strong coverage in the underlying knowledge graph. Performance degrades gracefully for long-tail entities with sparse graph connections.