Three-Dimensional Magnetotelluric Forward Modeling Using Multi-Task Deep Learning with Branch Point Selection

1. Introduction

Magnetotelluric(MT) is a geophysical exploration technique widely used in practical applications, primarily for detecting the electrical structure of underground materials [1,2,3]. It is extensively applied in mineral exploration, geothermal exploration, oil and gas exploration, and deep structural studiess [4,5,6]. Forward modeling in MT is a key technique, involving the use of predefined MT models and initial boundary conditions to solve various parameters of the theoretical field [7,8].

Traditional 2-D and 3-D MT forward modeling typically employs numerical computation methods such as the finite element method [9,10,11], finite difference method [12,13,14,15,16], finite volume method [17], and integral equation method [18]. Although these traditional methods can simulate underground electrical properties with high accuracy, they require significant computational resources and time when dealing with large-scale, high grid-density problems, making them unsuitable for practical operations. In particular, for 3-D problems, the grid size is extremely large, making efficient forward modeling highly costly. Therefore, some researchers have proposed using deep learning to perform MT forward modeling to improve efficiency [19,20,21,22,23,24,25,26]. In previous studies on MT forward modeling, Yu et al. [25] in 2023 introduced Recurrent Neural Networks (RNN) and Fully Convolutional Neural Network (FCN) models, which enhanced the computational efficiency of 2-D MT forward modeling, maintaining high accuracy even when dealing with more anomalous bodies. In 2023, Hu et al. [21] proposed using a Mix Transformer network model to improve 2-D MT forward modeling, further increasing computational speed while maintaining accuracy. In 2024, Wang et al. [26] introduced a 3-D neural network based on the Swin Transformer for MT forward modeling, achieving high-accuracy forward response predictions for anomalous geoelectric models with good generalization capabilities. This study also developed a 3-D MT forward modeling dataset, laying a solid foundation for future deep learning based 3-D MT forward modeling. In 2025,Yu et al. [19] developed an advanced 3-D MT forward modeling framework incorporating a current density divergence correction for arbitrary anisotropic conditions, which optimized convergence and computational performance.

However, in the studies mentioned above, simulating the forward response of different apparent parameters required training multiple networks, which resulted in inefficient training and inference for forward modeling. To address this, a multi-task learning approach can be employed. In addition to improving efficiency, the multiple branches of a multi-task network can be trained simultaneously, sharing some common model parameters during the training process. This sharing mechanism can extract correlations between tasks and learn common features among them, thereby enhancing the model’s generalization ability (as well as efficiency) [27,28,29,30]. In previous studies on multi-task learning, Zhang et al. [31] used YOLO-based object detection branch in the Unet network, adding building semantic segmentation layers and object detection layers, achieving high accuracy. Shan et al. [32] conducted research on 2-D MT multi-task learning, designing a multi-task network based on a Convolutional Neural Network (CNN), choosing the smallest size layer after multiple convolutions as the multi-task branch point. This yielded excellent results and demonstrated the potential of multi-task learning in MT. However, while these researchers successfully applied multi-task learning, methods for selecting multi-task branch points remain scarce. The total number of network parameters increases linearly with the number of branch parameters. Consequently, when each branch has a large parameter count, the overall network size becomes substantial. This is particularly challenging for 3-D MT forward modeling, where there are multiple apparent parameters and large data sizes, leading to high GPU demands for multi-task learning methods. If the network’s terminal layer is directly used as a branch point to reduce branch parameters, significant differences between tasks can cause the intermediate layers to learn conflicting features, resulting in decreased accuracy. In 2016, Lu et al. [33] studied branch points in multi-task learning, starting with the network’s terminal layer as a branch point and dynamically expanding the branches. They found that a compact network was more efficient compared to directly using the network’s terminal layer as a branch point.

To address the issues mentioned above, this paper proposes a method for selecting task branch nodes in multi-task learning. By increasing the shared representation layer of the network, the utilization of the dataset is improved, which in turn reduces the total number of network parameters, shortens training time, and enhances the network’s generalization and prediction accuracy. In the field of 3-D MT forward modeling, this paper designs a neural network model based on multi-task learning (MT-FeatureNet). To balance the learning weights between tasks in multi-task learning, we integrate a loss function based on heteroscedastic uncertainty. Experimental results show that, compared to traditional single-task network, this method outperforms in terms of computational resources, training time, and prediction accuracy. The main contributions of the paper are as follows:

• Achieved Synchronous Inference of 3-D MT Apparent Resistivity and Impedance Phase: We introduced multi-task learning into 3-D MT deep learning forward modeling, constructing a 3-D neural network with multiple branch decoders targeting the multi-apparent parameters of 3-D MT forward responses. First, we analyzed the characteristics of apparent parameters in different polarization directions of MT. Second, we explored the feature extraction capability and forward modeling effect when different depths of the decoder layers were used as multi-branch nodes. Finally, considering both network size and forward modeling effectiveness, we selected an appropriate branch point for training. The proposed network model reduced the loss value compared to traditional single-task networks, significantly improving all evaluation metrics and achieving faster convergence.
• Proposed a Method for Selecting Optimal Multi-task Branch Points through Feature Map Comparison: We constructed a multi-task learning network model. On the branch nodes of the decoder, we compared and analyzed the feature maps at each layer in single-task network models to reduce the number of training parameters and increase the shared representation layer to enhance model generalization. Based on this, we selected the appropriate network layer for multi-task branching. Compared to traditional multi-task networks, the proposed network model further reduced the loss value, achieved faster convergence, and decreased the consumption of computational resources and time.
• Introduced a Loss Function Suitable for 3-D MT Multi-branch Decoding (Multi-task Learning): We integrated a loss function based on heteroscedastic uncertainty to balance the learning weights between tasks, avoiding the potential issue of negative loss values found in traditional uncertainty-based loss functions, thereby improving the training accuracy for each task.

2. Related Research(3-D Magnetotelluric Single-Task Forward Modeling Based on Deep Learning)

The main purpose of MT forward modeling is to calculate the response of electromagnetic waves at different frequencies (such as electric field, magnetic field, apparent resistivity, and impedance phase) based on a known subsurface medium model. This provides a basis for subsequent inversion and geological interpretation [34,35,36].

There has been considerable research on the application of deep learning in 3-D MT forward modeling, especially for single-task training focused on the physical parameters of apparent resistivity and impedance phase in the XY and YX polarization directions [22,26]. This method integrates deep learning with 3-D MT forward modeling, using extensive training to fit forward modeling parameters. Thanks to the powerful learning and representation capabilities of deep learning, this approach can quickly and accurately achieve results highly similar to those of traditional forward modeling algorithms. However, research on multi-task learning in the field of 3-D MT is relatively scarce. Figure 1 provides a simple comparison between single-task and multi-task networks, demonstrating that constructing a multi-task network can efficiently solve multiple forward modeling response parameters simultaneously, further improving the efficiency of MT forward simulation.

Currently, deep learning based on deep convolutional neural networks performs very well. Since our research focuses on establishing a multi-task model, we use U-Net3D as the backbone network, which can improve research efficiency, and the results can be extended to other similar models. The U-Net3D network model used in this study is shown in Figure 2. U-Net3D consists of an encoder, decoder, bottleneck, skip connections, and an output layer [37]. The encoder is composed of multiple convolutional blocks, each containing a convolutional layer, an activation function (ReLU), and a pooling layer. The pooling layer is used for downsampling, gradually reducing the size of the data feature maps while increasing the number of channels, allowing for higher-level and more abstract feature learning.

The bottleneck is the part that connects the encoder and decoder, focusing on further feature extraction with additional convolutional layers. The skip connections pass feature maps from the encoder to the decoder, facilitating the retention of more detailed feature information when the decoder restores the data dimensions. The decoder stitches together the features from the skip connections and then gradually restores the input data dimensions through transposed convolutions. The output layer uses convolutional layers to compress the feature maps back to the input size, followed by the Sigmoid activation function to normalize the data.

In this study, U-Net3D was used as the backbone network to separately train the apparent resistivity and impedance phase in the XY and YX polarization directions. This served as a basis for comparing the multi-task learning network model designed in this study.

3. Three-Dimensional Magnetotelluric Multi-Task Forward Modelling

3.1. Overall Concept

In multi-task learning, the network structure includes a feature-sharing component and task-specific components for each subtask. Multi-task learning requires learning from the same dataset while extracting different features. An efficient multi-task learning model should balance both the shared feature part and the task-specific parts. In previous multi-task learning research, they often directly branched multiple decoders after the encoder [38,39], similar to the U-Net architecture. However, this approach has two significant limitations: (1) it often results in an excessive number of parameters, leading to inefficient model training and inference; and (2) the branching point may not be optimal, as the shared features extracted at the encoder’s output may not be equally suitable for all tasks. This study conducted extensive research on selecting the branching points for multi-task learning. Based on this, multiple multi-task learning networks were trained, and comparative experiments were conducted to validate the theoretical foundation proposed in this paper.

In multi-task learning, the choice of the loss function is also a crucial component. How to set the learning weights for each task is a significant issue, as an unreasonable weight design can lead to imbalanced learning among the subtasks. If the model focuses too much on tasks with higher weights, it may neglect those with lower weights. Multi-task learning relies on the shared information in the shared layers, and poorly designed weights can cause interference between tasks, affecting the quality of the shared information and the overall learning process.This study utilizes a loss function based on homoscedastic uncertainty to dynamically adjust the weights in multi-task learning. This allows the model to better learn each task, balancing the importance of tasks and improving overall performance.

3.2. Network Structure

The neural network architecture designed in this study is illustrated in Figure 3. This network model consists of a shared encoder, two polarization-direction shared decoders, and four independent decoders. The shared encoder is first split into two parts for each polarization direction, and then each polarization-direction decoder is further split into two, forming the overall structure.The shared encoder is based on a U-Net3D-style upsampling encoder design and also includes parts of a downsampling decoder. It features 3-D convolution modules, ReLU activation function modules, MaxPooling3-D pooling modules, and 3-D transposed convolution modules. The polarization-direction shared decoders use the same modules. The independent decoders additionally incorporate a Sigmoid activation function module.Although the structure of each parameter training module is consistent, they use different datasets, each loaded with data corresponding to parameters in different polarization directions.

In traditional multi-task U-Net3D network models, there’s usually one shared encoder and independent decoders equal to the number of tasks. The shared encoder handles the downsampling (similar to the U-Net3D encoder), while the independent decoders handle the upsampling. In this network, the presence of multiple independent encoders and the use of 3-D data significantly increase the number of training parameters, consuming substantial computational resources and time.Another common approach is to set the multi-task branching point at the end of the network. However, this results in minimal task-specific layers, limiting the ability to optimize each task individually. Additionally, if the shared layers are too extensive and tasks differ greatly, learning conflicts can arise, leading to inefficient learning and reduced accuracy.Thus, choosing an appropriate branching point is crucial. It not only reduces the number of training parameters and speeds up training but also allows the shared layers to facilitate complementary reinforcement in multi-task learning. Simultaneously, the task-specific layers can optimize individual features, enhancing the model’s generalization capability. To identify a balanced point for the multi-task branching, we first pre-trained four traditional single-task networks, each corresponding to one of the four parameters. This allowed us to analyze the feature maps output by each convolutional layer and examine the parameters at each level. Figure 4 shows the feature maps for images A, B, and C in the traditional single-task networks.The analysis revealed that the feature maps in A were relatively similar. In B, the feature maps for apparent resistivity and impedance phase in the same direction were quite similar, while the differences between directions were more pronounced. For C, the feature maps across parameters displayed significant variations.Based on this analysis, we selected A as the network branching point to avoid the training inefficiency that might occur due to the large differences in feature maps seen in B and C. From B, it’s clear that the feature maps for the two parameters in the same direction are quite similar, so we split B into two tasks for the XY and YX directions to maximize the utilization of shared parameters. In the next layer, we then generate branches for apparent resistivity and impedance phase in each direction.

3.3. Loss Function of Multi-Task Learning

In traditional multi-task learning, one core issue is how to optimize a model that includes multiple loss functions. A common approach is to directly linearly combine each loss function [40]. The formula for this can be defined as follows:

$$L_{\mathrm{total}}=\sum _{\mathrm{i}}{w}_i{L}_i$$

(1)

In this formula, L represents the loss value, and $\omega$ is the weight assigned to each task. While this straightforward approach can work well when the weights are set appropriately, the training performance can be significantly impacted if the weights are not properly configured. When the weights are well-chosen, the model can train effectively and produce high-quality results. However, if the weights are poorly selected, the training performance can be severely compromised.Thus, determining the appropriate weight values is a crucial and complex task [41]. It requires extensive experimentation and comparative analysis to find the optimal weight configuration. This process is computationally expensive, time-consuming, and highly challenging.

Alex Kendall et al [42]. proposed using homoscedastic uncertainty to adjust the weight coefficients in multi-task learning, which yielded excellent results. Homoscedastic uncertainty is an important type of uncertainty in Bayesian modeling. For regression-type tasks, this can be expressed as follows:

$$L_{\mathrm{total}}=\sum _i{\displaystyle \frac{1}{2{\sigma }_i^2}}{L}_i+log{\sigma }_{\mathrm{i}}$$

(2)

In this approach, $\sigma$ represents the observation noise, and it is trainable. L is the loss value. We can interpret that as $\sigma$ increases, the target loss function decreases, and the final term acts as a regularization term to prevent the loss from becoming zero.

By implementing this loss function, we can effectively balance the weights of each task without manually setting the weight parameters. However, during multiple training sessions, we found that the final loss value could become negative due to the regularization term $log\sigma$ potentially being negative. Therefore, we need to modify the formula to ensure that the final loss value is greater than zero. The first part of the formula is already greater than zero, but we need to ensure that the regularization term at the end is also positive. To do this, we can square $\sigma$, add one, and then take the logarithm. The final formula we designed is as follows:

$$L_{\mathrm{total}}=\sum _i{\displaystyle \frac{1}{2{\sigma }_i^2}}{L}_i+log\left({\sigma }_{\mathrm{i}}^2+1\right)$$

(3)

The regularization term is always greater than zero, which prevents the loss value from becoming negative and ensures that the multi-task model is trained in the correct direction.

4. Results and Analysis

4.1. Dataset

The training dataset used in this paper is provided by Wang et al [26]., and it contains 3208 3-D data samples, each with forward modeling data labels for both the XY and YX directions. The test validation set consists of 211 3-D data samples, also with forward modeling data labels for both XY and YX directions. The dataset is categorized into three types based on the number of anomalies: one anomaly, two anomalies, and three anomalies.

4.2. Data Preprocessing

For deep learning networks, reducing the range of data values helps the network find the correction points and fit the data quickly. Therefore, during the training process, we map all data to the range of (0, 1). For geoelectric models with a large range of values (ranging from 0.1 to 1000) and apparent resistivity data (ranging from 0.01 to 2000, though most values are below 500), we first apply a logarithmic transformation, followed by min-max normalization. For impedance phase data with a smaller range (ranging from 0 to 90), only min-max normalization is applied.The formula for min-max normalization is as follows:

$$x^{\prime }=\left(x-x_{min}\right)/\left(x_{max}-x_{min}\right)$$

(4)

For the network’s output results, the inverse of the above steps will be applied to restore the data to its original scale, allowing for the calculation of various quantitative metrics. This involves reversing the normalization process and logarithmic transformation to return the predicted values to their original range and units.

4.3. Experimental Environment and Parameter Settings

In the experiments, the batch size for both the training and test datasets is set to 8. The AdamW optimizer is used to optimize the model parameters, with the learning rate set to 0.0002, and the maximum number of training epochs is set to 120. The algorithm is implemented using the GPU version of the PyTorch library, which accelerates model training. Additionally, a dynamic learning rate adjustment algorithm is designed during training to prevent gradient vanishing, speed up convergence, and improve training stability.

During network training, we load the pre-trained one-to-one network parameters and assign corresponding parameters to the shared layer, XY independent encoder, and YX independent encoder for training. This approach accelerates the network’s convergence and enhances the model’s learning efficiency.

4.4. Evaluation Indicators

In this paper, the experimental accuracy of the NN forward modeling is quantitatively evaluated using SSIM, MRE, and MAE. SSIM focuses on the structural similarity of data and is well-suited to reflect the local structural variations and overall feature similarity of anomaly responses in 3-D MT data. The SSIM formula is as follows:

$$SSIM(x,y)=\underset{C_1={\left(K_{1}L\right)}^2,C_2={\left(K_{2}L\right)}^2}{\underbrace{{\displaystyle \frac{\left(2{\mu }_x{\mu }_y+C_1\right)\left(2{\sigma }_{xy}+C_2\right)}{\left({\mu }_x^2+{\mu }_y^2+C_1\right)\left({\sigma }_x^2+{\sigma }_y^2+C_2\right)}}}}$$

(5)

In the SSIM formula:x represents the ground truth values,y represents the network predicted values,${\mu }_x$ and ${\mu }_y$ are the mean luminance values of x and y,${\sigma }_x$ and ${\sigma }_y$ are the variance of luminance values of x and y,${\sigma }_{xy}$ is the covariance between the luminance of x and y,$C_1$ and $C_2$ are constants to stabilize the division and avoid division by zero,L is the actual maximum range of parameter values,$K_1$ and $K_2$ are constants controlling the sensitivity to contrast and brightness, typically set to 0.01 and 0.03, respectively.

MRE and MAE exhibit similar trends, but differ in their numerical ranges. MRE focuses on the proportion of the network’s prediction error relative to the true values, reflecting the reliability of the network’s predictions.The formula for MRE is as follows:

$$MRE(x,y)=\sum _{i=0}^N\left({\displaystyle \frac{\left|y_i-x_i\right|}{x_i}}\right)$$

(6)

In the MAE formula:x represents the true values,y represents the predicted values,N represents the total number of data points.

MAE focuses on the average of the absolute errors between the predicted values and the true labels, reflecting how much the predicted values deviate from the actual values. Combining MAE with MRE provides a more comprehensive analysis of the accuracy of the network’s forward modeling results.The formula for Mean Absolute Error (MAE) is as follows:

$$MAE(x,y)=\sum _{i=0}^N\left(\left|y_i-x_i\right|\right)$$

(7)

The parameters in the MAE formula have the same meaning as those in the MRE formula.

In addition to quantitative metrics, this paper also compares visual effects, such as the variation curves of the apparent resistivity values and impedance phases on 3-D MT slices. These visualizations help intuitively present the final results of the 3-D deep learning network’s forward modeling, and explore its advantages. The visual comparison allows for a clearer understanding of the network’s performance and the effectiveness of the proposed model in capturing and representing the underlying physical features of the MT data.

4.5. Comparative Results of Branch Point Selection Strategies

To validate our hypothesis that the designed multi-task network achieves optimal branching when starting from layer A to split a single network path into two independent directional paths and further splitting each of these paths at layer B, we trained two multi-task networks with branching points at layers A and B, respectively. For each network, pre-trained parameters were loaded into the shared layers. Table 1, Table 2, and Table 3 below compare the training performance on the validation set (in terms of SSIM, MAE, and MRE) for three architectures: (1) a traditional multi-task network where layer A directly splits into four paths, (2) branching initiated at layer A, with further splitting at layer B, and (3) branching initiated at layer B, with further splitting at layer C.

As can be seen, using layer A as the initial branching point achieves notable improvements in all three metrics (SSIM, MAE, and MRE) compared to both the traditional multi-task network and the network with layer B as the initial branching point. This suggests that selecting layer A as the initial branching point optimizes the multi-task learning process, leading to better performance in terms of both structural similarity and prediction accuracy.

Figure 5 below shows the comparison of validation set losses between the network structure designed in this paper, which uses layer E as the branching point and loads the pre-trained parameters in a one-to-one configuration, and the traditional multi-task learning model using U-Net3D. The losses have been logarithmically scaled for better visualization. As can be seen, the network structure proposed in this paper fits much faster and significantly reduces the loss, thereby improving computational accuracy. Compared to the traditional multi-task learning model, the loss is reduced by 21.96%. This demonstrates the efficiency and effectiveness of the proposed network structure in minimizing the loss and achieving faster convergence.

4.6. Comparative Results of Multi-Task and Single-Task Learning

To evaluate the performance of the network designed in this paper, we used SSIM, MAE, and MRE to measure the quality of the training structure at each epoch. Table 4, Table 5, and Table 6 show the results of SSIM, MAE, and MRE at the last training epoch. The results demonstrate a significant improvement, with both MAE and MRE values much smaller than those of the U-Net3D network. Additionally, SSIM approaches 1, while further improvements in accuracy are observed. These results highlight the effectiveness of the proposed network in improving both the structural similarity and the prediction accuracy compared to the traditional U-Net3D network.

Figure 6, Figure 7, and Figure 8 show the anomaly models under different types, along with the outputs of the multi-task network predictions, model labels, and the comparison of the curves for the labeled data, multi-task network predictions, and single-task network predictions for apparent resistivity and impedance phase. Since deep learning already achieves excellent accuracy in 3-D single-task networks, the differences between the single-task and multi-task networks in the 3-D images are minimal. Therefore, only the comparison between the multi-task network 3-D images and the labeled images are shown, while the curve graphs compare the labels, multi-task network predictions, and single-task network predictions for each parameter.From the comparison, it is evident that the network designed in this paper has excellent generalization ability, capable of handling different numbers of anomalies. Regardless of the number of anomalies, the model label curves and the network output prediction curves closely align, achieving outstanding results with no overfitting. Additionally, the multi-task network curves are even closer to the label curves compared to the single-task network, demonstrating that the use of multi-task learning has improved prediction accuracy.

4.7. Comparative Results of Loss Function Performance

The following Table 7, Table 8, and Table 9 present the experimental results comparing SSIM, MAE, and MRE using the Mean Squared Error (MSE) loss function and the uncertainty-based homoscedastic loss function.

It can be observed that the homoscedastic uncertainty-based loss function has a significant advantage over the MSE loss function in terms of SSIM, MAE, and MRE metrics. Figure 9 shows the comparison of the validation set loss between the network structure using the improved uncertainty loss function (Uncertain) and the MSE loss function. The loss values are presented on a logarithmic scale, and it can be seen that the uncertainty loss function achieves a lower final loss compared to the MSE loss function, reducing the loss by 16.35% relative to the MSE loss.

4.8. Time Consumption Comparison

In terms of training time, Table 10 compares the average training time of the one-to-one network with that of the multi-task network.

Theoretically, under the same hardware conditions, training four parameters with the traditional single-task network requires 682 seconds per training round. In contrast, the multi-task learning network structure designed in this paper reduces the time to 466 seconds per round. With a maximum of 120 training rounds, this results in a total reduction of 25,920 seconds, or approximately 7.2 hours. This demonstrates that the network structure designed in this paper not only significantly reduces computational resource consumption but also saves a considerable amount of training time, which has significant implications for both theoretical research and practical applications.

5. Discussion

Through experiments with data from Section 4, the following understanding can be gained.

• For our feature map-based branch point selection strategy, it has been very effective and significantly improved the prediction accuracy. In Figure 4, we conducted a comparison of network feature outputs, and ultimately chose to split from layer A in both XY and YX directions, then further divided each direction into apparent resistivity and impedance phase branches. Table 1, Table 2, and Table 3 present comparisons of different branch selection strategies across various metrics (SSIM, MAE, MRE), where our branch point selection strategy achieved the best performance in all three metrics with notable improvements. Furthermore, in Figure 5, we analyzed the loss per epoch for three experiments and found that our strategy maximizes resource utilization and reduces training loss. However, we observed that when using networks with a large number of complex layers, due to the increased number of potential branch points, selecting the optimal multi-task branch point from feature maps becomes complicated and challenging. Therefore, a method capable of dynamically selecting multi-task learning network branch points is a future development trend.
• By using the proposed strategy in this paper, the trained network was compared item-wise with traditional single-task networks, and it can be observed that the network proposed in this paper achieved significant improvements across all three metrics. Furthermore, based on the experimental results using an improved uncertainty-based loss function combined with the MSE loss function, Table 7, Table 8, and Table 9 show that the uncertainty-based loss function can better adjust the learning weights for each task in multi-task learning, thereby further improving the prediction accuracy, as shown in Figure 9, which also reduces the training loss. Table 10 indicates that the proposed multi-task network consumes less time than the total time of four single-task networks. This reduces the time cost and maintenance cost for practical tasks, making it more efficient.

6. Conclusion

Recent research has shown that traditional 3-D MT forward modeling is not only feasible but also exhibits significant effectiveness. However, Traditional 3-D MT forward modeling is time-consuming, and training multiple parameters requires the use of several separate networks, which significantly increases computational resource consumption and training time. With the growing progress in multi-task learning, it has proven to outperform single-task frameworks by leveraging shared representations and optimizing data utilization. Despite these advantages, applying multi-task learning to 3-D MT forward modeling poses significant challenges, particularly in terms of hardware requirements. Multi-task networks typically involve more parameters compared to single-task networks, demanding higher computational power and memory capacity to achieve efficient training and inference. In this paper, we propose a multi-task learning neural network model (MT-FeatureNet) with custom branching, which enhances the shared representation layer of the network while retaining the independent layers of single-task networks. This reduces the uncertainty to dynamically adjust the learning weights between tasks, effectively minimizing negative interference between tasks. Ultimately, through comparisons of SSIM, MAE, and MRE, our model achieves significant improvements over traditional single-task networks, with a 21.96% reduction in loss, while also effectively reducing training time and enhancing training accuracy. This demonstrates that the custom network branching approach proposed in this paper is highly effective, offering a novel perspective for designing multi-task learning network structures.