🤖 machine learning

KG-SoftMAP: Soft Knowledge-Graph Priors for Bayesian Network Structure Learning from Sparse Discrete Data

The paper introduces KG-SoftMAP, a Bayesian network structure learning method that leverages imperfect, confidence-weighted knowledge graph priors to effectively recover causal structures from sparse discrete data where traditional data-only approaches fail, demonstrating superior performance on synthetic benchmarks and providing calibrated, knowledge-consistent diagnostic models for real-world educational datasets.

Original authors: Guoliang Xu, James E. Corter

Published 2026-06-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Guoliang Xu, James E. Corter

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Missing Puzzle Pieces"

Imagine you are trying to solve a giant jigsaw puzzle to understand how a complex system works (like how different medical symptoms relate to diseases, or how different math concepts build on each other).

Usually, you need to see every piece of the puzzle together to figure out how they fit. But in many real-world situations, your data is sparse. This means for any single person or event, you only see a tiny, random handful of pieces.

The Analogy: Imagine trying to figure out the rules of a board game, but you only get to watch 5 seconds of a game once a week, and every time you watch, you only see 3 random players. You can't see who is influencing whom because you rarely see two players interacting at the same time.
The Result: Standard computer methods that rely only on this data get stuck. They can't find any patterns because the pieces are too scattered.

The Solution: A "Soft" Map from an Expert

The authors propose a new method called KG-SoftMAP. Instead of waiting for the data to be perfect, they bring in a "helper" in the form of a Knowledge Graph (KG).

The Analogy: Think of the Knowledge Graph as a rough sketch drawn by a domain expert (or an AI). The expert says, "I'm pretty sure Concept A leads to Concept B, but I'm not 100% certain."
The "Soft" Part: This is the crucial innovation.
- Hard Constraints (Old Way): Some methods treat the expert's sketch like a law. If the expert says "A leads to B," the computer must draw that line, even if the data later proves the expert wrong. This is brittle; if the expert makes a mistake, the whole model breaks.
- Soft Priors (KG-SoftMAP): This method treats the expert's sketch like a suggestion. It says, "The expert thinks A leads to B, so we'll start there. But if the data we do have strongly suggests otherwise, we will ignore the expert and follow the data." It's a "soft" nudge, not a hard rule.

How It Works: The "MAP" Objective

The computer uses a mathematical formula to find the best puzzle solution. It balances two things:

The Data Fit: How well does this puzzle match the few pieces we actually saw?
The Expert Suggestion: Does this puzzle look like the expert's rough sketch?

The computer tries to maximize both. If the data is very weak (which it usually is in these sparse cases), the expert's sketch guides the computer to a good starting point. If the data is strong enough to contradict the expert, the computer overrules the expert.

The "LLM" Trick: Building the Sketch

Often, we don't have an expert sketch ready. The paper shows how to use a Large Language Model (LLM) (like a smart AI chatbot) to create this sketch automatically.

The Process: You feed the AI a bunch of reference text (like textbook answers or medical guidelines). The AI reads it and says, "Okay, based on this text, here is a list of concepts and a map of how they probably connect, with a confidence score for each connection."
The Safety Net: Because the AI might hallucinate (make things up), the "Soft" nature of the method is vital. If the AI guesses a connection that the data proves is wrong, the data wins, and the connection is dropped.

What the Experiments Showed

The authors tested this in two ways:

1. The Synthetic Test (The "Controlled Lab")
They created fake puzzles where they knew the true answer.

The Result: When data was extremely scarce (only 5% of pieces visible), standard methods found almost nothing (0% success). KG-SoftMAP, using the expert sketch, managed to find a significant chunk of the correct structure (up to 96% success when data was slightly better).
The Lesson: The method works best when the expert sketch is mostly right but not perfect. If the sketch is random garbage, the method fails gracefully (it doesn't get confused; it just performs like it has no sketch at all).

2. The Real-World Test (Education Data)
They tried this on real student data (Short Answer Feedback), where students answer different questions, leaving huge gaps in the data.

The Goal: They weren't trying to prove the AI found the "true" structure (because no one knows the true structure of student learning). Instead, they checked if the model was useful for prediction and diagnosis.
The Result:
- Prediction: A simple "Logistic Regression" model (a standard, non-graph method) was slightly better at just guessing the right answer.
- Diagnosis: However, KG-SoftMAP provided something the simple model couldn't: a calibrated map. It could tell you, "If a student failed Concept A, there is a 70% chance they also failed Concept B," and it did this while respecting the logical flow of the subject matter.
- The Trade-off: If you only care about the final score, use the simple model. If you want to understand why a student is struggling and how concepts connect, use KG-SoftMAP.

Summary

KG-SoftMAP is a tool for learning from messy, incomplete data. It combines the "best guess" of an expert (or an AI reading a textbook) with the actual evidence from the data. It treats the expert's advice as a helpful guide rather than an unbreakable law, allowing it to find patterns in data that is too sparse for traditional methods to handle.

Key Takeaway: It's like having a GPS that knows the general layout of a city (the Knowledge Graph) but is smart enough to reroute you if it sees a roadblock in real-time (the Data), even if you can only see a few blocks ahead.

Technical Summary: KG-SoftMAP

Problem Statement

Learning the structure of Bayesian Networks (BNs) from data is a core but NP-hard problem. This task becomes particularly intractable when the data is sparse and discrete, a common scenario in domains like educational assessments, clinical registries, and sensor networks where each instance records only a small subset of variables (observation rates often below 5%). Under such sparsity, most variable pairs lack sufficient joint observations for reliable scoring or conditional independence testing. Consequently, classical and modern data-only methods (both constraint-based and score-based) fail to recover meaningful directed structure, often returning near-empty graphs or incorrect skeletons.

While purely data-driven approaches struggle, imperfect domain knowledge regarding variable relationships often exists outside the data. This knowledge can be expressed as a weighted, directed Knowledge Graph (KG). However, existing methods for injecting such knowledge tend to sit at extremes: either using hard constraints (fixing or forbidding edges), which are brittle when the knowledge is uncertain, or prompting Large Language Models (LLMs) to generate a graph directly, effectively discarding the data. There is a lack of a general framework to treat a weighted KG as a soft, data-overridable structure prior that can guide learning in sparse regimes without forcing incorrect edges.

Methodology: KG-SoftMAP

The authors propose KG-SoftMAP, a two-stage approach that learns BN structures by maximizing a Maximum A Posteriori (MAP) objective. The method is designed to be agnostic to the source of the KG (expert-curated or LLM-extracted).

1. Knowledge Graph Construction (Stage 1)

If a curated KG is not available, the authors provide an optional pipeline using an LLM to extract a weighted KG from domain reference text.

Schema: The LLM extracts Concept nodes and PRECEDES relations (forming the KG edges) with confidence weights $w_{uv} \in [0, 1]$ .
Instance-Level Deviations: If instance-level text (e.g., student responses) is available, the LLM also identifies concept-specific mistakes with severity weights. These are used only during data preprocessing to refine the observation matrix, not as part of the structural prior.
Output: A weighted directed graph $K = \{(u, v, w_{uv})\}$ .

2. Structure Learning (Stage 2)

The core algorithm learns a Directed Acyclic Graph (DAG) $G$ by maximizing the log-posterior:
$G^* = \arg \max_G \left[ \log P(D | G) + \log P(G | K) \right]$

Data Likelihood ( $\log P(D | G)$ ): The authors use the BDeu score (Bayesian Dirichlet equivalent uniform). To handle sparsity, they employ complete-case scoring, where only rows where a node and all its parents are observed contribute to the counts. This avoids imputation bias. An optional EM extension is available for Missing At Random (MAR) settings but is computationally expensive and not the default.
KG Structure Prior ( $\log P(G | K)$ ): The authors introduce an edge-factored soft prior. The log-odds of including an edge $(u \to v)$ $(u \to v)$ is a linear function of its KG confidence weight $w_{uv}$ $w_{uv}$ :
$\text{logit}(\theta_{uv}) = \beta_0 + \beta_1 \cdot w_{uv}$
Here, $\theta_{uv}$ $θ_{uv}$ is the prior probability of edge inclusion.
- Edges absent from the KG ( $w_{uv}=0$ ) have a low base inclusion probability $\theta_0$ (acting as a sparsity penalty).
- Edges present in the KG with high confidence ( $w_{uv}=1$ ) have a high inclusion probability $\theta_{max}$ .
- Data-Overridable: The prior is "soft." If data evidence strongly supports an edge not in the KG, it can be added; conversely, if data strongly penalizes a KG edge, it can be dropped.
Search Procedure: The optimization is solved via a greedy coordinate ascent (hill-climbing) with three operations: edge addition, edge swap (replacing a parent), and edge deletion. The search enforces acyclicity and a maximum in-degree constraint.

Key Contributions

Soft, Data-Overridable KG Prior: The paper casts a weighted KG as an edge-factored structure prior where the inclusion log-odds are linear in the KG confidence. This allows the method to leverage imperfect domain knowledge while remaining robust to errors, as the data can override the prior.
Controlled Evidence of Recovery: On synthetic benchmarks with ground-truth DAGs, KG-SoftMAP demonstrates the ability to recover partial directed structure even at extreme sparsity ( $\rho=0.05$ ), where data-only baselines (including modern differentiable methods like GOLEM, DAGMA, and SDCD, and discrete permutation searches like GRaSP) recover near-zero directed structure. Performance improves substantially as observation rates increase ( $\rho \ge 0.2$ ).
Robustness to KG Quality: Ablation studies show that performance degrades gracefully as KG quality drops (via corruption or noise). Crucially, when provided with a "zero-signal" random KG (containing no true edges), the method recovers no structure (Directed-F1 $\approx 0.00$ ), confirming that gains are driven by the structural signal in the KG, not the search algorithm alone.
Deployment Evaluation on Real Data: The authors evaluate the method on three real-world sparse educational datasets (SAF, ASSISTments, Eedi). Since these lack ground-truth DAGs, evaluation focuses on deployment metrics: prediction, calibration, and KG-consistency.
- SAF (LLM-extracted KG): The learned BN provides a calibrated, interpretable diagnostic model with KG-consistent edges. It trails logistic regression (LR) by a small margin in prediction (0.03 F1) but offers the ability to infer from arbitrary observed subsets and provides a reusable concept graph.
- ASSISTments (No meaningful KG): When instantiated with heuristic priors (name-similarity, random), the BN performs poorly compared to LR, confirming the method requires a meaningful structural signal.
- Eedi (Expert Ontology): Using an independent expert ontology as a prior increases the share of within-topic edges (structural consistency) but does not improve prediction, highlighting the trade-off between structural coherence and discriminative power.

Results

Synthetic Benchmarks: At $\rho=0.2$ , KG-SoftMAP achieves Directed-F1 scores between 0.46 and 0.96 across various network sizes, significantly outperforming all data-only baselines (which remain below 0.44). At $\rho=0.05$ , it recovers partial structure (DF1 0.14–0.29) while baselines fail.
Ablations:
- Prior vs. Search: The gain is primarily attributed to the soft prior. Even standard learners (like MMHC or BNLEARN) equipped with the same KG prior recover substantial structure, though KG-SoftMAP's sparse-aware scoring provides an additional margin.
- Corruption: Recovery degrades smoothly with KG corruption rates.
- Hard vs. Soft: Under high-quality KGs and sufficient data, a near-hard prior performs comparably to the soft prior. The soft prior's advantage lies in safety when KG fidelity is uncertain.
Real Data: The method produces a "diagnostic model" that is well-calibrated and KG-consistent. However, for pure cell-level prediction, discriminative models (like Logistic Regression) remain the stronger baseline, especially when no meaningful KG exists.

Significance and Claims

The paper positions KG-SoftMAP as a solution for extreme sparsity where data alone is insufficient for structure learning. Its significance lies in:

Bridging the Gap: It occupies a middle ground between brittle hard constraints and data-agnostic LLM generation, treating domain knowledge as a probabilistic prior that can be overridden by data.
Defining Scope: The authors are modest about their claims. They explicitly state that structural recovery claims are confined to synthetic data with ground truth. For real data, the method is presented as a diagnostic tool that provides interpretable, calibrated joint probabilities and KG-consistent structures, rather than a superior predictor for cell-level outcomes.
Practical Utility: The method is most valuable when a meaningful (even if imperfect) KG exists. If no such knowledge is available, the paper concludes that discriminative baselines are preferable.

The work does not claim to solve general causal discovery or replace discriminative prediction; rather, it offers a robust framework for learning interpretable probabilistic models in data-scarce regimes by effectively leveraging external, imperfect domain knowledge.