Structure-aware World Model for Probe Guidance via Large-scale Self-supervised Pre-train

Abstract

The complex structure of the heart leads to significant challenges in echocardiography, especially in acquisition cardiac ultrasound images. Successful echocardiography requires a thorough understanding of the structures on the two-dimensional plane and the spatial relationships between planes in three-dimensional space. In this paper, we innovatively propose a large-scale self-supervised pre-training method to acquire a cardiac structure-aware world model. The core innovation lies in constructing a self-supervised task that requires structural inference by predicting masked structures on a 2D plane and imagining another plane based on pose transformation in 3D space. To support large-scale pre-training, we collected over 1.36 million echocardiograms from ten standard views, along with their 3D spatial poses. In the downstream probe guidance task, we demonstrate that our pre-trained model consistently reduces guidance errors across the ten most common standard views on the test set with 0.29 million samples from 74 routine clinical scans, indicating that structure-aware pre-training benefits the scanning.

1 Introduction

Cardiovascular diseases are the leading cause of death in worldwide [18, 19]. Echocardiography is the most commonly used method in clinical practice to assess heart conditions. However, the structure of the heart is extremely complex. According to [12], up to seven anatomical structures need to be identified in a single plane, such as the Parasternal Long-axis Plane (Fig. 1 Left). Recognizing these anatomical structures is crucial for diagnosis, and understanding their spatial relationships in a two-dimensional plane also helps the sonographer fine-tune the probe to obtain the best quality images. Additionally, up to 27 different planes need to be examined, requiring the sonographer to understand the spatial relationships between these planes in three-dimensional space and to finely adjust the ultrasound probe’s position to reach the target location (Fig. 1 Right). Due to the aforementioned reasons, cardiac ultrasound examinations are extremely challenging. This also results in a long training period for ultrasound medical personnel, as they need to spend a lot of time familiarizing themselves with both 2D and 3D structures. Consequently, there is a significant talent shortage in this field, especially in regions with scarce medical resources, such as Africa.

Figure 1: Diagram illustrating the capabilities of a cardiac world model. We aim to develop a cardiac world model that can understand both two-dimensional and three-dimensional structures. (Left) The world model needs to recognize various structures in two-dimensional planes and understand their spatial relationships for in-plane probe adjustment. (Right) Understanding the three-dimensional structure of the heart, specifically the spatial relationships between different planes, is crucial for out-of-plane probe adjustment. The images used in the diagram are sourced from [12].

With the rise of deep learning [4, 6, 7, 21], AI technology has shown great potential in improving the efficiency of echocardiography analysis. For example, Ouyang et al. [14] proposed a video-based deep learning algorithm, i.e., EchoNet-Dynamic, that can automatically make accurate assessments of cardiac function. More importantly, with the rise of large language models [2, 16, 17, 20] and multi-modal learning [8, 10, 11, 15], the interpretation of echocardiograms [3] has shown increasing improvement, demonstrating excellent performance in tasks such as pulmonary artery pressure estimation, left ventricular hypertrophy, heart failure, and left atrial enlargement. These AI-assisted diagnostic tools have demonstrated performance almost comparable to human experts. However, this comes with a prerequisite: the acquisition of high-quality echocardiograms. In regions with scarce medical resources, there are often no sonographers available to obtain high-quality echocardiograms. In such cases, the powerful capabilities of these AI-assisted diagnostic tools cannot be fully utilized. Few works [5, 9, 13] have focused on how to use AI technology to assist inexperienced sonographers in accurately acquiring target planes. Recently, Jiang et al. [9] proposed a cardiac dreamer, which only focuses on supervised learning of the 3D-structure of the heart and serves as a "heart map" for the probe guidance task. This work demonstrates a potential AI-assisted scanning method that is expected to improve the scanning skills of novices.

According to the clinical experience of doctors, understanding both the 2D and 3D structures is crucial for efficient scanning. For example, when you need to adjust the probe’s viewing angle on a particular plane to capture specific anatomical structures, you must have a good understanding of the spatial positions of those anatomical structures on that plane (Fig.1 Left). Thus, in this paper, we propose a 2D-3D joint structure-aware pre-training framework to obtain a data-driven cardiac world model that benefits ultrasound scanning. Specifically, we require the world model to learn important spatial relationships in the following ways: (1) for understanding two-dimensional structures, we use a masking approach that requires the world model to predict features at adjacent positions in the two-dimensional plane; (2) for understanding three-dimensional structures, we provide information on the positional changes of two planes in 3D space, requiring the world model to predict the features of the target plane after the positional change. We further collected expert operational data on acquiring the most common ten standard planes from 364 routine clinical scans, resulting in 1.36 million sample pairs gathered by three certified sonographers, to enhance the model’s learning of generalizable cardiac structure knowledge. Results on the downstream probe guidance tasks indicate that the proposed structure-aware pre-training method learns useful knowledge for assisting in the acquisition of echocardiograms.

2 Method

In this section, we describe the proposed structure-aware pre-training framework, illustrated in Fig. 2. We first introduce the prior I-JEPA work [1] , on which we based our method, in Section 2.1. Next, we discuss how we construct a pre-training framework that simultaneously learns 2D and 3D structural information by introducing three-dimensional spatial information in Section 2.2.

2.1 Preliminary

In the context of echocardiogram analysis, accurately understanding 2D structural information is fundamental for making correct diagnoses and conducting efficient scans. Recently, Assran et al. proposed a Joint-Embedding Predictive Architecture (I-JEPA) [1] to learn highly semantic image representation. The key idea for I-JEPA is learning the representation of images by predicting features of target blocks based on the non-overlap** context block and the positional embedding in the same image. This paradigm requires the model to understand the spatial relationships of different semantic structures in the 2D plane. For example, if a context block contains the head of a dog, it is likely that the body of the dog will be located below the context block. For echocardiogram analysis, this paradigm also enables the model to learn the spatial relationships of fine structures in the two-dimensional plane. For instance, the left ventricle (LV) is located below the interventricular septum (IVS), as shown on the left side of Fig.1. Therefore, it is highly suitable for modeling two-dimensional structural information in cardiac ultrasound images. Next, we briefly introduce the modeling and training method.

Targets. Specifically, the I-JEPA model employs a Vision Transformer [4]. Firstly, the input image $\mathbf{I}\in\mathbb{R}^{H\times W}$ is divided into $N$ non-overlap** patches, which are further processed by the target encoder $F^{{}^{\prime}}_{\theta}$. Then, we select a portion of the patches to form $M$ target blocks, which may have overlap** regions. We denote the ground-truth target blocks’ features as $\mathbf{Y}^{t}_{i}\in\mathbb{R}^{L_{t}\times C},i\in\{1,2,\cdots,M\}$.

Context and Prediction. For the context features, patches not belonging to the target blocks are randomly selected and input into the context encoder to obtain features $\mathbf{Z}_{c}\in\mathbb{R}^{L_{c}\times C}$. Subsequently, the predictor $W_{\theta}$ (world model) generates the features of target blocks based on the context features and positional embedding which indicate the location of context and target blocks. We denote the predicted target blocks’ feature as $\hat{\mathbf{Y}}^{t}_{i}\in\mathbb{R}^{L_{t}\times C},i\in\{1,2,\cdots,M\}$.

Loss. With the ground-truth and predicted features of target blocks, the model is optimized using the following loss:

$$\mathcal{L}_{total} =\frac{1}{M}\sum_{i=1}^{M}\sum_{j=1}^{L_{t}}\mathcal{L}_{\mathrm{SmoothL1}}(\mathbf{Y}^{t}_{i,j},\hat{\mathbf{Y}}^{t}_{i,j}).$$ (1)

Figure 2: Diagram illustrating the pre-training method and downstream task. (a) The world model and encoder are required to predict features on the target plane based on the spatial relationships in both two-dimensional and three-dimensional spaces. (b) Following cardiac dreamer [9], we choose the probe guidance task to evaluate our method. The cardiac images used in the diagram are sourced from [12].

2.2 2D-3D Joint Structure-aware Pre-training

Echocardiography involves acquiring high-quality echocardiogram and subsequently conducting analysis and diagnosis based on these images. However, previous researches [3, 14] have primarily focused on understanding and analyzing the two-dimensional plane, neglecting how AI can assit in acquiring high-quality echocardiogram. While understanding the 2D structure of a plane can aid in scanning, especially when some features of the standard view are already visible, this alone is insufficient. Particularly when transitioning from one standard view to another, a deep understanding of the heart’s 3D structure is essential. Therefore, we propose a 2D-3D joint structure-aware pre-training framework, as illustrated in Fig.1 (a). The core insight is predicting the visual features of structures at target locations based on given 2D and 3D positional conditions, thereby learning the map** relationship between spatial positions and visual features. Next, we provide details of the modeling and training method.

Input. Given a source image $\mathbf{I}^{s}\in\mathbb{R}^{H\times W}$, we select a target image $\mathbf{I}^{t}\in\mathbb{R}^{H\times W}$ from the same individual’s images. To ensure sufficient spatial pose variation between the two images, there must be an interval of at least 150 frames between their sequence numbers. Then, both images are divided into $N$ non-overlap** patches $\mathbf{I}^{s}_{p},\mathbf{I}^{t}_{p}\in\mathbb{R}^{N\times h\times w}$.

Context and Target. First, we randomly select a rectangular region from the patches of the source image $\mathbf{I}^{s}$ as the context information. These patches are then fed through the context encoder $F_{\theta}$, resulting in the context feature $\mathbf{Z}^{s}$:

$$\mathbf{Z}^{s} =F_{\theta}(\mathbf{I}^{s}_{p}\odot\mathbf{B}^{s}),\ \ \mathbf{Z}^{s}\in\mathbb{R}^{L_{s}\times C},$$ (2)

where $\mathbf{B}^{s}\in\mathbb{R}^{N\times 1\times 1}$ is a binary mask indicating the selected patches, $L_{s}$ is the number of context patches, and $C$ is the hidden dimension. Next, we select $M$ non-overlap** regions on the target image as the target blocks, requiring the model to understand both the three-dimensional structure of the heart and the spatial relationships of heart structures in the two-dimensional plane. If we choose regions on the target image that have the same spatial position as the context block, the model only needs to understand the three-dimensional spatial relationships. The target features are obtained as follows:

$$\mathbf{Y}^{t}_{i} =F_{\theta}(\mathbf{I}^{t}_{p})\odot\mathbf{B}^{t}_{i},\ \ \mathbf{Y}^{t}_{i}\in\mathbb{R}^{L_{t}\times C},i\in\{1,2,\cdots,M\},$$ (3)

where $\mathbf{B}^{t}_{i}\in\mathbb{R}^{N\times 1}$ is a binary mask indicating the selected patches and $L_{t}$ is the number of target patches.

Condition and Prediction. The three-dimensional spatial relationship between the source image and the target image is denoted as $\mathbf{a}\in\mathbb{R}^{6}$, which encapsulates the translation and rotation changes in the x, y, and z directions. This vector $\mathbf{a}$ is encoded by $P_{\theta}$ to obtain the pose embedding, denoted as $\mathbf{P}\in\mathbb{R}^{1\times C}$. Then, the 2D positional embeddings indicating the location of the context and target block are denoted as $\mathbf{Q}^{s}\in\mathbb{R}^{L_{s}\times C}$ and $\mathbf{Q}^{t}_{i}\in\mathbb{R}^{L_{t}\times C},i\in\{1,2,\cdots,M\}$. Finally, the predicted features from the world model $W_{\theta}$ are obtained as follows:

$$\mathbf{Z}^{s^{\prime}} =\mathbf{Z}^{s}+\mathbf{Q}^{s},$$ (4)

$$\mathbf{Z}^{t^{\prime}}_{i} =\mathbf{Z}^{t}_{i}+\mathbf{Q}^{t}_{i},$$ (5)

$$\hat{\mathbf{Y}}^{t}_{i} =W_{\theta}(\text{concat}[\mathbf{Z}^{s^{\prime}},\mathbf{P},\mathbf{Z}^{t^{\prime}}_{i}]),$$ (6)

where $\mathbf{Z}^{t}_{i}\in\mathbb{R}^{L_{t}\times C},i\in\{1,2,\cdots,M\}$ is a set of learnable parameters representing the features to be predicted. These parameters interact with the context feature and condition information to generate the target feature. Then, the models are optimized according to the loss function defined in Eq. (1).

3 Experiments

Figure 3: Anatomic and 2D ultrasound images of ten standard planes. The cardiac images used in the diagram are sourced from [12].

3.1 Implementation Details

Dataset. In this paper, we collected data and conducted experiments on the ten most common standard planes [12], as shown in Fig. 3. The ultrasound images and scan data were acquired following the procedure described in [9]. Ultimately, we amassed data from 364 routine clinical scans, resulting in a total of 1.36 million image and 3D pose data pairs. The whole data collection process was approved and supervised by The University Science and Technology Ethics Committee. We split the dataset into 290 scans (1.07 million samples) for training and 74 scans (0.29 million samples) for testing. For pre-training, only the training set was utilized. Both the training and test sets were employed for the downstream probe guidance task. It is important to note that the individuals in the training and test sets are different to avoid information leakage and fairly validate the model’s generalization performance.

Pre-training. The context and target encoders were implemented using ViT-Small/16, and the world model utilized a custom vision transformer with a depth of 6 layers and a hidden dimension of 384. The entire model was trained for 50 epochs with a batch size of 1024 on 8 Nvidia RTX-4090 GPUs. The training included a 7-epoch warmup period (starting learning rate is 1e-4), followed by a learning rate of 5e-4, using a cosine scheduler with a final learning rate of 5e-7. The implementation details for generating the context block and target blocks, as well as the hyper-parameters, followed the procedures described in [1].

Downstream task. For the probe guidance task, we adopted the framework and procedure proposed in [9], as shown in Fig.2 (b). The input for this task is an ultrasound image, and the output is the probe position adjustment needed to achieve a specific standard view. This task aims to assist junior ultrasound medical personnel in scanning, enhancing the success rate and quality of view acquisition. During fine-tuning, our pre-trained world model $W_{\theta}$ was loaded and optimized for 5 epochs with a batch size of 1024 on 8 Nvidia RTX-4090 GPUs. The learning rate was set to 1e-4, using a cosine scheduler with a final learning rate of 1e-6. The optimizer was set to AdamW.

Evaluation Metrics. The metric used for probe guidance task is the Mean Absolute Error (MAE) between the predicted and ground truth probe poses.

3.2 Results

Table 1: Evaluation of the probe guidance task. We report MAE results that represent probe guidance errors (lower is better) across ten standard planes.

Plane	Model	Translation (mm)			Rotation (degree)
		x	y	z	rx	ry	rz
PLAX	Cardiac Dreamer [9]	8.66	8.14	5.63	6.60	5.42	8.23
	$+$ Our Pre-train	8.39 (-3.15%)	8.02 (-1.53%)	5.53 (-1.80%)	6.46 (-2.12%)	5.34 (-1.42%)	7.89 (-4.07%)
PSAX-AV	Cardiac Dreamer	7.26	6.63	4.51	5.28	6.26	7.43
	$+$ Our Pre-train	7.06 (-2.73%)	6.57 (-0.97%)	4.34 (-3.85%)	5.28 (-0.11%)	6.03 (-3.63%)	7.18 (-3.25%)
PSAX-PV	Cardiac Dreamer	7.59	6.58	4.71	5.47	5.67	8.54
	$+$ Our Pre-train	7.49 (-1.38%)	6.56 (-0.28%)	4.66 (-1.08%)	5.46 (-0.21%)	5.52 (-2.73%)	8.21 (-3.88%)
PSAX-MV	Cardiac Dreamer	7.81	6.42	4.89	6.68	5.87	9.11
	$+$ Our Pre-train	7.51 (-3.89%)	6.42 (0.04 %)	4.77 (-2.53%)	6.55 (-1.92%)	5.78 (-1.63%)	8.78 (-3.67%)
PSAX-PAP	Cardiac Dreamer	7.18	6.07	4.43	6.35	5.39	8.53
	$+$ Our Pre-train	6.92 (–3.57%)	5.97 (-1.63%)	4.35 (-1.78%)	6.11 (-3.88%)	5.33 (-1.04%)	8.42 (-1.31%)
PSAX-APEX	Cardiac Dreamer	6.98	5.98	4.25	5.69	4.89	7.33
	$+$ Our Pre-train	6.85 (-1.88%)	5.77 (-3.41%)	4.11 (-3.28%)	5.55 (-2.43%)	4.90 (0.07%)	7.25 (-1.11%)
A4C	Cardiac Dreamer	7.72	7.17	5.45	5.64	4.89	8.91
	$+$ Our Pre-train	7.55 (-2.14%)	7.00 (-2.49%)	5.36 (-1.70%)	5.48 (-2.84%)	4.86 (-0.55%)	8.60 (-3.43%)
A5C	Cardiac Dreamer	7.38	6.65	5.54	5.83	5.91	12.03
	$+$ Our Pre-train	7.15 (-3.19%)	6.46 (-2.85%)	5.40 (-2.41%)	5.80 (-0.49%)	5.88 (-0.54%)	11.88 (-1.18%)
A3C	Cardiac Dreamer	7.21	6.52	5.21	5.92	6.29	9.81
	$+$ Our Pre-train	6.90 (-4.34%)	6.34 (-2.81%)	5.08 (-2.44%)	5.77 (-2.63%)	6.12 (-2.74%)	9.52 (-2.98%)
A2C	Cardiac Dreamer	7.28	6.89	4.95	8.47	5.46	14.51
	$+$ Our Pre-train	7.04 (-3.32%)	6.70 (-2.84%)	4.83 (-2.35%)	8.37 (-1.20%)	5.33 (-2.45%)	13.98 (-3.66%)

Figure 4: Ablation of the pre-training objectives. The figure shows the relative change in MAE across six degrees of freedom for ten standard views, comparing different pre-training objectives with Cardiac Dreamer [9]. Smaller values indicate better performance. (a) Our proposed 2D-3D Joint Structure-aware pre-training. (b, c) Pre-training focused only on 2D or 3D structures.

Comparison with SOTA. To validate that the proposed structure-aware pre-training method benefits the acquisition of echocardiogram, we conducted comprehensive evaluations on the downstream probe guidance task. As shown in Tab.1, the pre-trained model consistently achieved better or comparable results across all dimensions in the ten standard views. Notably, the highest observed improvement observed was up to 4.34%. The pre-trained model demonstrated slight weakness compared to the Cardiac Dreamer in only one dimension for the PSAX-MV and PSAX-APEX planes. Despite this minor shortfall, the overall performance indicates that the model has effectively learned valuable information about the 2D and 3D structures of the heart. This acquired knowledge enhances the precision of probe guidance during cardiac ultrasound scanning tasks, thereby potentially supporting less experienced sonographers in acquiring high-quality echocardiograms in the future.

Ablations. To demonstrate the importance of each component in our proposed joint 2D-3D modeling approach, we decoupled the 2D and 3D modeling, either focusing solely on 2D structure or 3D structure modeling. As shown in Fig.4, our method achieves improvements in almost all dimensions, whereas using only 2D or 3D modeling results in poorer performance in rotation dimension. In summary, while either 2D or 3D modeling alone enhances the model’s performance to some extent, combining both in joint pre-training achieves the best results. This conclusion aligns with practical experience as well. First, 3D modeling helps in understanding the 3D structure of the heart, while 2D modeling enables the model to more accurately identify anatomical landmarks on standard views, thereby providing crucial guidance for probe positioning adjustments.

4 Conclusion and Discussion

In this work, we propose a 2D-3D joint structure-aware pre-training framework to enhance the cardiac world model’s understanding of spatial relationships within two-dimensional structures on a single view and the three-dimensional spatial relationships between different views. We innovatively designed a self-supervised learning task that predicts visual features based on both two-dimensional and three-dimensional spatial information. To support large-scale self-supervised pre-training, we collected over a million ultrasound image and 3D pose data pairs. After pre-training on the large-scale dataset, considerable improvement were observed in downstream probe guidance tasks across the ten standard views.. In the future, we will attempt to validate our probe guidance model in real clinical settings, aiming to directly translate algorithmic improvements into enhanced medical outcomes or increased efficiency. Additionally, this algorithm has the potential to serve as the decision-making brain for autonomous ultrasound scanning robots, promoting the realization of fully autonomous cardiac scanning.

4.0.1 Acknowledgement.

This work was supported in part by the National Key R&D Program of China (2021ZD0140407) and the NSFC (62321005).