Vision Transformers for the Self-Supervised Classification of Images of Rock Samples

The subsurface disciplines are always in a quest for analogs, whether it be geologic, petrophysical or engineering related. The classification of rock sample images provides a valuable start point in the search for analogs by providing an understanding of microstructural similarities and differences between different plays. Convolutional Neural Networks (CNN) are the dominant approach for these applications and are well-suited for these tasks given the inductive biases inherent to the process of training, which is invariance of the images under translation or rotation. Processing these images, however, can take an inordinate amount of time because of the need for proper annotation, which is extremely time consuming and resource intensive with conventional supervised learning approaches. Self-supervised learning, remarkably, is capable of distilling the information content of images to find digital analogs rapidly. Our dataset comprises 22,440 FIB-SEM images collected from 22 shale plays across North and South America. We explore a vision transformer-based architecture for the classification of these SEM images. Our goal is to identify the source of the image from among the 22 different unconventional plays. From a training set, we successively feed unlabeled images to independent teacher and student vision transformer (ViT) networks that learn from each other and progressively get better at identifying similarities and distinctions in the microstructure. We then exploit this by pairing the ViT-based algorithm with a k-Nearest Neighbors (k-NN) classifier that groups the images. The surprising result was the fidelity with which the features associated with any given play are automatically distilled from the images over the training period. This enabled a very high accuracy and low misclassification rate among 22 different plays. Without prior training on labeled data, the model successfully identified key image features that guide its classification through attention maps. Interestingly, there were a few misclassified samples from pairs of plays that share common features and form digital analogs. This work constitutes the first approach for self-supervised classification of images of rock samples. The self-supervision automatically extracts meaningful information that we discuss in the paper in the form of attention maps. These maps contain microstructural details that enable rapid and efficient classification, but more importantly, provide a pathway for interpretability for the petrologist or geologist concerned.