Document world models: a cognitive-inspired auxiliary training framework for document-level relation extraction

Document-level Relation Extraction (DocRE) requires models to integrate relation evidence scattered across multiple sentences. Yet existing encoder-based methods perform relation classification in a single pass over a static final representation, causing long-range evidence to decay before it can inform the classification decision. Inspired by theories of incremental reading comprehension, we introduce the Document World Model (DocWM), a lightweight training-only auxiliary framework that encourages progressive entity-state maintenance and relation-belief stabilization in standard DocRE encoders. DocWM introduces two auxiliary training objectives. Entity State Tracking (ES) maintains an entity-specific recurrent hidden state, updated sentence by sentence and aligned with entity embeddings pooled from mentions up to each sentence to preserve early evidence. Relation Belief Dynamics (BD) supervises relation predictions at each reading step via a temporally weighted loss, guiding relation beliefs to stabilize progressively toward the correct label. Experiments on Re-DocRED and HacRED show consistent improvements over the compared baselines, most clearly on long-range cross-sentence relations.