A modified vision transformer architecture with scratch learning capabilities for plutonic rock classification

Abstract The classification of plutonic rocks is essential for geological research, mineral exploration, and petrological analysis. Traditional methods rely on expert interpretation of thin-section images, making the process time-consuming and necessitates expertise. Vision Transformers (ViTs) provide promising tools to address these challenges by enabling automated classification. However, previous studies have focused primarily on model accuracy without considering generalization and often lack a large plutonic rock dataset curated by petrologists. Thus, this study introduces the Petrographic Vision Transformer (PViT), a family of ViT models designed to automatically classify plutonic rock types in Thailand with high accuracy and generalization. The proposed models incorporate a new Gaussian Error Linear Unit (new GELU) activation in the position-wise feed-forward network (FFN) and an optimized multi-head attention mechanism with merged Query (Q), Key (K), and Value (V) projections. Training using a dataset of 2,886 thin-section images representing 15 plutonic rock types, PViT-Base achieved 98.34% accuracy, exceeding the standard ViT-Base by 0.62 percentage points when trained from scratch. Moreover, PViT achieved slightly higher accuracy than the Dual Attention Vision Transformer (DaViT) and EfficientNetV2, both state-of-the-art (SOTA) models originally benchmarked on ImageNet-1K. Grad-CAM visualizations generated by our model highlighted influential regions for classification, corresponding to key petrographic features, which closely resemble the interpretive approach used by petrologists. These findings demonstrate the potential of PViT to advance automated classification and offer strong potential for future geoscience applications.