Beyond Language: Applying MLX Transformers to Engineering Physics

TL;DR

This paper introduces a physics-informed Transformer model for solving 2D heat conduction problems, leveraging MLX framework for efficient training and prediction on personal hardware, demonstrating high accuracy on unseen conditions.

Contribution

It presents the first application of Transformers to engineering physics problems using MLX, enabling accessible and efficient physics-based modeling on personal devices.

Findings

High accuracy in predicting steady-state temperature fields

Efficient training on modest hardware using MLX framework

Successful generalization to unseen boundary conditions

Abstract

Transformer Neural Networks are driving an explosion of activity and discovery in the field of Large Language Models (LLMs). In contrast, there have been only a few attempts to apply Transformers in engineering physics. Aiming to offer an easy entry point to physics-centric Transformers, we introduce a physics-informed Transformer model for solving the heat conduction problem in a 2D plate with Dirichlet boundary conditions. The model is implemented in the machine learning framework MLX and leverages the unified memory of Apple M-series processors. The use of MLX means that the models can be trained and perform predictions efficiently on personal machines with only modest memory requirements. To train, validate and test the Transformer model we solve the 2D heat conduction problem using central finite differences. Each finite difference solution in these sets is initialized with four random Dirichlet boundary conditions, a uniform but random internal temperature distribution and a randomly selected thermal diffusivity. Validation is performed in-line during training to monitor against over-fitting. The excellent performance of the trained model is demonstrated by predicting the evolution of the temperature field to steady state for the unseen test set of conditions.