⚡ electrical engineering

Weight Space Correlation Analysis: Quantifying Feature Utilization in Deep Learning Models

This paper introduces Weight Space Correlation Analysis, a methodology that quantifies feature utilization in deep learning models by measuring the alignment between classification heads, thereby verifying that medical imaging models for spontaneous preterm birth rely on clinically relevant signals rather than confounding metadata when not artificially biased.

Original authors: Chun Kit Wong, Paraskevas Pegios, Nina Weng, Emilie Pi Fogtmann Sejer, Martin Grønnebæk Tolsgaard, Anders Nymark Christensen, Aasa Feragen

Published 2026-02-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chun Kit Wong, Paraskevas Pegios, Nina Weng, Emilie Pi Fogtmann Sejer, Martin Grønnebæk Tolsgaard, Anders Nymark Christensen, Aasa Feragen

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 "Cheat Sheet" Trap

Imagine you are taking a difficult history exam. You study hard, but you also notice a weird pattern: every time the test paper is printed on blue ink, the answer is "A," and every time it's on red ink, the answer is "B."

You don't actually know the history facts. Instead, you start guessing based on the ink color. You get a perfect score on the practice tests because the teacher kept using that pattern. But the moment you take the real exam printed on green ink, you fail miserably.

In the world of medical AI, this is called Shortcut Learning.

The Goal: The AI is supposed to learn to diagnose a disease (like preterm birth) by looking at the anatomy in an ultrasound image.
The Shortcut: The AI might accidentally learn that "if the image comes from Scanner Model X, the patient is likely healthy," or "if the image has text on the side, it's a specific type of scan."
The Danger: The AI isn't looking at the baby; it's looking at the "ink color" (the scanner model or hospital logo). If you move the AI to a new hospital with a different scanner, it might fail because its "cheat sheet" no longer works.

The New Tool: The "Weight Space Correlation" (WSC) Analysis

The authors of this paper wanted to answer a very specific question: Just because the AI sees the cheat sheet, does it actually use it to make its decision?

Previous methods could only say, "Hey, the AI can guess the scanner model just by looking at the image." But that doesn't prove the AI is using that guess to diagnose the patient.

The authors invented a new method called Weight Space Correlation (WSC). Here is how it works, using an analogy:

The Analogy: The Chef and the Ingredients

Imagine the AI is a Chef (the model) trying to cook a specific dish (the clinical diagnosis, like "Preterm Birth").

The Ingredients are the features in the image (the baby's anatomy, the texture of the tissue).
The Recipe is the AI's internal logic (its "weights").

The authors wanted to know: Is the Chef using the Fresh Vegetables (the real medical signal) or the Salt Shaker (the scanner model metadata) to decide if the dish is good?

Step 1: The "Taste Test" (Embedding): First, they checked if the Chef could even taste the salt. They found that yes, the Chef could tell which salt shaker was used just by looking at the ingredients. The "salt" (scanner info) was definitely in the bowl.
Step 2: The "Recipe Check" (The WSC Analysis): This is the magic part. They compared the Chef's Recipe for the Dish against the Recipe for identifying the Salt Shaker.
- If the two recipes are identical (high correlation), it means the Chef is using the same mental steps to identify the dish as they are to identify the salt. This is a shortcut. The Chef is cheating.
- If the two recipes are completely different (low correlation), it means the Chef is using one set of rules for the dish and a totally different set of rules for the salt. This is good. The Chef is ignoring the salt and focusing on the vegetables.

What They Tested

The team tested this method in two ways:

1. The "Fake Cheat" Experiment (Validation)
They took a dataset and deliberately messed it up. They made sure that Scanner A only took pictures of Heads, and Scanner B only took pictures of Torsos.

Result: The AI learned the shortcut. When they ran the WSC analysis, the "Recipe for the Body Part" and the "Recipe for the Scanner" were perfectly aligned. The method successfully caught the AI cheating.

2. The Real-World Test (SA-SonoNet)
They applied this to a real, advanced AI model used to predict Spontaneous Preterm Birth (sPTB) from cervical ultrasound images.

The Question: Does this real-world AI rely on the scanner model or the hospital ID to make its prediction?
The Result:
- Good News: The AI's "recipe" for predicting preterm birth was highly aligned with Cervical Length (a real, medically important factor). It was also aligned with Pixel Spacing (a technical setting that matters for image quality).
- Crucial Finding: The AI's "recipe" for predicting preterm birth was completely unaligned (decoupled) with the Scanner Model.
- Conclusion: Even though the AI could tell which scanner took the picture, it was not using that information to make its diagnosis. It was looking at the actual medical features.

Why This Matters

This paper provides a "trust check" for medical AI.

Before this, if an AI performed well, we didn't know if it was a genius doctor or a cheater looking at ink colors.
Now, we have a tool to look inside the AI's brain and say: "We checked your logic. You are not using the scanner model as a shortcut. You are actually looking at the patient."

Summary

The authors created a mathematical "lie detector" for AI. They proved that while medical images often contain hidden clues about the equipment used (like scanner models), a well-trained AI can be designed to ignore those clues and focus only on the real medical facts. Their method confirms that the specific AI they tested is trustworthy because it isn't relying on the "cheat sheet."

Technical Summary: Weight Space Correlation Analysis for Quantifying Feature Utilization

Problem Statement
Deep learning models in medical imaging are susceptible to "shortcut learning," where models rely on confounding metadata (e.g., scanner model, acquisition protocol) rather than genuine clinical signals. While previous research has demonstrated that such confounding information is often encoded within image embeddings (e.g., Gichoya et al., 2022; Glocker et al., 2023), a critical gap remains: the mere presence of a feature in the embedding space does not prove the model actively utilizes it for final predictions. Existing methods, such as analyzing embedding distributions or generating counterfactual test sets, do not directly address whether the classification head leverages these confounders. This paper addresses the fundamental question: Does the model's decision boundary actively align with clinically irrelevant factors?

Methodology: Weight Space Correlation (WSC) Analysis
The authors propose Weight Space Correlation (WSC) analysis, an interpretable methodology to quantify feature utilization by measuring the alignment between the classification heads of a primary clinical task and auxiliary metadata tasks. The method proceeds in three steps:

Linear Representation: The method treats the weights of the linear classification head ( $W$ ) as "attention vectors" acting upon the latent embedding space. For a task with $C$ classes, the weight matrix $W \in \mathbb{R}^{C \times d}$ defines the decision directions for each class.
Projection onto Data Manifold: To ensure comparisons are made on meaningful dimensions, classifier weights are projected onto the intrinsic data manifold using Principal Component Analysis (PCA). The authors extract the top $k$ principal components (enforcing a minimum of 50 and capturing 99% of variance) to form a projection matrix $P$ . The weights are transformed as $W' = W P^\top$ . This step filters out orthogonal, non-informative dimensions.
Quantifying Alignment: The method calculates the cosine similarity between the projected weight vectors of the primary task ( $W'_A$ ) and auxiliary metadata tasks ( $W'_m$ ). A high correlation between the decision boundaries of a clinical task and a metadata task (e.g., scanner type) serves as a quantitative proxy for shortcut learning, indicating the model relies on the same latent features for both predictions.

Key Contributions

Novel Metric: Introduction of WSC analysis to distinguish between the encoding of metadata and its active utilization in inference.
Validation via Induced Bias: The authors validated the method by creating a controlled "Induced Bias Dataset" where a strong correlation was artificially introduced between fetal plane classes and scanner models. WSC analysis successfully detected the resulting high alignment in weight vectors, whereas the "Balanced Dataset" (no correlation) showed low alignment, confirming the method's sensitivity to shortcut learning.
Application to Clinical Models: The methodology was applied to the SA-SonoNet model for Spontaneous Preterm Birth (sPTB) prediction, a state-of-the-art model in obstetric ultrasound.

Experimental Results
The study utilized two private clinical ultrasound datasets: a fetal anatomical classification dataset and a cervix dataset for sPTB prediction.

Encodability vs. Utilization: In the fetal dataset, baseline models confirmed that metadata (scanner, hospital ID) was highly encodable in the image embeddings (high AUROC when probing). However, WSC analysis revealed that the weight vectors for the primary plane classification task were largely orthogonal (low correlation) to those of the scanner classification task. This held true even for multi-task models explicitly trained to encode metadata, suggesting the primary task did not rely on scanner shortcuts.
Induced Shortcut Detection: When the dataset was biased to correlate specific planes with specific scanners, WSC values between plane and scanner weights increased significantly. This confirmed that WSC can detect when a model is forced to adopt a shortcut strategy.
SA-SonoNet Analysis: Applied to the sPTB prediction model, the analysis showed:
- Clinical Relevance: The sPTB classifier weights were moderately correlated with clinically relevant factors like cervical length and pixel spacing.
- Decoupling from Confounders: The weights were decoupled from clinically irrelevant acquisition factors, specifically scanner type.
- Caveat: While birth weight showed high WSC with sPTB, the authors note this was functionally meaningless because the probing head failed to predict birth weight effectively, indicating the information was not robustly encoded.

Significance and Claims
The paper claims that WSC analysis provides a necessary tool for verifying model trustworthiness in medical imaging. By inspecting whether a classifier utilizes features unrelated to the genuine clinical signal, researchers can move beyond simple performance metrics to understand the internal decision logic. The authors conclude that while confounding factors are often present in medical image embeddings, models like SA-SonoNet can be designed or trained such that their decision boundaries remain aligned with clinically meaningful signals and decoupled from spurious correlations like scanner models. This offers a method to confirm that a model is learning the intended pathology rather than exploiting dataset artifacts.

Limitations
The authors acknowledge that the method assumes a linear relationship between the embedding space and class predictions, which aligns with standard deep learning conventions but may not capture complex non-linear interactions. Additionally, the method currently analyzes individual metadata attributes; real-world shortcuts may involve non-linear combinations of multiple confounders, which future work aims to investigate. The paper also notes that low WSC does not guarantee robustness to domain shifts caused by other factors, such as changes in image quality or feature shifts, but it does rule out shortcut learning as the cause of such failures.