Results and discussion

Spatiotemporal variability and compound-specific partitioning of ECs across basins

The dataset comprised 5085 paired sediment–water concentration records for PFASs (n = 1687), ABs (n = 1268), and EDCs (n = 2130) from 1093 sampling sites across China's seven major river basins, providing a robust basis for characterizing large-scale partitioning patterns (Supplementary Fig. S3). Inter-basin variability [48], quantified as the variance of basin-level averages of z-score standardized log_K_d, differed among EC classes (PFASs, 0.28; ABs, 0.12; EDC, 0.22; Fig. 1a–c). PFASs and EDCs displayed nearly twice the spatial divergence observed for ABs, indicating that their sediment–water partitioning is more sensitive to basin-scale environmental heterogeneity. Within-class variability remained pronounced, particularly for PFASs and EDCs (Fig. 1a–c), suggesting strong interactions with geochemical and basin-specific environments. In contrast, ABs exhibited comparatively uniform partitioning across basins, consistent with the dominant role of intrinsic molecular interactions that are less affected by basin-specific settings. The statistical results implied that assuming a uniform sediment–water partitioning mechanism, even within a single EC class, can mask environmentally driven variability and undermine predictive accuracy [49].

Fig. 1: Basin- and season-dependent variation in sediment–water partitioning. a,c,e, Distributions of log10-transformed sediment–water partition coefficients (logKd) among river basins for per- and polyfluoroalkyl substances (PFASs; a), antibiotics (ABs; c), and endocrine-disrupting chemicals (EDCs; e). Points represent individual records, violin plots show the distributions, and embedded boxplots indicate the median and interquartile range (IQR), with whiskers extending to the most extreme values within 1.5 × IQR. b,d,f, Corresponding distributions of logKd between wet and dry seasons for PFASs (b), ABs (d), and EDCs (f). Mean values are annotated; horizontal dashed lines indicate the first quartile, median, and third quartile, from bottom to top. SR, LR, HR, YR, HuR, YTR, and PR represent the Songhua, Liaohe, Haihe, Yellow, Huaihe, Yangtze, and Pearl River basins, respectively. Wet and dry seasons were defined based on a monthly precipitation threshold of 100 mm.

Seasonal stratification revealed clear temporal variability in sediment–water partitioning (Fig. 1d–f). An independent-samples test confirmed a highly significant difference in _K_d between the dry and wet seasons (p < 0.0001). Basin average log_K_d values were consistently lower during the wet season for all EC classes, reflecting enhanced dilution, increased suspended sediment transport, and greater hydrodynamic mobilization under high-flow conditions [50]. ABs exhibited the lowest seasonal variation, at approximately 6%, whereas PFASs and EDCs showed higher fluctuations, at approximately 13% and 27%, respectively. Overall, the pronounced spatial and seasonal heterogeneity demonstrated that sediment–water partitioning mechanisms varied across EC classes, requiring class-specific predictive frameworks that coupled molecular descriptors with multi-scale environmental drivers, particularly for PFASs and EDCs [51].

Cross-scale MB-MHA framework for enhanced Kd prediction

Building on the inter-class divergence and cross-basin heterogeneity identified in Section 3.1, we designed a cross-scale multi-branch multi-head attention (MB-MHA) framework to explicitly capture the multi-scale drivers underlying sediment–water partitioning (Fig. 2a). Unlike conventional ANN architectures that merge heterogeneous predictors indiscriminately into a single input stream, the MB-MHA framework adopted a mechanistic-inspired design that explicitly encoded micro- (i.e., molecular descriptors), meso- (i.e., sediment and water medium properties), and macro- (i.e., basin-scale environmental parameters) features into distinct yet interactive branches (Fig. 2a) [26]. Across all three EC classes, MB-MHA outperformed the baseline ANN, with _R_2 values of 0.76 versus 0.71 for PFAS, _R_2 = 0.92 versus 0.83 for ABs, and _R_2 = 0.89 versus 0.81 for EDCs (Fig. 2b). These results suggest that explicitly incorporating molecular-, medium-, and basin-scale variables improved model performance to characterize heterogeneous sediment–water partitioning behaviors across EC classes.

Fig. 2: Architecture and predictive performance of the multi-branch multi-head attention model. a, Architecture of the multi-branch multi-head attention (MB-MHA) framework. Molecular descriptors, medium properties, and basin characteristics are processed through separate micro-, meso-, and macro-scale branches, respectively. The branch representations are stacked and passed through a multi-head attention layer before predicting logKd. Q, query; K, key; V, value, denote the mathematical variables used in the attention mechanism. b, Predictive performance of MB-MHA and a conventional artificial neural network (ANN) for per- and polyfluoroalkyl substances (PFASs), antibiotics (ABs), and endocrine-disrupting chemicals (EDCs), evaluated using the coefficient of determination (R2) and root mean square error (RMSE). c, t-distributed stochastic neighbor embedding (t-SNE) visualization of feature representations without and with multi-head attention. Colors indicate logKd value ranges determined by equal-frequency binning for each dataset: red for low values, blue for medium values, and green for high values. The corresponding thresholds are < −0.42, −0.42 to 0.38, and > 0.38 for PFAS; < −0.58, −0.58 to 0.5, and > 0.5 for ABs; and < −0.52, −0.52 to 0.24, and > 0.24 for EDCs. Each point represents the two-dimensional projection of one high-dimensional feature representation.

This architecture mirrored the hierarchical drivers that govern EC partitioning under the MB approach, thereby embedding domain knowledge into the model. Furthermore, the integration of MHA introduced a second layer of mechanistic interpretability by enhancing the structural organization of latent embeddings [52]. Each attention head learned a complementary subspace, capturing distinct cross-scale interaction patterns, such as micro–meso coupling or macro–context modulation [42]. While t-SNE provided only qualitative visualization, clearer cluster separability was observed, particularly for PFASs and ABs (Fig. 2c), consistent with the improved predictive accuracy. By aggregating heterogeneous rational subspaces, the MB-MHA framework produced a noise-resilient representation manifold in which samples sharing similar partitioning behavior converge, while diffuse or misaligned samples were assigned to appropriate environmental clusters [39]. In conclusion, embedding domain hierarchies and attention-driven relational reasoning into the model architecture, improved the organization of heterogeneous features and the modelling of cross-scale relationships for log_K_d prediction.

Model interpretability of MB-MHA and MD-based mechanistic validation

Beyond predictive performance improvements, MB-MHA provided intrinsic interpretability by revealing how predictive reliance shifts across micro-, meso-, and macroscale drivers, as indicated by the attention matrices (Fig. 3a) [53]. For PFASs, attention heads showed different weighting patterns, including balanced weights (0.32, 0.42, and 0.40), macro-dominant weights (0.49), and macro-suppressed weights (0.11), indicating that basin-scale drivers act as context-dependent modulators influencing partitioning. ABs exhibited highly polarized weight distributions with strong microscale dominance (e.g., 0.52) and suppressed macroscale contributions (e.g., 0.11), consistent with structurally governed partitioning and weaker environmental modulation. EDCs displayed more homogeneous allocations across heads, with weights of 0.28–0.38 for microscale, 0.27–0.34 for mesoscale, and 0.31–0.37 for macroscale drivers, suggesting coordinated contributions from molecular attributes, sediment–water physiochemistry, and basin characteristics. Overall, the attention matrices showed that molecular features served as a stable backbone for prediction, while sediment–water properties and basin-scale environments acted as conditional modulators, varying with EC classes [54].

Fig. 3: Model interpretation of the multi-branch multi-head attention (MB-MHA) model. a, Attention-weight matrices for per- and polyfluoroalkyl substances (PFASs), antibiotics (ABs), and endocrine-disrupting chemicals (EDCs). Rows and columns correspond to the micro-, meso-, and macro-scale branches, and color intensity indicates attention weight. b, Branch-level contributions derived from ablation analysis (upper pie charts) and feature-level contributions derived from permutation importance analysis (lower radial charts). Percentages indicate normalized contributions to model performance; features contributing less than 0.1% are grouped as indicated. c, Partial-dependence relationships between logKd and the most influential predictor in each branch for PFAS, ABs, and EDCs. Solid lines show partial-dependence estimates, shaded areas show predictor distributions, and vertical marks indicate data density. Only predictors with permutation feature importance greater than 5% are shown. NRB, number of rotatable bonds; MW, molecular weight; HBA, number of hydrogen bond acceptors; TEMP, temperature; TOC, total organic carbon content; PREC, precipitation; UWD, urban wastewater discharge.

Ablation experiments further reinforced these insights quantitatively by measuring the decline in predictive accuracy after removing each branch (Fig. 3b) [55]. AB predictions relied heavily on microscale descriptors, with a contribution of approximately 72.1%. PFASs exhibited moderate dependence on meso-scale sediment–water properties (approximately 34.5%) and minimal sensitivity to macro-scale basin factors (approximately 5%). EDCs displayed the most evenly distributed reliance across microscale, mesoscale, and macroscale branches, at 45.4%, 26.1%, and 28.5%, respectively, consistent with their heterogeneous chemistries. Post hoc permutation feature importance (PFI) and PDP analyses further validated these trends (Fig. 3c). Molecular descriptors remained the dominant predictors for ABs, while environmental and basin variables gained prominence for PFASs and EDCs (Fig. 3b). The convergence between built-in (attention/ablation) and model-agnostic interpretations (PFI/PDP) strengthened confidence in the multi-scale attributions captured by the MB-MHA framework.

Feature-level interpretation was broadly consistent with physicochemical expectations and was corroborated by MD simulations. For PFASs, the prominence of log_K_ow reaffirmed hydrophobic partitioning as the dominant mechanism [11,56], as evidenced by the positive PDP trend between log_K_d and log_K_ow (Fig. 3c) [57,58]. Model-derived sensitivity to temperature (approximately 27%) and pH (approximately 18%) implied that PFAS partitioning is not governed solely by hydrophobic interactions but is jointly modulated by electrostatic and thermodynamic processes [16,59], rendering PFAS partitioning responsive to redox and interfacial chemistry [60]. This interpretation aligns with MD-derived interfacial interaction mechanisms. MD simulations showed that both PFOA and PFBA were stabilized through coupled electrostatic and hydrophobic interactions rather than simple hydrophobic embedding alone. In both cases, the –COO− headgroup formed Na+-bridged inner-sphere complexes with mineral surfaces, while the perfluorinated tail was associated with nearby organic phases (Fig. 4 and Supplementary Fig. S4). Compared with PFBA, PFOA exhibited stronger hydrophobic association due to its longer fluorinated chain, whereas PFBA showed relatively greater aqueous mobility (Fig. 4 and Supplementary Fig. S4). Despite these differences in interaction strength, both PFAS compounds exhibited similar ion-mediated interfacial interactions, jointly governed by ion speciation, surface charge, and Na+-bridging dynamics, which were distinct from the partitioning behaviors of ABs and EDCs [27].

Fig. 4: Molecular-dynamics analysis of contaminant interactions with humic acid. a, Representative molecular-dynamics simulation configurations for perfluorooctanoic acid (PFOA), tetracycline (TC), and bisphenol A (BPA) interacting with humic acid and the model sediment surface. b, Temporal evolution of free monomers, free aggregates, bound monomers, and bound aggregates during the simulation time. c, Interaction energies between humic acid and PFOA, TC, or BPA, decomposed into van der Waals (vdW) and Coulombic contributions. Dark-blue bars show hydrogen bond acceptors on the right axis.

For ABs, model interpretations indicated strong dependence on molecular descriptors, particularly _D_max and molecular weight [61]. The PDP curves revealed a clear threshold behavior in _D_max: log_K_d remained nearly constant for smaller molecules (_D_max < 7.1 Å) but increased sharply beyond the limit, suggesting enhanced steric compatibility with sediment matrices and the formation of stabilizing noncovalent interactions. The relationship between log_K_d and molecular weight exhibited a rise–fall pattern, with a peak at approximately 370 Da, reflecting a trade-off among molecular mobility, sorption-site accessibility, and the solvation effects of large molecules [62]. MD simulations further revealed that TC and OFX, which are representative of ABs, primarily formed free aggregates within montmorillonite interlayers. TC with a larger _D_max exhibited a closer interfacial association with sediment surfaces than OFX, supporting the ML-derived threshold behavior in _D_max. TC and OFX also exhibited stronger affinities toward humic substances than PFASs or EDCs (Fig. 4 and Supplementary Fig. S4), confirming that organic matter mediated AB retention primarily through noncovalent stabilization [63]. Nevertheless, the limited contribution of total organic carbon in the MB-MHA model suggested that organic matter primarily enhanced local aggregation microenvironments rather than exerting a quantitatively dominant influence at the meso-scale.

For EDCs, molecular flexibility, represented by the number of rotatable bonds, had the greatest micro-level impact, positively correlating with log_K_d. This pattern suggests that more conformationally adaptable molecules could better fit heterogeneous sorption sites [64]. Water COD was identified as the dominant meso-scale driver, indicating the important role of organic matter in regulating EDC partitioning. MD simulations further showed that both BPA and 4-NP remained relatively dispersed in the absence of organic matter, whereas the introduction of organic phases promoted molecular aggregation and interfacial association (Fig. 4 and Supplementary Fig. S4). At the macroscale, the impervious surface area was positively correlated with log_K_d, indicating enhanced emissions and downstream accumulation in more urbanized basins (Fig. 3).

The integrated interpretability workflow combined attention matrices, branch-level ablation, PFI/PDP analyses, and MD simulations. The complementary perspectives delineated a mechanistic spectrum of sediment–water partitioning across EC classes: ABs were primarily governed by molecular descriptors, PFASs were conditionally modulated by environmental factors, and EDCs exhibited more distributed, multiscale controls. These insights demonstrate that the MB-MHA framework improved predictive skill and mechanistically aligned explanations consistent with principles of molecular interactions and basin-scale environmental dynamics.

Future projections and spatiotemporal risk mapping of logKd in the Greater Bay Area

To evaluate short-term dynamics of representative ECs, we integrated LSTM-forecasted water quality parameters and regression-predicted basin-scale variables into the trained MB-MHA model to generate one-year log_K_d projections across the Greater Bay Area (Fig. 5 and Supplementary Fig. S5). The MB-MHA model showed high predictive performance on the test set for the input variables (_R_2 = 0.78–0.99; Supplementary Fig. S6), thereby ensuring reliable downstream spatiotemporal projections. Using PFOA, PFBA, TC, OFX, BPA, and 4-NP as representative ECs, the projected log_K_d mapping exhibited spatiotemporal variability among the EC classes. PFASs and EDCs showed more pronounced seasonal variability, whereas ABs exhibited relatively limited temporal fluctuation with log_K_d differences of only approximately 0.1 for TC and approximately 0.4 for OFX across months (Fig. 5 and Supplementary Fig. S5). This pattern aligned with the ML models: the log_K_d of ABs was primarily governed by molecular-scale descriptors, whereas PFASs and EDCs were more sensitive to variations in environmental and basin-scale conditions (Fig. 3). Distinct spatial distribution tendencies in log_K_d among ECs were also evident. PFASs displayed higher log_K_d in estuary and coastal zones throughout the year, whereas EDCs exhibited reduced partitioning in estuary zones (Fig. 5 and Supplementary Fig. S5) [4,11]. This contrast reflected their distinct partitioning mechanisms: PFAS were dominated by hydrophobic interactions that favored interfacial accumulation, whereas EDCs were more susceptible to dilution and variations in organic matter in dynamic estuarine environments [11].

Fig. 5: Seasonal predictions of sediment–water partitioning across the Greater Bay Area. Spatial predictions of logKd for representative per- and polyfluoroalkyl substances, antibiotics, and endocrine-disrupting chemicals—perfluorooctanoic acid (PFOA), tetracycline (TC), and bisphenol A (BPA), respectively—in January, April, July, and October 2026. Colors along the river network indicate predicted logKd. A separate color scale is used for each compound and held constant across the four months.

Spatial mapping of the projected log_K_d provided a basis for assessing potential environmental risks associated with contaminant accumulation and mobility. Regions characterized by persistently high log_K_d values, particularly in the central Pearl River Delta, indicate greater risks of sedimentary accumulation and long-term retention. Conversely, areas with low log_K_d values represent zones of higher aqueous mobility and downstream transport potential. These patterns suggest that management priorities should differentiate between accumulation-prone and migration-prone regions [1,2,19]. Strengthened monitoring in estuarine and urban clusters, optimized discharge scheduling, and adaptive wastewater treatment strategies are recommended to mitigate both local accumulation and cross-regional dispersion risks [65].