cs.LGOct 27, 2026

Simulated petrological labels enable probabilistic mineralogy prediction from whole-rock geochemistry

Authors: Matthew J. Liu, Aaron D. Kaplan, Michael L. Whittaker, Kristin A. Persson

Organizations: Materials Sciences Division, Lawrence Berkeley National Laboratory, Berkeley, CA, USA · Energy Geosciences Division, Lawrence Berkeley National Laboratory, Berkeley, CA, USA

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Modal mineralogy, the mineral identities and abundances in rocks, guides geological studies and industrial resource extraction. While measuring modal mineralogy is time-intensive, whole-rock chemical measurements furnish elemental concentrations at higher throughput. Predicting mineralogy from elemental data, however, poses a one-to-many inverse problem, and labeled mineralogical datasets are too sparse to train on at scale. We address both challenges by simulating petrological labels that express multiple formation conditions per composition and at a scale unattainable by measurements. Nearly one million simulated labels supervise the middle component of a three-part model: a transformer encoder self-supervised on roughly 788,000 whole-rock compositions, the simulation-trained heads that predict distributions over mineral assemblages, and a recalibration layer that aligns the predictions with measured mineralogy. Zero-shot on 640 volcanic rocks, the model discriminates the presence of most mineral categories and predicts which rocks are richer in each mineral. However, the predicted abundances are systematically offset from measured values. Recalibration fit on as few as 25 labeled rocks recovers 76 percent of the likelihood improvement from fitting on the full 368-rock fitting set, while abundance error responds more slowly (11.0% zero-shot and 9.5% with the full fitting set). This work demonstrates how a machine learning model in the geochemical sciences can both leverage an imperfect simulator and mitigate its systematic biases.

Figures & tables

Figure 1: Zero-shot predictions over 640 volcanic rocks. (a) Model inputs. (b) Presence AUROC per mineral category, where the dotted line represents a random classifier. Whiskers span the 25th to 75th percentile across bootstrap resamples of stratigraphic units. (c) Violin plots for predicted vs. measured mineral abundance when the heads are given the measured mineral list as their support (i.e., teacher-forced). In (a) and (c), violin width reflects the density of rocks at each value and the tick marks the median. Mineral abbreviations in Table B9 .
Figure 2: Recalibration improves mineral predictions. Zero-shot vs. recalibrated predictions on the test set rock with the median recalibrated joint NLL. (a) Model input. (b, c) Zero-shot presence probability and mineral abundance. (d, e) Analogous predictions after a 25-rock fit. Abundance panels (c and e) each summarize 2,048 candidate assemblages sampled from the presence probabilities of (b and d), respectively. Points denote the mean across candidates, shaded bands indicate the 25th to 75th percentile, and gray lines connect predicted to measured values. Mineral abbreviations in Table B9 .
Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

FamilyChallenge addressedAdopted in this work
Site-specific tabular modelsPredicting quantities that are costly to measure at an operating site–
Grain-scale classifiersAutomated mineral identification from spot chemistryStoichiometric naming rules (Appendix B.4 )
Normative and unmixing calculationEstimating mineralogy from bulk compositionThe prediction task and precedent for probabilistic output
Thermodynamic simulatorsComputing equilibrium assemblages at fixed conditionsThe label generator
Simulation-to-real transferCarrying simulation-trained models to real-world measurementsRecalibration paradigm
Appendix
Table A1: Prior families of work in geochemical prediction, the challenge each addresses in its respective field, and what this work adopts from each.
GroupCountMembersfrac
Majors9Si, Al, Ca, Mg, Fe, K, Na, Ti, Crwt% / 100
Traces13Sr, Zr, Ba, Y, Rb, Ni, Nb, V, La, Ce, Nd, Yb, Thppm / 10610^{6}
Redox1Fe 3+ /total Feratio, unscaled
Appendix
Table B1: The fixed elemental inputs read by the encoder.
StatesrcfracAttended
Padding (element not reported)00no
Populated1–118>0>0yes
Masked for pretraining1190yes
Fe-redox ratio120ratio in [0,1][0,1]yes
Appendix
Table B2: Possible states for each token of the input sequence.
HyperparameterMeaningValue
d_modelDimensionality of latent representation256
n_layersTransformer encoder layers4
n_headsAttention heads per layer8
dim_feedforwardFeed-forward width, 4×4\times d_model1024
activationNon-linear activation functionGELU
norm_firstPre-norm residual arrangementtrue
Appendix
Table B3: Encoder configuration.
HyperparameterMeaningValue
mask_probFraction of populated rows masked0.281
batch_sizeRocks per step32
lrAdam learning rate8.18×10−58.18\times 10^{-5}
weight_decayAdam weight decay1.98×10−31.98\times 10^{-3}
epochsCeiling, not a target60
early_stop_patienceEpochs without val improvement3
Appendix
Table B4: Pretraining configuration and outcome.
PhaseAbbreviationBenchmark categoryPresence rate (%)
plagioclaseplagFeldspar79.3
alkali feldsparafsFeldspar20.0
clinopyroxenecpxPyroxene76.5
orthopyroxeneopxPyroxene25.4
amphiboleampAmphibole53.3
biotitebiBiotite27.0
Appendix
Table B5: The 21-phase vocabulary of the assemblage head, with each phase’s presence rate, defined as the percent of training labels in which the phase is present (i.e., holds more than 0.1 wt% of the crystal mass). The ten phases above the divider map onto the 7 benchmark categories of Table B9 .
SplitCompositionsDraws eachMelt fractions read (wt%)Labels
Train9,9464850, 30949,228
Validation21812850, 3055,576
Test27412850, 3069,521
Appendix
Table B6: Size of the simulated label corpus by split.
HeadPredictsDistributionTrunkParameters
Assemblagephase presence and abundanceBernoulli ×\times Dirichlet3×5123\times 5121,346,090
Compositionper-phase oxide compositionDirichlet per phase1×1281\times 12865,404
ConditionsTT , PP , dissolved H 2 OGaussian mixture per axis2×5122\times 512418,843
Appendix
Table B7: The three prediction heads, where trunk width and depth were tuned by Bayesian search alongside the loss weights and optimizer settings (Appendix B.6 ).
HyperparameterMeaningRangeSelected
lrAdamW learning rate10−410^{-4} to 10−210^{-2}4.70×10−44.70\times 10^{-4}
weight_decayAdamW weight decay10−710^{-7} to 10−210^{-2}2.46×10−52.46\times 10^{-5}
λabund\lambda_{\mathrm{abund}}Abundance weight, nominal0.0048 to 0.0480.0281
λbalance\lambda_{\mathrm{balance}}Mass-balance weight10−310^{-3} to 10410^{4}3.08×10−33.08\times 10^{-3}
dropoutTrunk dropout, all heads0 to 0.40.260
mode_hiddenAssemblage trunk width64, 128, 256, 512512
Appendix
Table B8: Joint search space and optimized values over the training environment and assemblage-head architecture. The remaining weights were fixed during the search at λcomp=1.0\lambda_{\mathrm{comp}}=1.0{} and λcond=0.01\lambda_{\mathrm{cond}}=0.01{} .
FeldsparPyroxeneAmphiboleBiotiteOlivineFe-Ti oxidesQuartz
fsppxampbiolopqq
plagioclaseclinopyroxeneamphibolebiotiteolivinespinelquartz
alkali feldsparorthopyroxene   ilmenite 
Appendix
Table B9: Mapping predicted phases onto the 7 mineral categories quantified in [ 14 ] and [ 15 ] .
Reported mineralCategoryDensity (g cm -3 )
plagioclasefeldspar2.68
alkali feldsparfeldspar2.57
feldsparfeldspar2.63
quartzquartz2.65
silica phasesquartz2.32
pyroxenepyroxene3.40
Appendix
Table B10: Densities used to convert reported volume percent to mass fractions.
Figure C1: Composition and conditions predictions against simulated labels. Predictions for one validation composition, an andesite, was selected because it had the median of the per-composition assemblage NLL at the 50%-melt readout. (a, b) Predicted compositions of the most probable anhydrous phase (clinopyroxene, a) and the most probable hydrous phase (amphibole, b), as ranked by the presence branch. Predicted points represent the mean over 2,048 candidate compositions from the predicted Dirichlet and whiskers span the 25th to 75th percentile. Simulated label points reflect the median of the simulated labels across the composition’s 128 draws and whiskers span the 25th to 75th percentile. (c–e) Predicted formation temperature, pressure, and dissolved water content over 512 assemblages sampled from the presence branch, against the sampled conditions of the same 128 draws.
Figure C2: Zero-shot composition and conditions predictions for a measured rock. Predictions for one benchmark rhyolite, the rock of Fig. 2 , selected for having the median of the recalibrated joint NLL over the test set. (a, b) Predicted composition of the most probable anhydrous phase (alkali feldspar, panel a) and the most probable hydrous phase (biotite, panel b), as ranked by the presence branch. Points represent the mean over 2,048 candidate compositions from the predicted Dirichlet and whiskers span the 25th to 75th percentile. (c–e) Predicted formation temperature, pressure, and dissolved water content over 512 assemblages sampled from the presence branch. The benchmark reports neither mineral compositions nor formation conditions, so no measured counterpart appears.
Figure D1: Recalibration tiers across fit sizes. The three recalibration configurations described in Appendix B.7 . (a) Median joint NLL. For reference, the zero-shot joint NLL is 13.4 nats, above the plotted range. (b) Abundance MAE, reported as the mean of the 7 per-category MAE values, where the dashed line marks the zero-shot value. For both panels, error bars span the interquartile range over the 10 replicates at each fitting size, evaluated on the fixed validation set of 111 rocks spanning 39 stratigraphic units. Axes are shared with Figs. D2 and D3 .
Figure D2: Heads trained on measured abundances across fit sizes. The three heads differing only in their input: the frozen encoder’s latent representation, the composition the latent representation was computed from, or a zero vector. (a) Median joint NLL. For reference, the zero-shot joint NLL is 13.4 nats, above the plotted range. (b) Abundance MAE, reported as the mean of the 7 per-category MAE values, where the dashed line marks the zero-shot value. For both panels, error bars span the interquartile range over the 10 replicates at each fitting size, evaluated on the fixed validation set of 111 rocks spanning 39 stratigraphic units. Axes are shared with Figs. D1 and D3 .
Figure D3: Simulated and measured supervision across fit sizes. The recalibration configuration is the mixture tier described in Appendix B.7 , which corrects predictions supervised by simulated labels. Meanwhile the head is the one given the latent representation, which is trained on measured abundances alone. Both represent the most performant from their respective groups in Figs. D1 and D2 . (a) Median joint NLL. For reference, the zero-shot joint NLL is 13.4 nats, above the plotted range. (b) Abundance MAE, reported as the mean of the 7 per-category MAE values, where the dashed line marks the zero-shot value. For both panels, error bars span the interquartile range over the 10 replicates at each fitting size, evaluated on the fixed validation set of 111 rocks spanning 39 stratigraphic units. Axes are shared with Figs. D1 and D2 .
MajorSiAlCaMgFeKNaTiCr    
R2R^{2}0.920.920.880.920.860.870.790.90–    
TraceSrZrBaYRbNiNbVLaCeNdYbTh
R2R^{2}0.800.800.770.820.890.800.850.830.930.940.940.840.88
Appendix
Table E1: Per-element imputation performance of the pretrained encoder. To score an element, its value is hidden from the encoder and a decoder recovers it from the frozen pooled latent representation. Each entry is the log-space R2R^{2} between recovered and measured concentrations over 157,544 held-out compositions. Cr occupies the ninth major-element row of the input vocabulary but is never reported in the pretraining compilation. Its row therefore always enters as padding, and no measured value exists to score against.
 FeldsparPyroxeneAmphiboleBiotiteOlivineFe-Ti oxidesQuartz
AUROC0.8700.7800.5750.7700.8290.6380.853
Appendix
Table E2: Zero-shot presence discrimination per mineral category. AUROC of the predicted presence probability, computed over every rock in the benchmark.
 FeldsparPyroxeneAmphiboleBiotiteOlivineFe-Ti oxidesQuartz
Spearman ρ\rho0.4560.3960.350-0.0440.7620.2190.667
nn60129231232279559148
Appendix
Table E3: Zero-shot abundance ranking per mineral category. Spearman rank correlation between predicted and measured abundance, computed over the nn rocks that report a quantity for each category.
 FeldsparPyroxeneAmphiboleBiotiteOlivineFe-Ti oxidesQuartz
Rocks over-predicted (%)7237462941684
Median pred/measured1.140.760.920.610.880.342.01
nn60129231232279559148
Appendix
Table E4: Zero-shot abundance bias per mineral category. The percentage of rocks where predicted abundance exceeds the measured abundance; the median across rocks of the predicted-to-measured abundance ratio; and the number of rocks nn that report a quantity for the category.
Figure F1: Composition encoder pretraining loss. Masked-element reconstruction loss (dimensionless) over pretraining of the composition encoder, where the loss is the standardized log MSE as presented in Eq. 3 . The validation loss is evaluated on held-out compositions from the pretraining set. Training was halted by early stopping once the validation ceased to improve 0.001/epoch.
Figure F2: Per-head training loss. Loss over joint-heads training, where each head is saved at its own best-validation epoch (vertical dashed line). The validation loss is evaluated on the 55,576 rows of the validation compositions (Table B6 ). (a) Assemblage head, the summed presence and abundance loss. (b) Composition head. (c) Conditions head.
Figure F3: Joint NLL decomposed into presence and abundance factors. The three heads trained on measured abundances, differing only in their input (Appendix D ). (a) The presence (Bernoulli) factor of Eq. 1 . (b) The abundance (Dirichlet) factor, where all three heads are given the measured mineral list as their support (i.e., teacher-forced). (c) The joint NLL, the sum of the presence and abundance factors. Note that this panel corresponds exactly to Fig. D2 a. Points are medians and error bars span the interquartile range over the 10 replicates at each fitting size, evaluated on the benchmark validation set. Note that the gap between curves is a difference of medians because each curve summarizes one head separately.
Figure F4: Water sensitivity of amphibole crystallization. Amphibole occurrence and abundance in the simulated label corpus, pooled over the validation compositions.
More in this topic

Earth Science

See all 26 papers

Explore similar work

CardsList

  1. GeoChemAD: Benchmarking Unsupervised Geochemical Anomaly Detection for Mineral Exploration

    Mar 13, 2026Yihao Ding, Yiran Zhang, Chris Gonzalez +2Unsupervised Anomaly DetectionAnomaly Detection

    The University of Western Australia, Perth, WA, Australia

  2. Expert lithological labels do not improve learning when used to predict marine sediment properties

    Oct 26, 2026John M. Aiken, William Gilpin, Thorsten W. BeckerMaterials Property PredictionSelf-Supervised Learning

    University of Oslo, Oslo, Norway · Expert Analytics, Oslo, Norway · Department of Physics, University of Texas at Austin, Austin, TX, USA +2

  3. Minerals in the Wild: A Hyperspectral-XRF Dataset for Elemental Composition Estimation

    Aug 31, 2026Eleftheria Tetoula-Tsonga, George Arvanitakis, Theodoros GiannakasSpectral UnmixingHyperspectral Imaging

    Institute of Communication and Computer Systems, Athens, Greece · Geonova, Athens, Greece