Enhancing Spatial Reasoning in Large Language Models for Metal-Organic Frameworks Structure Prediction

The paper introduces MOF-LLM, a novel framework that enhances the spatial reasoning capabilities of a Qwen-3 8B language model through spatial-aware continual pre-training, supervised fine-tuning, and reinforcement learning to achieve state-of-the-art, high-efficiency block-level 3D structure prediction for Metal-Organic Frameworks.

The Gist

The Big Picture: Building with Molecular LEGO

Imagine Metal-Organic Frameworks (MOFs) as incredibly complex, microscopic structures made of "LEGO bricks." These aren't plastic bricks, but tiny clusters of metal atoms and organic molecules that snap together to form a porous, sponge-like crystal. Scientists love them because they can be used to catch carbon dioxide from the air or deliver medicine inside the body.

The problem? There are millions of ways to snap these bricks together. Trying to find the perfect, stable structure by building them one by one in a lab is like trying to find a specific needle in a haystack by looking at every single piece of hay. It takes too long and costs too much.

For a long time, computers tried to solve this by looking at every single atom (like counting every grain of sand in a castle). But MOFs are so big and complex that this approach is too slow and confusing for computers.

The New Idea: Teaching a Language Robot to Build

This paper introduces a new tool called MOF-LLM. Think of a Large Language Model (LLM) like a super-smart robot that has read every book in the library. Usually, it's great at writing stories or answering questions, but it's terrible at 3D geometry—it doesn't "see" space well.

The researchers asked: Can we teach this language robot to build these molecular LEGO structures?

The answer is yes, but only if we teach it a new way of thinking. Instead of asking the robot to describe every single atom (which is like asking it to write a novel about every grain of sand), they taught it to think in blocks.

How They Did It: A Three-Step Training Camp

To turn a text-reading robot into a 3D builder, the team used a three-step training process:

1. The "Spatial Awareness" Class (Continual Pre-Training)
First, they gave the robot a crash course in geometry. They didn't just show it the chemical names of the bricks; they gave it a "mass-weighted bounding box" description.

• The Analogy: Imagine you are blindfolded and trying to stack boxes. If someone just says "Box A," you don't know how big it is. But if they say, "Box A is 5 inches wide, 3 inches tall, and weighs 2 pounds," you can start to visualize it.
• What they did: They fed the robot data about the size, shape, and weight of the molecular blocks, plus how they connect. This helped the robot understand the "shape" of the pieces before it even tried to build.

2. The "Assembly Line" Class (Supervised Fine-Tuning)
Next, they taught the robot how to actually put the pieces together.

• The Analogy: Now that the robot knows what the boxes look like, they taught it the instructions: "Take Box A, move it 2 inches to the right, and rotate it 45 degrees."
• What they did: They trained the model to predict the exact position and rotation (using something called Euler angles, which are like describing a turn as "roll, pitch, and yaw" instead of complex math) for each block to build a stable crystal.

3. The "Quality Control" Class (Reinforcement Learning)
Finally, they let the robot practice, but with a strict judge.

• The Analogy: The robot builds a structure. If the structure collapses or the blocks crash into each other, the judge gives it a "thumbs down" (a low score). If the structure looks exactly like a perfect, stable crystal, the judge gives a "thumbs up" (a high score). The robot learns from these scores to stop making mistakes.
• What they did: They used a system called SAPO (Soft Adaptive Policy Optimization). If the robot built a structure that was close to the real thing, it got a bonus. If it built something unstable, it was gently corrected. This helped the robot learn to avoid "crashes" and build stable structures.

The Results: Fast and Accurate

The team tested their new robot, MOF-LLM, against other computer programs that try to build these structures.

• Accuracy: MOF-LLM was the best at its job. It successfully predicted the correct structure about 36% of the time (which is a huge win in this field), beating all other methods.
• Speed: This is where it really shines. Other methods take seconds or even minutes to build one structure because they have to do complex math over and over. MOF-LLM is like a speed-reader; it generates a structure in 0.04 seconds. It's so fast it could theoretically build thousands of structures in the time it takes a human to blink.

Why This Matters

The paper claims that by treating these complex molecules as "blocks" and teaching a language model to understand 3D space, they have created a tool that is both smarter and faster than anything else currently available.

They didn't just make a robot that guesses; they made a robot that understands the geometry of the building blocks. This allows scientists to skip the slow, expensive trial-and-error in the lab and instantly see which molecular designs are likely to work, potentially speeding up the discovery of new materials for cleaning the air or curing diseases.

In short: They taught a text-bot to become a master architect of molecular LEGO, making the search for new materials significantly faster and more accurate.

Technical Summary

Technical Summary: Enhancing Spatial Reasoning in Large Language Models for Metal-Organic Frameworks Structure Prediction

Problem Statement

Metal-Organic Frameworks (MOFs) are porous crystalline materials with significant applications in carbon capture, drug delivery, and water harvesting. However, accurately predicting their 3D structures remains a formidable challenge due to their high structural complexity, often involving hundreds of atoms per unit cell. While Large Language Models (LLMs) have shown promise in generating crystal structures for simpler bulk materials, their direct application to MOFs is hindered by two primary factors:

1. Context Length: Representing MOFs at the atom level results in excessively long token sequences that exceed the context limits of current LLMs.
2. Spatial Reasoning Deficits: LLMs struggle to understand complex 3D geometries and precise rotational orientations required to assemble building blocks (metal nodes and organic linkers) without collisions or physical implausibility. Existing LLM-based approaches often rely on 1D string identifiers or external solvers, failing to explicitly perceive 3D spatial relationships or directly assemble precise atomic structures.

Methodology

The authors propose MOF-LLM, the first framework to adapt LLMs for block-level MOF structure prediction. The approach treats MOF generation as an autoregressive assembly task, predicting the lattice parameters and the roto-translations (position and orientation) of pre-defined building blocks rather than individual atoms.

The framework employs a three-stage training pipeline using a Qwen-3 8B backbone:

1. Text Formatting and Representation

To bridge the gap between 3D geometry and LLM text processing:

• Blocks: Building blocks are represented by canonical SMILES strings to leverage chemical semantic priors.
• Geometry: To compensate for the loss of 3D information in 1D strings, the authors augment the input with spatial priors: molecular weight, PCA-based spatial spans (bounding box dimensions), and topology codes (RCSR).
• Transformations: Lattice parameters are converted to scalar values. Crucially, 3D rotation matrices are converted to Euler angles (roll, pitch, yaw), which are found to be more intuitive for LLMs than quaternions or axis-angle vectors.

2. Three-Stage Training Pipeline

• Spatial-Aware Continual Pre-training (CPT): The model is pre-trained on a curated dataset containing block connectivity, geometry, and topology information. This stage injects explicit spatial priors, enabling the LLM to understand the intrinsic geometry of blocks and their potential arrangements.
• Structural Supervised Fine-Tuning (SFT): The model is fine-tuned to autoregressively generate the complete 3D configuration (lattice parameters, translation vectors, and Euler angles) given a set of building blocks. This stage focuses on the assembly logic.
• Matching-Driven Reinforcement Learning (RL): To address structural instability (e.g., block collisions), the authors employ Soft Adaptive Policy Optimization (SAPO). The model generates a group of candidate structures, which are evaluated against ground-truth references using a structural matching reward (based on StructureMatcher and RMSE). The reward function provides bonuses for high-precision matches and penalizes structural failures, guiding the policy toward generating stable, physically plausible MOFs.

Key Contributions

1. First LLM Framework for MOFs: MOF-LLM is the first work to apply LLMs directly to MOF structure prediction, moving beyond atom-level text representations to a block-level generative paradigm.
2. Enhanced Spatial Reasoning: By integrating spatial-aware CPT with explicit geometric descriptors (PCA spans, topology codes) and Euler angle representations, the framework significantly improves the LLM's ability to reason about 3D block assembly.
3. Efficient and Accurate Prediction: The method achieves state-of-the-art performance while maintaining exceptional computational efficiency, generating structures in a single autoregressive pass.

Experimental Results

The model was evaluated on a dataset of 324,426 hypothetical MOFs (Boyd et al. [3]).

• Accuracy: MOF-LLM achieved a match rate of 35.78% (at strict tolerance $stol=0.5$) and 93.25% (at $stol=1.0$), outperforming both denoising-based baselines (MOF-BFN, MOFFlow) and other LLM-based approaches (PLaID++). It also demonstrated lower Root Mean Square Error (RMSE) in atomic positions compared to baselines.
• Efficiency: The inference time is 0.04 seconds per structure, significantly faster than denoising-based methods which require iterative sampling (e.g., MOF-BFN takes 0.21s for 5 samples).
• Scalability: The model maintains high performance as the number of atoms and building blocks increases, showing a widening performance gap over baselines for larger systems (>800 atoms).
• Ablation Studies:
  • Euler Angles: Outperformed axis-angle vectors, suggesting Euler angles are more LLM-friendly.
  • Spatial CPT: Removing spatial descriptors (topology, PCA spans) from the CPT stage caused a significant drop in match rate and an increase in structural invalidities (atomic overlaps and lone molecules).
  • RL: The SAPO stage significantly reduced structural implausibilities and improved alignment with ground truth.

Significance and Claims

The paper claims that MOF-LLM establishes a principled pathway for adapting generic LLMs to intricate scientific systems where spatial reasoning is critical. By overcoming the scalability bottlenecks of traditional first-principles methods and the context limitations of atom-level LLM approaches, this work offers an accurate and highly efficient alternative for accelerated MOF discovery. The authors position this framework as a foundational stepping stone toward next-generation MOF design, where natural language prompting could eventually enable versatile material design. The work contributes to reticular chemistry by demonstrating that LLMs, when properly guided with spatial priors and reinforcement learning, can effectively solve complex 3D assembly problems.