To address the complexity of geological setting and the large volume of geological data inherent in 3D geological modelling, we employ spatial-domain GAT and GraphSAGE to aggregate node features, thereby harnessing the strong nonlinear integration capabilities of GNNs.

The tetrahedral meshing of the sampled points is performed to create graph structured data. All the sampling points and the other densifying points are both included into the graph data. As shown in Fig. 2, the base and the lateral boundaries are determined by the modeling extent whereas the top boundary is determined by the digital elevation model (DEM). Each sampling point is a vertex of a tetrahedron, and the topological connection of the sampling point is maintained during tetrahedron subdivision. Densifying points are created based on the given modelling resolution; overlapping points are eliminated. The resulting mesh has a multi-resolution structure, with tetrahedra around sampled points having smaller volumes and higher densities. Graph nodes are represented as tetrahedral vertices, and graph edges are represented as tetrahedral edges. Finally, the faults are encoded as a feature of nodes in relation to fault planes, which is combined with other node features as model inputs (Gao and Wellmann, 2025).

The multi-scale graph neighborhood aggregation mechanism is formulated into the proposed framework. GAT with its learnable attention mechanism gives a specific weight to the neighboring nodes (Veličković et al., 2017), which is especially useful in a geological modeling task where local heterogeneity and sensitivity to the boundaries are considered to be a significant issue. When processing nodes near faults, GAT learns to down-weight connections to neighbors on the opposite side of the fault, effectively preventing spurious propagation of stratigraphic samples across the discontinuity. At stratigraphic boundaries, GAT preferentially attends to aggregate nodes within the same unit, allowing for sharper delineation of geological discontinuities. For node i and its neighbor node j, the attention coefficient αij is calculated, which reflects the importance of node j in propagating features to node i, as follows: 1αij=exp⁡(LeakyReLU(aT[Whi‖Whj]))∑k∈N(i)exp⁡(LeakyReLU(aT[Whi‖Whk])) where a is the learned attention parameter vector, ‖ represents the concatenation of feature vectors, W is the weight matrix, and hi and hj are the feature representations of nodes i and j, respectively.

GraphSAGE statistically incorporates neighboring node features into the central node through a mechanism primarily designed to capture regional geological characteristics and spatial continuity patterns (Hamilton et al., 2017). For scalar value prediction tasks, this approach demonstrates robust capability in capturing spatial gradient patterns, given that scalar properties within a continuous domain generally exhibit gradational transitions rather than abrupt discontinuities. In lithology classification tasks, GraphSAGE exploits the geological principle that spatially proximate nodes tend to share similar features, thereby incorporating the lithological distribution within the statistical neighborhood to inform classification decisions at each target node. GraphSAGE aggregates the features of neighboring nodes layer by layer through a multi-layer neural network, with the embedding update formula for each layer's nodes, as follows: 2hvk=σ(Wk⋅AGGREGATEk({h1k-1,∀u∈N(v)})⊕hvk-1)(k=1,…,K) where hvk is the embedding vector of node v at the k layer, N(v) is the neighborhood of node v, and AGGREGATEk is the aggregation function of the k layer. Common choices for aggregation functions include pooling aggregation and long short-term memory (LSTM) aggregation. Given the similarity of geological features within the same region in geological modeling, the average pooling aggregation function is used in this study to embed the features. Wk represents the learnable weight matrix, and ⊕ denotes the vector concatenation operation.

Geological data are distinguished by a high degree of spatial heterogeneity, thereby similar characteristics in different geological environments may have entirely distinct meanings. Such domain-specific subtleties are not reflected in conventional attention mechanisms, which only compute weights using feature similarity. To address this weakness, we present a stratigraphy-aware attention (SAN) process that involves the prior constraints of stratigraphic sequences. Through the incorporation of lithological similarity and stratigraphic sequence, the mechanism gains a meaningful combination of stratigraphy knowledge for the proposed framework (Fig. 3).

First, the SAN receives the following inputs: features aggregated from the graph neighborhood, query vectors Q derived from graph nodes, observation point features K from the graph, and feature V used for aggregation. The input features undergo an initial linear transformation that projects them into the attention space. Secondly, the projected features are partitioned into H attention heads, each processing dk=d/H dimensional features. For nodes i and j, the base score of the h head is: 3Sbase,ij(h)=1dk∑l=1dkQi,h,l⋅Kj,h,l We first define the mapping function from lithology to geological age as ai=Age(ri), as shown in Fig. 4. For node (i,j), the output Sage(h) denotes the value for stratigraphic age similarity between node (i,j) at the h head, as follows: 4Sage(h)(i,j)=y([aiaj|ai-aj|]) where y(⋅) is the age mapping function.

Then, we define a learnable lithological similarity matrix L∈RK×K, where K is the number of lithological categories. Adjacent strata show stronger lithological similarity, with stratigraphic quantitative difference correlating inversely to lithological resemblance. For nodes i and j, lithological similarity Sseq is: 5Sseq(h)(i,j)=σ(Lri,rj)if 1≤ri,rj≤Kcotherwise where σ(x)=exp⁡(-d(i,j)/τ), and d(i,j)=|i-j| denotes that the quantitative difference between the ordinal numbers of two lithological categories within the stratigraphic sequence, and τ is a stratigraphic decay length governing how rapidly lithological affinity is lost with increasing stratigraphic separation. Setting τ=3 assigns adjacent units a similarity of e-1/3≈0.72 and units three steps apart e-1≈0.37, which is consistent with the initialized similarity structure recovered in Fig. 5.

When predicting scalar values, for the h attention head, the enhancing attention for node pair (i,j) is as follows: 6As_ij(h)=softmaxSbase,ij(h)+Sage(h)+Scontact(h) where Scontact(h) is the similarity of adjustable contact relationship.

When lithology classification, for the h attention head, the enhancing attention for node pair (i,j) is as follows: 7Al_ij(h)=softmaxSbase,ij(h)+Sseq(h) The SAN mechanism begins by computing standard base attention scores through linear projection of node features and their division into multiple attention heads. Its key innovation is the concurrent incorporation of stratigraphic sequence domain knowledge as prior constraints. For any node pair i,j, where Sbase is the standard query key feature similarity used by conventional attention, Sage is an explicit function of the geological ages of the two nodes' lithologies, Sseq is the learnable lithological similarity that decays with ordinal stratigraphic distance, and Scontact encodes the conformable or unconformable contact relationship. Because these terms are added before the softmax, each attention weight admits an exact additive attribution into a learned feature similarity component and named stratigraphic prior components; the contribution of every prior to any aggregation is therefore transparent and auditable. The two prior terms are monotone in geological proximity. Sage and Sseq raise the weight for node pairs that are close in the stratigraphic sequence and lower it for distant pairs. After the softmax, aggregation is thus biased toward stratigraphically coherent neighbours irrespective of incidental feature resemblance. When Sbase is ambiguous (two candidate neighbours with similar features but on opposite sides of a fault or across an unconformity), the prior terms dominate the sum and steer attention to the geologically correct neighbour, which is the mechanism by which cross fault propagation is suppressed and within unit aggregation is sharpened.

Following the extraction of spatial perceptual features by the SAN, a dual-task learning head is employed to simultaneously perform lithology classification and scalar field prediction, as illustrated in Fig. 6.

The two branches are coupled at three levels rather than merely sharing a backbone. (1) Reading the same SAN-refined node embedding, so gradients from the two tasks jointly shape the shared representation, forcing features to be simultaneously predictive of the scalar value and of the lithology. (2) The stratigraphy-aware attention injects lithological similarity, geological-age and contact similarity, so the aggregation is tied to stratigraphic structure for both tasks. (3) The cross-task loss explicitly links the two predictions, the KL term (soft) aligns the classifier distribution with the distribution implied by the predicted scalar field, and the compatibility term (hard) penalises scalar predictions that fall outside the admissible interval of the predicted lithology.

Geo-SAN predicts the scalar field fv, which is constrained by stratigraphic interfaces and orientation points like dip angle and strike. We can define 3D space as a scalar function f(v), where represents the scalar field values at each point v in 3D space. A series of H stratigraphic interfaces can be articulated as: 8f1>⋯>fH where f1 to fH represent the scalar values of the oldest to newest stratigraphic interfaces, respectively.

For a certain stratigraphic interface I with a scalar field value fI, the subset of graph nodes sampled from it can be denoted as: 9I={v∈V|fv=fI} The loss functions Lscalar can be used to find the difference between the scalar field at the stratigraphic interface sampling point and the predicted scalar field by Geo-SAN, as follows: 10Lscalar=1N∑i=1N(f^i-fi)2

Nodes with attitudes can be extracted from geological maps, which represent the local morphology of the strata at these nodes. A unit vector n in the normal direction of the stratigraphic interfaces can be transformed from the attitudes: 11n=sin⁡θ⋅sin⁡δcos⁡θ⋅sin⁡δcos⁡δ where θ and δ are the strike and dip angles of the stratum, respectively.

To measure predicted gradient and the actual direction of angular error, the gradient of the predicted scalar field is gained at these graph nodes. This study employs the first-order Taylor series approximation of the scalar field in the node's first-order neighborhood to estimate the scalar field gradient, as follows: 12PvTPv∇fv=PvTfu-fv⋮ where the angle between the gradient at node v and the known vector ∇fv=(∂fv/∂x,∂fv/∂y,∂fv/∂z) can be used to measure the angular error, and u is the neighbor of node v. 13Pv=xu∗-xv∗yu∗-yv∗zu∗-zv∗⋮⋮⋮,∀u∈Nv where x∗, y∗, and z∗ are the normalized spatial coordinates of the given node.

The estimated angle between the gradient ∇fv at node v and the known vector αv in the observed attitude node-set o is utilized to formulate the loss function, as follows: 14Lorientation=∑v∈ocos⁡∇fv,αv=∑v∈oαv⋅∇fv‖αv‖‖∇fv‖

In the lithology node-set R, comprising Nclasses lithology categories, each lithology node is associated with a one-hot encoded vector lv, denoting its respective lithology category. The loss function employed for lithology nodes is the cross-entropy loss function, as follows: 15Llitho=-∑v∈R(lv⋅log⁡(l^v)) where l^v is the predicted probability distribution for node v, and the predicted lithology category is determined based on the highest probability from the softmax distribution of l^v.

The KL divergence loss function (soft constraint) is utilized to measure the discrepancy between lithology classification and scalar field prediction, thus encouraging models to adhere to stratigraphic sequence patterns, as follows: 16LKL=KL(l^‖lf^)=∑i=1N∑j=1Kl^ijlog⁡l^ijlf^ij where l^ij denote that model's predicted probability distribution over jth lithological category at ith node, and lf^ij denote that the probability distribution derived from stratigraphic scalar field.

The compatibility loss function evaluates whether expected scalar field conform to the acceptable range for lithological category, imposing penalties on forecasts that contravene geological principles, as follows: 17Lcompat=1N∑i=1N1-exp⁡-Δf(i)F where F denotes exceeding the upper or lower threshold. The deviation Δf(i) measures how far the predicted scalar value is outside the valid range: 18Δf(i)=max⁡(0,f^i-tli)+max⁡(0,bli-f^i) where tli and bli are the top and bottom interface scalar values according to the observed lithological category at node i. When bli≤f^≤tli, then Δf(i)=0 (perfect compatibility). The cross-task loss function is expressed as: 19Lcross=LKL+Lcompat The overall loss function is as follows: 20L=αLscalar+βLorientation+γLlitho+δLcross In this study, the GradNorm (Gradient Normalization) method is introduced to optimize the training process (Chen et al., 2018). GradNorm adaptively adjusts the weights of each task's loss function to ensure that the gradients of each task have similar magnitudes during training, improving the stability and efficiency of training.