FlowLLM: Flow Matching for Material Generation with Large Language Models as Base Distributions

Authors: anuroop  sriram
benjamin kurt miller
ricky t q chen
brandon m wood fair
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:35:50

Abstract: Material discovery is a critical area of research with the potential to revolutionize various fields, including carbon capture, renewable energy, and electronics. However, the immense scale of the chemical space makes it challenging to explore all possible materials experimentally. In this paper, we introduce FlowLLM, a novel generative model that combines large language models (LLMs) and Riemannian flow matching (RFM) to design novel crystalline materials. FlowLLM first fine-tunes an LLM to learn an effective base distribution of meta-stable crystals in a text representation. After converting to a graph representation, the RFM model takes samples from the LLM and iteratively refines the coordinates and lattice parameters. Our approach significantly outperforms state-of-the-art methods, increasing the generation rate of stable materials by over three times and increasing the rate for stable, unique, and novel crystals by $\sim50$% – a huge improvement on a difficult problem. Additionally, the crystals generated by FlowLLM are much closer to their relaxed state when compared with another leading model, significantly reducing post-hoc computational cost.

Main paper: 
 introduction: 
 Material discovery holds transformative potential across numerous industries including carbon capture [38] , batteries [28] , photovoltaics [9] , and energy storage [1] . However, the vastness of the chemical space has hindered experimental synthesis of the majority of possible materials. Generative models offer a promising avenue for exploring this untapped potential.Generating crystalline materials is particularly challenging as it involves simultaneously generating both discrete (atomic types) and continuous values (atomic positions and lattice geometry). While existing approaches, namely autoregressive large language models (LLMs) [11, 6] and denoising models, e.g., denoising diffusion and flow matching [47, 16, 49, 48, 30, 26, 17] , have demonstrated success, they exhibit complementary strengths and weaknesses. LLMs excel at modeling discrete values, but they can struggle with continuous values due to their reliance on finite precision representations. Conversely, denoising models more effectively handle continuous values and can easily ensure equivariances, but they face challenges with discrete elements.LLMs also offer the distinct advantage of natural language prompting, enabling versatile and intuitive conditional generation. This capability is further enhanced by training LLMs on vast corpora of chemistry text, equipping them with valuable prior knowledge to generate chemically valid outputs. Queries like "Generate materials with a high bandgap and thermal stability" or "Propose a novel perovskite structure for efficient solar energy conversion" can be directly integrated into the LLM Figure 1 : FlowLLM generative process: the fine-tuned LLM is first prompted with an unconditional query to generate an initial material representation. This material is then iteratively transformed by the RFM model to update its atom positions and lattice parameters. The atom types are static in RFM. prompt, while denoising models typically require bespoke changes to the architecture and training procedure to handle conditional generation.To harness the strengths of both paradigms, we introduce FlowLLM, a novel hybrid approach that uses an LLM to generate an initial material representation, which is iteratively refined with a Riemannian Flow Matching (RFM; [2] ) model. This synergistic approach allows us to effectively bridge the gap between discrete and continuous modeling, resulting in a significant improvement in the rate of generation of stable, unique, and novel (S.U.N.) materials. Such materials expand the limited knowledge we have of "material space" and are much more likely to be synthesizable than unstable generations. Our experiments demonstrate that FlowLLM generates stable materials at over 300% higher rate, and S.U.N. materials at ∼ 50% higher rate compared to prior models, while retaining the LLM's ability to be prompted with natural language instructions. We offer two interpretations for the effectiveness of our approach. 1) The LLM learns a good base distribution for RFM: the LLM's output distribution serves as a learned base distribution for RFM, replacing the common practice of using the uniform base distribution. Since the LLM has been trained on material data, this learned base distribution is closer to the target distribution, greatly simplifying integration with RFM. 2) RFM refines the output of the LLM: The LLM generates an approximate material representation due to its finite precision when handling continuous values. The RFM then refines this approximation through iterative denoising, to generate a much more accurate representation.Our contributions are as follows:• We introduce FlowLLM, a novel hybrid approach for materials generation that combines LLMs and RFM, effectively leveraging their complementary strengths. • We demonstrate that FlowLLM significantly outperforms existing state-of-the-art generative models in generating novel and stable materials. • We show through ablation experiments that our method of combining LLM and RFM models through FlowLLM significantly outperform simpler combination approaches.Code for training the FlowLLM model is available at https://github.com/facebookresearch/flowmm. 
related work: 
 In the past, computational materials discovery relied on generating numerous candidate materials through random atomic substitutions in known materials [42] , followed by computationally expensive quantum mechanical screening [19] to assess stability. Genetic algorithms [8, 33] , and machine learning models trained to predict energies [37, 25] have accelerated this process, but the fundamental bottleneck of brute force search remains.Recent research has focused on generative models that directly produce stable materials, bypassing brute-force search. Diffusion models, either combined with Variational Autoencoders (VAEs) for partial variable prediction [47] or jointly diffusing all variables [16, 48, 49] have shown promise. Additionally, Riemannian Flow Matching [26] , Normalizing Flows [45] , and Variational Autoencoders [34] have also been adapted for material generation.A parallel line of work utilizes autoregressive Large Language Models (LLMs) for material generation [6, 11] , representing materials as a sequence of discretized tokens. Pretraining these models on natural language imbues them with powerful prior knowledge not attainable by other approaches. 
preliminaries: 
 Our approach models probability distributions over crystal lattices, defined as periodic arrangements of atoms in three-dimensional space. A crystal lattice is created by tiling a fundamental unit cell, where the unit cell contains a specific atomic configuration, forming the entire lattice when repeated.In this section, we present a high-level overview of crystal representations, building up to explain our model in section 4. Background details for the crystal representation are in appendix A.In the paper, we represent an n ∈ N atom crystal in a product space: c := (a, f , l) ∈ C, indicating the atom types, positions and unit cell geometry, respectively [47, 26] . The atom types are represented by a matrix of categorical vectors: a := a 1 , . . . , a n , where a i ∈ A. The atomic coordinates are represented using fractional coordinates within the unit cell, f := f 1 , . . . , f n , where f i ∈ F = T 3 with T denoting the unitary length, flat torus manifold, i.e., the fractional coordinates satisfy periodic boundary conditions; that is, the atoms "wrap around" the unit cell. The unit cell geometry is defined using lattice parameters l ∈ L, where L is the space formed by a 6-tuple of three side lengths (a, b, c) ∈ R + (Å, i.e. Angstrom) and three internal angles(α, β, γ) ∈ [60 • , 120 • ].This representation is not unique as the same crystal can be produced by different choices of unit cell. To make the representation unique, we select the minimum-volume unit cell and employ Niggli reduction [10] that uniquely determines the unit cell parameters.Equivariance & Invariance Given a group G with g• denoting a group action for some g ∈ G,a function f : X → Y is called G-equivariant if ∀x ∈ X , ∀g ∈ G, f (g • x) = g • f (x), while it is called G-invariant if ∀x ∈ X , ∀g ∈ G, f (g • x) = f (x). Since a crystal is not uniquely defined by any particular representation c but an infinite set, we know that the data distribution has a G-invariant density, where G represents symmetries of a crystal.Symmetries of crystals Concretely, our crystal representation exhibits multiple symmetries that we detail here. The symmetric group S n on n atoms permutes the atom indices: σ • c = a σ(1) , . . . , a σ(n) , f σ(1) , . . . , f σ(n) , l . The special Euclidean group SE(3) consists of orientation preserving rigid rotations and translations: (Q, T ) where Q ∈ SO(3) and T ∈ [-1 2 , 1 2 ] 3×1 denote 3D rotations and translations respectively. This element transforms the crystal as: (Q,T) • c = (a, f + τ 1 -⌊f + τ 1⌋, l).We emphasize that the representation c is completely invariant w.r.t. Q because lattice parameters do not contain orientation information. Since these represent symmetries fundamental to crystals, the data distribution q(c) is invariant to these group operations.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] Justification: We do not include any theoretical results in the paper. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
method: 
 Our goal is to fit a parametric generative model p(c; θ) to approximate the distribution of known meta-stable materials q(c) using a dataset of samples. The distributions p and q are defined on the Riemannian manifold C. Our FlowLLM model generates samples from the parametric distribution using a two-step procedure (see figure 1 ). First it samples the LLM, then it refines the LLM output using RFM, like so:EQUATIONEQUATIONwhere p LLM is modeled using a large language model [6, 11] , and p RFM is modeled using Riemannian Flow Matching (RFM) [3, 26] , and θ = (θ 0 , θ 1 ). Both the LLM and RFM frameworks are trained to estimate the data distribution over meta-stable crystals on samples from the Materials Project [15] .Overview of training First, we fine-tune an LLM to generate string representations of meta-stable materials [11] . Once trained, we can sample the LLM distribution using next token prediction, optionally conditioning on a prompt (see figure 2 ). Next, we train the RFM model using the FlowMM objective [26] where, conditioned on the chemical formula, will learn to transport between the LLM's model distribution and the data distribution. The full training process is described in Algorithm 1.We give the standard prompt to the LLM and allow it to do next token prediction until it produces a stop token. As long as all atom types are actual elements and the lattice parameters are physical, we move forward. Otherwise we reject the sample. Then, we convert the text to a crystal representation that serves as the initial sample. This sample's fractional coordinates f and lattice parameters l are iteratively refined by the RFM model to produce the final sample of FlowLLM. This sampling process is illustrated in figure 1 .LLMs define a distribution over sequences through an autoregressive decomposition, T t=1 p(w t+1 |w 0:t ), where each p(w t+1 |w 0:t ) follows a categorical distribution conditioned on all previous tokens (w 0:t ) in the sequence. Our LLM model closely follows Gruver et al. [11] .Tokenization Language models interact with strings in text datasets after the string is converted into a sequence of tokens. The choice of tokenizer can have a large impact on the performance of the language model. In terms of tokens, we represent a crystal c using fixed precision numberstwo decimal places for fractional coordinates, and one for lattice lengths. Angles are represented as integers. Atom types are represented as discrete tokens. We use LLaMA-2 models [41] for our LLM architecture since these models break numbers into a sequence of digits, which has been shown to dramatically improve performance on arithmetic tasks [23] .Training We rely on the extensive pretraining of LLaMA-2 models to instill useful biases over numerical operations. To train p LLM , we fine-tune a pre-trained LLaMA-2 model on a dataset of crystal structures represented as strings along with a prompt indicating that the model should generate bulk materials by writing the lattice in lengths and angles along with atom types and coordinates. An example of such a representation along with a prompt is shown in figure 2 .The flexibility of LLMs allows us to optionally include different kinds of conditional information in the prompt such as the chemical formula. We can also solve other tasks such as infilling by making changes to the prompt. For this hypothetical conditional generation, the prompt could include a desired Sample c i 0 ∼ p LLM (•|θ 0 ) using a prompt conditioned on the formula of c i c t := exp c0 (t log c0 (c 1 )) 12: L(θ 1 ) = ∥v θ1 t (c t ) -u t (c t |c 1 )∥ 2 13:Take gradient descent step on ∇ θ1 L(θ 1 ) 14: end while chemical formula, material properties, or a combination of such information. In this work, we used the same conditioning used in Gruver et al. [11] , and we leave a more detailed study of this to future work.Sampling To generate sequences from the model, the conditional distribution is sampled sequentially. The sampling procedure is modulated to control the diversity and sampling speed using the temperature (τ ) and nucleus size (P ) hyperparameters of nucleus sampling [13] . Temperature controls the entropy of the conditional distributions, introducing a trade-off between diversity and mode sampling. The nucleus size limits the number of tokens that can be sampled. Given a nucleus size P with 0 < P ≤ 1, sampling is restricted to the most probable tokens with cumulative probability P .The LLM architecture does not inherently produce a symmetric density, i.e., the distribution of meta-stable crystals that the LLM learns is not symmetric according to the fundamental properties of crystals. We perform no fractional coordinate data augmentation via translation, and no token permutation data augmentation. Unlike the other symmetries, rotation invariance holds for the learned LLM distribution due to our choice of representing the unit cell with lattice parameters.Riemannian Flow Matching RFM produces a Continuous Normalizing Flow [3] , i.e., a continuous, parametric, diffeomorphism between the LLM base distribution p 0 := p LLM and an approximation to our target distribution p 1 ≈ q. To model p RFM := p 1 , we fit a time-dependent vector field v θ1 t that has been adapted to crystals and is implemented using a neural network with parameters θ 1 . Continuous Normalizing Flows are computationally expensive to train using maximum likelihood, but an alternative objective called Conditional Flow Matching [22] is more stable and scales better. The objective was generalized to Riemannian manifolds [2] , and specifically to labeled point clouds with periodic boundary conditions, i.e. crystals, by Miller et al. [26] .Concretely, each point c ∈ C has an associated tangent space T c C with an inner product ⟨u, v⟩ for u, v ∈ T c C, enabling the definition of distances, volumes, angles, and minimum length curves (geodesics). The geodesics for any C that we consider can be written in closed form using the exponential and logarithmic maps. The geodesic connecting c 0 ,EQUATIONwhere exp □ and log □ are the exponential and logarithm maps for the manifold C. These geodesics help define the supervision signal used to train RFM.Our RFM generative model v θ1 t : [0, 1] × C → T C is parameterized as a time-dependent, smooth vector field. Training proceeds by regressing onto conditional vector fields u t (c|c 1 ) that generate single data points c 1 . For the geodesic path, this corresponds to u t (c|c 1 ) = -1 1-t log c1 (c). The general RFM training objective is then:EQUATIONSince we only use flat manifolds, ∥•∥ is the Euclidean norm. At the optimal solution, v θ1 t generates p t with endpoints p 0 = p, p 1 = q. At sampling time, we draw a sample from c 0 ∼ p and solve the ordinary differential equation d dt c t = v θ1 t (c t ) with initial value c 0 at t = 0; the solution at t = 1 is then the sample from our RFM model.Geometry of F We apply the conditional vector field for a point cloud living on a n×3-dimensional product of flat tori invariant to global translations, i.e. fractional coordinates with periodic boundary conditions [26] . This is a geodesic path, which may cross the periodic boundary:EQUATIONwhereEQUATIONḟ i ∈ T f i F i for i = 1, . . . , n.Computing the geodesic of n atoms amounts to an atom-wise application of log f 0 on f 1 and exp f on ḟ ∈ T f F respectively. Additionally, following Miller et al. [26] we address translation-invariance by removing the mean torus translation:EQUATIONGeometry of L The space of lattice parameters,L := R +3 × [60, 120]3 , is a Euclidean space with boundaries. We can ignore these boundaries for the lattice lengths in R +3 since (i) the data does not lie on the boundary (a, b, c > 0) and (ii) we can clamp our base distribution to be positive with rejection. The boundary issue for the lattice angles α, β, γ can be addressed [26] using a diffeomorphism φ : [60 • , 120 • ] → R to unconstrained space, applied element-wise to each angle:EQUATIONwhere σ(.) and logit are the sigmoid and the log-odds functions, respectively. We directly apply RFM in the unconstrained space, and for sampling, we map the angles back into [60• , 120 • ] using φ -1 .The RFM training objective With this formulation, our training objective based on (4) becomes:EQUATION+ λ l 6 v L,θ1 t (c t ) + l 0 -l 1 2 ,where we now use p LLM as the base distribution, and c t = (f t , l t , a). The loss coefficients λ f , λ l ∈ R + are hyperparameters. We use a graph neural network (GNN) inspired by [36, 16, 26] for v θ1 t (c). This GNN enforces equivariance to atom permutations via message passing, invariance to atom translation by featurizing graph edges as relative displacements of nodes, and invariance to rotations by our choice of lattice representation. See appendix B for more details about the GNN architecture.Model symmetries Just like the LLM, the orientation-invariant representation of the unit cell leads to global rotation invariance. However, permutation and translation symmetries are not so simple. If the parameterization of the RFM velocity field is G-equivariant, and the base distribution is G-invariant, then the model density is G-invariant [18] . We use graph neural networks [36, 40, 27, 44, 7, 21, 31, 51] , and additional projections [26] , to ensure that the RFM velocity predictions are Gequivariant to both permutation and translation. However, we will generally not recover a translation invariant density because the base distribution defined by the LLM is not invariant to translation. The density will be permutation invariant in our RFM representation because the each atom is a node in an unordered point cloud and the LLM ordering is ignored by the RFM, but the density will not be permutation invariant in the text representation, due to the LLM's lack of token permutation invariance.Empirically, we do not find the lack of exact invariance to be a problem, and FlowLLM outperforms methods with exact invariance (section 5). This is because an LLM trained to generate crystals is approximately invariant to crystal symmetries. This was verified by Gruver et al. [11] who proposed a new metric, Increase in Perplexity under Transformation (IPT), to quantify this approximation:EQUATIONwhere g * = arg min PPL(t g * (s)), and PPL is the perplexity of the sequence, the exponent of the length-normalized cross entropy loss, PPL(s) = 2 CE(s)/n . They find that a well-trained LLM obtains a small IPT value, implying that it is approximately invariant.The LLM base distribution is not constrained to C, i.e. the LLM can generate invalid crystals. We find that this is extremely rare and easy to detect. In such cases, we simply reject that sample, and draw a new sample until we get a valid crystal. Empirically, we found this rejection rate to be ∼ 0.5% with a softmax temperature of 0.7.Text is not continuous in L or F The LLM base distribution only takes non-zero values over a small number of discrete points due to the use of finite precision representations. For example, we represent fractional coordinates with only 2 decimal places, so they can only take one of 100 distinct values. We can mitigate this problem by adding a small amount of random zero-mean gaussian noise to all continuous values predicted by the LLM. Empirically, we do not observe any noticeable difference in performance due to this added noise (see appendix F).Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: Our method has been described in detail (section 4) with all relevant hyperparameters (appendix C). The MP-20 dataset we use is publicly available. We have made our code available on Github. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discuss societal impact in section 6. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
setup: 
 We trained our model on the widely used MP-20 dataset1 of inorganic crystalline materials [47] . MP-20 comprises 45,231 materials, a subset of the Materials Project [15] containing up to 20 atoms known to be metastable (see section 5.2).We first train our LLM independently using the various prompting strategies described in Section 4.Unless otherwise specified, we employed a pretrained LLaMA-2 70B model [41] for all experiments, that was fine-tuned with the Low-Rank Adapters (LoRA) method [14] using PyTorch [32] and Transformers [46] .Next, we trained the RFM model using the fine-tuned LLM (with frozen weights) as the base distribution and the MP-20 dataset as the target distribution. For computational efficiency, we sampled a large number (N tr ) of examples from the base distribution in advance, and used the same set for all of our training runs. To create this set, we sampled N tr materials, with replacement from MP-20, and queried the LLM with a prompt conditioned on the chemical formula of each of these materials. This results in a set of N tr pairs, {(c i 0 , c i 1 )} Ntr i=0 , of LLM generated materials and ground truth materials that constitutes the training set for the RFM model. We list the hyperparameter values used in our experiments in appendix C.To generate new samples, we first generate a material from the LLM using an unconditional query. We then perform an integration with the RFM model, starting from this LLM-generated material. During sampling, we can adjust hyperparameters such as temperature τ , nucleus probability P , and the number of integration steps to achieve different trade-offs between diversity, accuracy, and efficiency.Our primary metrics are Stability Rate, the percentage of generated materials that are thermodynamically stable, a key indicator of synthesizability, and the S.U.N. rate, the percentage of materials that are stable, unique and novel. Since computing stability is computationally expense, Xie et al. [47] proposed a number of proxy metrics. We explain these metrics in more detail in appendix D.One key difference in evaluation between the proxy metrics and the stability metrics is the use of pre-relaxation and relaxation techniques. Proxy metrics are computed on raw samples without any further processing. Stability metrics are computed on structures that are first pre-relaxed using CHGNet [5] then relaxed using Density Functional Theory.Density Functional Theory is extremely expensive, even with speedups using pseudo-potentials [20] .Ideally, the generative model can generate many S.U.N. structures that are already close to their relaxed ground state. Generating structures close to ground state may also indicate that the model has done a better job capturing the data distribution. It can also speed up or obviate the need for relaxing the generated structures, which has huge computational benefits. We include several additional metrics to measure the closeness of generated and corresponding ground state structures, that are described in appendix E.We compare our model to four prior methods: CD-VAE [47] , a hybrid Variational Autoencoder & diffusion model; DiffCSP [16] , a diffusion model; FlowMM [26] , a Riemannian Flow Matching model; and CrystalLLM [11] , which fine-tunes a LLaMA-2 model on materials represented as sequences.The LLM and RFM components of FlowLLM closely resemble the formulations in CrystalLLM and FlowMM , respectively. To compare different models, we generate 10,000 new structures from each model and compare the metrics described in section 5.2.Our main results are presented in table 1. On the most important metrics, namely the Stability & S.U.N. rates, FlowLLM significantly outperforms all prior methods across various LLM sampling parameters. For our best FlowLLM model (τ = 0.7, P = 0.9), 17.82% of the generated structures are stable, out of which 48% are novel (not similar to any training or validation structure). Of the remaining structures, 58% are unique, leading a to a S.U.N. rate of 4.92%. FlowLLM obtains a ∼ 300% higher stability rate and ∼ 50% higher S.U.N. rate than the best prior model! Figure 3a shows histograms comparing the E hull values of generated materials from FlowLLM compared to prior models. Clearly, FlowLLM generates many more materials with lower E hull values than the other models.The results on proxy metrics, on the other hand, remain mixed. Diffusion and flow matching methods excel on Coverage Recall, while CrystalLLM has the best Composition Validity. FlowLLM achieves the best compromise between coverage and validity, potentially explaining its superior Stability & S.U.N. rates. It is important to note that many of these metrics have become saturated, offering limited discriminatory power for evaluating state-of-the-art models. As a result, we anticipate a decreased reliance on these metrics in future research.Comparison of generated and relaxed structures While the stability rate and S.U.N metrics capture whether the generated structures can be relaxed to stable / S.U.N. states, they do not address the question: How close are the generated structures to their relaxed state? To answer this question, we compared generated structures to those same generated structures after relaxation using CHGNet, computing the following metrics between generated and CHGNet relaxed states: Match Rate and RMSD, as defined by StructureMatcher, along with the ∆-Energy and the average Num steps between the states. Definitions for these metrics can be found in appendix E. Table 2 shows a comparison of FlowMM and FlowLLM. The samples generated by FlowLLM are significantly closer to ground state compared to FlowMM, according to our metrics.Importance of learned base distribution One motivation for a hybrid LLM-RFM model is to leverage the LLM's superior ability to generate accurate atom types compared to denoising models. To isolate this effect, we trained the FlowLLM-Types model, following a similar procedure as FlowLLM but using simple base distributions for lattice parameters and fractional coordinates identical to those used in FlowMM [26] . Thus, the LLM only contributes to atom type prediction in this model. Despite this simplification, FlowLLM-Types still surpasses prior models on the Stability Rate metric (table 1 ), highlighting the benefits of employing an LLM for atom type prediction. The stability rate of FlowLLM-Types remains considerably lower than that of FlowLLM, underscoring the substantial value of using learned base distributions.The number of distinct element types in a material is called the N -ary value of that material. Figure 3b compares the distribution of N-ary values for different models with the target data distribution. FlowMM and FlowLLM match the data distribution better than the diffusion models, which tend to generate too many materials with high n-ary.Number of RFM integration steps Compared to diffusion and flow matching models which require hundreds or thousands of integration steps, FlowLLM is able to converge in as little as 50 steps (figure 3c ). This is not surprising given our use of a learned base distribution. 
discussion: 
 The discovery of novel, stable materials holds the potential to help revolutionize numerous industries, but progress has been slow due to the high computational costs involved. Widely used random structure search methods [33] yield less than a 1% success rate in identifying stable materials. Given the substantial cost of validating generated structures using density functional theory, improving this rate is of paramount importance.Recent breakthroughs with denoising models [16, 26] and large language models [11] have increased the stability rate to ∼ 5%, a significant improvement over traditional approaches. In this work, we propose a novel generative model which harnesses the strengths of both paradigms to further increase this number by over 3×, representing a major advancement in the field.Limitations While FlowLLM excels at generating stable materials, a key limitation is its lack of end-to-end differentiability. This hinders its direct application to inverse design, where generative models are optimized to generate material with specific properties, as explored in prior work using denoising models [49, 47] . Future research could investigate extending FlowLLM for inverse design.Broader impact This work can accelerate the discovery of new materials for renewable energy, electronics, and carbon capture, ultimately benefiting society by enabling more efficient and sustainable technologies. However, the adoption of generative models also raises concerns, such as the creation of harmful substances and access inequalities.Atomic types The representation of atomic number is dependent on the model processing the data.In the LLM, the name of the element can be written into the text representation directly. This can be a string or single token, depending on LLaMA-2's tokenization. In the RFM framework, we applied a one-hot representation.Unit cell geometry Throughout the paper and in our implementation, we represent the unit cell using lengths and angles; however, there is another representation relevant for defining the fractional coordinates and better expressing crystal symmetries. The unit cell can be defined by a matrix of Cartesian column vectors l := l1 , l2 , l3 ∈ L = R 3×3 . This representation has strictly more information than l, since it also defines the orientation of the unit cell. This orientation is irrelevant in our paper, since we want rotation invariance. That's why we choose l in the first place.Fractional coordinates Now that we have the representation l, we can define fractional coordinates.Recall, atomic positions are typically represented using Cartesian coordinates x := x 1 , . . . , x n ∈ X = R 3×n with coordinates in the rows and atoms in the columns. The Fractional coordinate representation is definedf := l-1 x = f 1 , . . . , f n ∈ F = [0, 1) 3×n .In this section, we describe the graph neural network used in our RFM model. Our GNN model is inspired by the GNNs used in FlowMM [26] and DiffCSP [16] , which in turn adapted the EGNN [36] model for fractional coordinates,EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONwhere m ij (s) , m i (s) represent messages at layer s between nodes i and j, h j (s) represents hidden representation of node j at layer s; φ m , φ h , ϕ h (0) , φ ḟ , φ l represent parametric functions with all parameters noted together as θ. Finally, we define SinusoidalEmbedding(x) := (sin(2πkx), cos(2πkx)) T k=0,...,n f req ,where n f req is a hyperparameter. We standardized the l input to the network with z-scoring. We also standardized the outputs for predicted tangent vectors ḟ , l. Models were trained using the AdamW optimizer [24] .We used N tr = 3.3 × 10 6 and trained the model for 20 epochs with early stopping. To generate the N tr training pairs, we used temperature τ = 0.9 and nucleus probability P = 0.99. While other values might be explored, the high computational cost of experimentation limited our exploration of these parameters. For training the RFM model, we swept over a few values of learning rates: {1e-3, 7e-4, 5e-4, 3e-4, 1e-4}. To compute the loss function, we used loss weights λ f = 200, and λ l = 1 in the training objective (equation ( 8)). These values were chosen by running a grid search over λ f ∈ {100, 200, 300, 400}, λ l ∈ {1}. Additional hyperparameter settings are given in table 3 .The LLM was trained for 10 epochs with a batch size of 16, and cosine annealed learning rate of 0.0005, with LoRA rank = 8 and α = 32.We trained our LLM model on 8x 80GB A100 GPUs for roughly 1 day. We used 4-bit quantization and LoRA to optimize training. Sampling from the trained LLM required a total of ∼ 250 A100 GPU days, that were parallelized over 300 A100 GPUs.Each of our RFM models were trained for 2 days on a single 32GB V100 GPU. All experiments were performed on an internal GPU cluster.Evaluations required running DFT computations that were run on a large internal CPU cluster with 5000 nodes, each equipped with a 26-core Intel Cooper Lake-SP CPU, and 64GB memory. Each DFT computation took about 1 hour of compute on a single node, and we ran nearly 50,000 such computations to evaluate all of our models. Thermodynamic stability is a key indicator of synthesizability, and generating novel stable materials is of keen interest in material science. Stability is determined by comparing a material's energy to those of competing crystals with the same elements. Formally, stability is measured by constructing a convex hull of all competing materials from a reference set and computing the distance from this hull (called Energy above the Hull, or E hull ). Stable materials have E hull < 0, while materials with E hull < 0.08 eV/atom are called metastable [39] With this defintion of stability, we define our Stability Rate metric as the percentage of generated materials that are stable (E hull < 0, and n-ary ≥ 2). For our reference set of materials, we use the Materials Project database recorded by [35] in February 2023.Following Miller et al. [26] , we compute E hull values by running structure relaxations on the generated structures with the CHGNet model [5] followed by density functional theory (DFT) [19] calculations.While stability rate is an important metric, it does not capture novelty. Therefore, we define a second metric, the S.U.N. rate which measures the percentage of generated structures which are Stable, Unique, and Novel. To determine novelty, we exclude generated structures that are similar to any structure in the training dataset. Similarity is measured using pymatgen's StructureMatcher [29] with default settings. A generated structure that is not similar to any training data structure is considered novel.To compute uniqueness, we use StructureMatcher to do pairwise comparisons between all generated structures, and group similar structures into equivalence classes. Each group is only counted as a single unique structure for the purpose of computing the S.U.N. rate. Formally, StabilityFor many practical applications in chemistry, it is important to find the local energy minimum of a generated structure. This is done by performing computationally expensive structure relaxations. Thus, it is beneficial to generate structures close to their ground state. To compare how close the generated structures are to their ground state (i.e. local energy minimum), we define 4 additional metrics (shown in table 2 ):1. Match Rate: What fraction of generated structures and corresponding ground state structures are similar (where similarity is computed using pymatgen's StructureMatcher with default settings). 2. RMSD: Average RMS distance between generated structures and corresponding ground state structures computed using pymatgen's StructureMatcher whenever there is a match. 3. ∆-Energy: Difference in energy between the generated structure and ground state structure of the DFT relaxation. This measures the reduction in energy during the structure relaxation process. 4. Num Steps: Number of optimizer steps needed to pre-relax the generated structure using CHGNet.Table 4 shows the effect of adding noise to the base distribution. We do not see a significant impact from the added noise.We compare the time to generate 10,000 materials between FlowLLM with FlowMM. Inference for both models was run on a machine with a 32 core Intel(R) Xeon(R) Platinum 8488C CPU, and a single 80GB A100 GPU. FlowMM used 750 integration steps, and the RFM step of FlowLLM used 250 integration steps. With this setup, the FlowMM model takes 65.1 minutes to generate 10,000 materials, while FlowLLM takes 89.6 minutes, which is comparable to FlowMM.A more useful metric is the time to generate a S.U.N. material, which is computed by dividing the inference time by the number of generated S.U.N. materials. With this metric, FlowMM takes 16.14 seconds to generate S.U.N. material, while FlowLLM takes only 10.9 seconds.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We do not plan to release trained models or any new datasets. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We use a publicly available dataset and we cite the appropriate source for the data (section 5). Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.38th Conference on Neural Information Processing Systems (NeurIPS 2024).Publicly available at https://github.com/txie-93/cdvae/tree/main/data/mp_20 
 
 Appendix:  


Decision: The paper presents a hybrid approach that combines LLMs and conditional flow models for material design. LLMs are used to generate a noisy version of the material and flow models are used to refine the initial generation. The reviewers marked the strong empirical results as one of the main strengths of the work and found the idea of combining LLMs and Flows interesting. However, they also raised concerns regarding the novelty, computational cost, and lack of comparisons against force field-based approaches which were mostly addressed in the rebuttal. Given these we are happy to recommend the paper for acceptance.

Reviews: review: paper_summary: paper_summary: This paper presents a generative model for crystals that uses a LLM as a base distribution and a flow matching model to refine the 3D structure. Both the LLM and flow matching parts are mostly borrowed from prior studies. However, the authors demonstrate a significant improvement in the stability of generated materials by using LLM generated crystal as the prior for the flow matching model.
strength_weakness: strength_weakness: Despite the significant improvement in empirical results, the method presented is a straightforward combination of 2 prior works. The method is also unprincipled in some cases. For example, the LLM prior doesn’t always generate valid crystals in the defined domain. It may also generate lattices with angles outside the 60-120 range. The authors say that they simply reject invalid samples from LLM. It would be nice to report the rejection rate. 

In Table 1, the SUN rate is significantly lower than the stability rate for FlowLLM. It means either the model is generating non-novel or non-unique crystals. Can the authors also report the novel and unique rates?

The wdist for density is significantly higher than DiffCSP and FlowMM. It indicates that the density of the generated crystals is less accurate.
questions: questions: Have the authors looked at the RMSD of the generated crystals to their equilibrium structures? The wdist for density is high for FlowLLM. It might indicate that the model is generating structures far away from equilibrium
limitations: limitations: The authors noted the key limitation of FlowLLM is that it lacks end-to-end differentiability. However, the authors didn’t discuss potential negative societal impact of the model.
paper_summary: paper_summary: The paper proposed the FlowLLM, a hybrid approach for material generation that combines LLMs and RFMs, effectively leveraging their complementary strengths. Namely, this method generates samples from the parametric distribution using a two-step procedure, using 1) LLM and 2) RFM. By adopting this hybrid approach, experiments demonstrated that it exhibited better generation capabilities especially in terms of SUN and stability compared to the existing baseline.
strength_weakness: strength_weakness: 1. Since the output of the LLM is used as the starting point for the RFM denoising process, the generation time is longer compared to using a simple data distribution if there is no prior sampling of a large number ($N_{tr}$) of examples from the base distribution in advance. I would appreciate a comparison of the generation times between using 1) RFM with a simple base distribution and 2) FlowLLM without prior sampling from the LLM. Additionally, since the fewer required integration steps in the structural validity-number of integration steps plot could be attributed to the benefits of RFM, it would enhance credibility to compare this with FlowMM.
2. Using the trained LLM as the base distribution for the RFM necessitates additional LLM training costs.
3. The Preliminaries section of the paper overlaps significantly with that of the FlowMM paper. The structure and content of the Preliminaries section are very similar, which raises some questions.
4. It is insufficient to verify the performance of the proposed method on only the MP-20 dataset. I believe it would be fairer to cover all the datasets addressed in other baselines.
5. This appears to simply combine the models of CrystalLLM and FlowMM, making the novelty seem somewhat weak. If there are contributions beyond this combination, please let me know.
questions: questions: 1. Are there any experiments supporting the claim that LLMs are weak with continuous type data and denoising models are weak with discrete type data?
2. In FlowMM, bit representation was applied to atom types. Is there a specific reason for not using this approach?
3. Approximately what is the rate of generating invalid crystals?
4. Could you explain the process of how $\lambda_f$ and $\lambda_s$ were determined?
5. Could you please explain in more detail about the invariance of flat tori and fractional coordinates?
limitations: limitations: Please refer to the Weaknesses section.
paper_summary: paper_summary: This paper proposes using a second-step Riemannian flow matching (RFM) model, separately trained, to improve the quality of crystal structures generated by a pre-trained large language model (LLM), closely following Gruver et al. [1].

Specifically, the authors first follow Gruver et al. to fine-tune an LLM to enable the generation of crystal structures, and then train an RFM model from scratch to improve the stability, or in other words, quality, of the generated crystals.

The motivation of this work is clear. The reason for introducing this kind of post-processing machine learning model is to address the issue that LLMs cannot deal with real values, while atom positions and lattice parameters in crystals are usually real values.

The experimental results show that introducing this second-step RFM model can improve the stability of generated crystals. However, it faces limitations, such as the inability to generate materials with specific properties, as mentioned in the limitations section. Additionally, there are other well-established pipelines to refine the generated crystal structures that may need to be compared and discussed.


References:

[1] Nate Gruver et al. “Fine-Tuned Language Models Generate Stable Inorganic Materials as Text”. In: arXiv preprint arXiv:2402.04379 (2024).
strength_weakness: strength_weakness: ## Weaknesses:

1. Missing discussions with other second-stage refinement strategies that are well-established. The generated crystal structures from LLMs can also be refined by machine learning force fields (MLFF), such as M3GNet, CHGNet (which the authors have used when calculating stability), or MACE-MP. For these MLFFs, one can directly use them without training a separate RFM model from scratch to refine the generated crystal structures.

2. The contribution may be a little limited. Given that there are well-established methods for the proposed issue in this paper, the contribution of introducing an RFM model to increase stability is somewhat limited.

3. The stability rate drops significantly when removing the duplicates (to SUN rate), up to more than three times (from 17.8 to 4.92). This means a majority of the generated stable crystals are similar or the same as each other.

4. Inability to generate materials with specific properties as mentioned in the limitations section. If the ability of this model is limited to generating stable crystal structures, there are other computationally cheaper models like DiffCSP that the authors have compared with for this task. Thus, it would be beneficial to add more discussions.
questions: questions: ## Questions:
Most of my concerns are listed above in the weaknesses. Addressing the weaknesses points and the following questions may impact the final score.

1. Why not just use MLFFs like CHGNet to increase the stability of generated crystals from LLMs?

2. Are a majority of the generated stable crystals similar or the same as each other? If not, why does the stability rate drop significantly when removing the duplicates (to SUN rate), up to more than three times (from 17.8 to 4.92)?

3. Why only compare with methods without this kind of second-step refinement module?
limitations: limitations: Limitations are discussed in the main paper.

Comments: title: Thanks for the rebuttal
Thanks for your response
Thanks for the rebuttal
comment: Thanks for the detailed response and summary.

As they mentioned, it would be beneficial to further emphasize why using LLM and FlowMM together creates a synergistic effect. One question I have is regarding datasets like Carbon, which involves isotopes (though they mentioned the dataset is unstable). It would be interesting to explore whether using LLM is still advantageous for such datasets. 

Additionally, my last question was asking for a more detailed explanation of flat tori. I wanted to know the rational behind for choosing this manifold, and that has now been clarified.

I have increased my score from 5 to 6.
Thank you again for your valuable feedback! Since we are close to the end of the discussion period, we would greatly appreciate your feedback on our response. We would be happy to provide more details, or answer any additional questions.
Thanks the authors for detailed response and additional results. They've address most of my concerns, so I'd like to increase my score from 5 to 6. 

The improved RMSD is a nice result. It would be helpful if the authors can provide additional insights on the source of improvements. Can it be attributed to the better prior or the better flow matching process?

Although I understand that authors use CHGNet to get a fast result on RMSD, DFT result is preferred for the final version of the paper. 

I also appreciate the authors adding a paragraph about potential negative societal impact.
Thank you for your responses. In summary, my concerns about following points have been addressed in a proper way.

1. My concern about “why not apply post process MLFF to other baselines” has been addressed, due to the fact that when calculating SUN scores, results from all baseline methods have been further relaxed using MLFF.

2. My concern about the contribution has been addressed, since further applying a flow model as a post process refining module is indeed helpful.

My questions about the drop from stability rate to SUN rate, and the ability to do property conditioned generation have been answered, and these points seem to be valid limitations of the current proposed method.

Overall, this seems to be an interesting paper to be accepted. And I increased score from 5 to 6.