arXiv Science⌕ Search

arXiv subjects

Shriram Chennakesavalu

Publications and source records attributed to Shriram Chennakesavalu.

11 recordsLinked to original sources

Training Large Language Models for Small-Molecule Design with Synthetic Task Scaling

Designing viable drug candidates requires searching a combinatorially large and rugged chemical space for molecules that satisfy multiple, often competing, objectives. Large language models (LLMs) provide a useful generative prior for this problem because of their representational capacity, reasoning ability, and flexibility when incorporating information from the external environment. While reinforcement learning from verifiable rewards (RLVR) can be used to improve the capabilities of LLMs, many chemically relevant scoring functions require hours or even days per evaluation, making them prohibitively expensive to use directly during online training. Here, we investigate whether LLMs can learn molecular design strategies from cheaper synthetic tasks that generalize to expensive molecular lead optimization settings. We find that curriculum-based training recipes that gradually incorporate more challenging synthetic design tasks enable strong performance that surpasses that of much larger frontier models on structure-based lead optimization. Our results suggest that scaling post-training using synthetic tasks is an effective strategy for adapting LLMs to high-cost experimental scenarios that are too expensive to directly train on.

cs.LG↗

Frontier LLMs are effective batch optimizers: Assessing reasoning models in continuous and discrete settings

Frontier large language models (LLMs) have become attractive priors for optimization due to their large-scale pretraining that enables them to navigate a variety of optimization settings. However, the effectiveness of modern reasoning LLMs in batch optimization settings remains underexplored. Here we investigate the performance of the current generation of frontier LLMs as batch optimizers in both continuous and discrete settings. We find that while LLMs are competitive zero-shot batch optimizers for numerical test functions, their performance is brittle compared to classical non-LLM optimization approaches. However, LLM priors are significantly better in semantically rich settings, indicating that their batch optimization behavior is highly effective when navigating and reasoning over the discrete spaces most similar in structure to their pretraining data.

cs.LG↗

Evaluating the Progression of Large Language Model Capabilities for Small-Molecule Drug Design

Large Language Models (LLMs) have the potential to accelerate small molecule drug design due to their ability to reason about information from diverse sources and formats. However, their practical utility remains unclear due to the lack of benchmarks that reflect real-world scenarios. In this work, we introduce a suite of chemically-grounded tasks spanning molecular property prediction, molecular representation transformations, and molecular design. Importantly, we formulate these tasks as reinforcement learning (RL) environments, enabling a unified approach for evaluation and post-training. Across three model families, we find that frontier models are increasingly proficient at chemical tasks, but that there is significant room for improvement, especially in experimental settings with low data. Critically, we show that RL-based post-training can substantially improve performance. A smaller model post-trained on our environments becomes competitive with state-of-the-art frontier models, despite a significantly weaker base model. This suggests a practical route toward employing LLMs in drug discovery; by combining carefully-designed evaluation tasks with targeted post-training, we can both elucidate and close critical capability gaps.

cs.LG↗

Scaling Transferable Coarse-graining with Mean Force Matching

Coarse-grained molecular dynamics often sacrifices accuracy and transferability for computational efficiency, but the use of machine learned potentials is helping coarse-grained models attain performance on par with atomistic molecular dynamics. Nevertheless, developing representations of the coarse-grained potential energy surface faces severe scaling challenges due to the extreme data demands of widely used "bottom-up" coarse-graining objectives. In this work, we show that mean force matching, a strategy for training thermodynamically consistent coarse-grained models, requires 50x fewer training samples and 87% less total atomistic simulation time, while obtaining better accuracy on the potential of mean force for unseen proteins compared to other commonly used objectives. By systematically removing noise from the objective function, we demonstrate that it is possible to scale machine learning architectures for coarse-graining, enabling highly accurate and transferable models. We show the advantages of mean force matching both theoretically and through exhaustive benchmarking using thermodynamic consistency as the primary metric of accuracy.

physics.chem-ph↗

Aligning Transformers with Continuous Feedback via Energy Rank Alignment

Searching through chemical space is an exceptionally challenging problem because the number of possible molecules grows combinatorially with the number of atoms. Large, autoregressive models trained on databases of chemical compounds have yielded powerful generators, but we still lack robust strategies for generating molecules with desired properties. This molecular search problem closely resembles the "alignment" problem for large language models, though for many chemical tasks we have a specific and easily evaluable reward function. Here, we introduce an algorithm called energy rank alignment (ERA) that leverages an explicit reward function to produce a gradient-based objective that we use to optimize autoregressive policies. We show theoretically that this algorithm is closely related to proximal policy optimization (PPO) and direct preference optimization (DPO), but has a minimizer that converges to an ideal Gibbs-Boltzmann distribution with the reward playing the role of an energy function. Furthermore, this algorithm is highly scalable, does not require reinforcement learning, and performs well relative to DPO when the number of preference observations per pairing is small. We deploy this approach to align molecular transformers and protein language models to generate molecules and protein sequences, respectively, with externally specified properties and find that it does so robustly, searching through diverse parts of chemical space.

cs.LG↗

Data-efficient generation of protein conformational ensembles with backbone-to-side chain transformers

Excitement at the prospect of using data-driven generative models to sample configurational ensembles of biomolecular systems stems from the extraordinary success of these models on a diverse set of high-dimensional sampling tasks. Unlike image generation or even the closely related problem of protein structure prediction, there are not currently data sources with sufficient breadth to parameterize generative models for conformational ensembles. To enable discovery, a fundamentally different approach to building generative models is required: models should be able to propose rare, albeit physical, conformations that may not arise in even the largest data sets. Here we introduce a modular strategy to generate conformations based on ``backmapping'' from a fixed protein backbone that 1) maintains conformational diversity of the side chains and 2) couples the side chain fluctuations using global information about the protein conformation. Our model combines simple statistical models of side chain conformations based on rotamer libraries with the now ubiquitous transformer architecture to sample with atomistic accuracy. Together, these ingredients provide a strategy for rapid data acquistion and hence a crucial ingredient for scalable physical simulation with generative neural networks.

cond-mat.stat-mech↗

Adaptive nonequilibrium design of actin-based metamaterials: fundamental and practical limits of control

The adaptive and surprising emergent properties of biological materials self-assembled in far-from-equilibrium environments serve as an inspiration for efforts to design nanomaterials and their properties. In particular, controlling the conditions of self-assembly can modulate material properties, but there is no systematic understanding of either how to parameterize this control or how \emph{controllable} a given material can be. Here, we demonstrate that branched actin networks can be encoded with \textit{metamaterial} properties by dynamically controlling the applied force under which they grow, and that the protocols can be selected using multi-task reinforcement learning. These actin networks have tunable responses over a large dynamic range depending on the chosen external protocol, providing a pathway to encoding ``memory'' within these structures. Interestingly, we show that encoding memory requires dissipation and the rate of encoding is constrained by the flow of entropy -- both physical and information theoretical. Taken together, these results emphasize the utility and necessity of nonequilibrium control for designing self-assembled nanostructures.

cond-mat.stat-mech↗

A unified, geometric framework for nonequilibrium protocol optimization

Controlling thermodynamic cycles to minimize the dissipated heat is a longstanding goal in thermodynamics, and more recently, a central challenge in stochastic thermodynamics for nanoscale systems. Here, we introduce a theoretical and computational framework for optimizing nonequilibrium control protocols that can transform a system between two distributions in a minimally dissipative fashion. These protocols optimally transport a system along paths through the space of probability distributions that minimize the dissipative cost of a transformation. Furthermore, we show that the thermodynamic metric -- determined via a linear response approach -- can be directly derived from the same objective function that is optimized in the optimal transport problem, thus providing a unified perspective on thermodynamic geometries. We investigate this unified geometric framework in two model systems and observe that our procedure for optimizing control protocols is robust beyond linear response.

cond-mat.stat-mech↗

Ensuring thermodynamic consistency with invertible coarse-graining

Coarse-grained models are a core computational tool in theoretical chemistry and biophysics. A judicious choice of a coarse-grained model can yield physical insight by isolating the essential degrees of freedom that dictate the thermodynamic properties of a complex, condensed-phase system. The reduced complexity of the model typically leads to lower computational costs and more efficient sampling compared to atomistic models. Designing ``good'' coarse-grained models is an art. Generally, the mapping from fine-grained configurations to coarse-grained configurations itself is not optimized in any way; instead, the energy function associated with the mapped configurations is. In this work, we explore the consequences of optimizing the coarse-grained representation alongside its potential energy function. We use a graph machine learning framework to embed atomic configurations into a low dimensional space to produce efficient representations of the original molecular system. Because the representation we obtain is no longer directly interpretable as a real space representation of the atomic coordinates, we also introduce an inversion process and an associated thermodynamic consistency relation that allows us to rigorously sample fine-grained configurations conditioned on the coarse-grained sampling. We show that this technique is robust, recovering the first two moments of the distribution of several observables in proteins such as chignolin and alanine dipeptide.

cond-mat.stat-mech↗

Cooperative multi-agent reinforcement learning for high-dimensional nonequilibrium control

Experimental advances enabling high-resolution external control create new opportunities to produce materials with exotic properties. In this work, we investigate how a multi-agent reinforcement learning approach can be used to design external control protocols for self-assembly. We find that a fully decentralized approach performs remarkably well even with a "coarse" level of external control. More importantly, we see that a partially decentralized approach, where we include information about the local environment allows us to better control our system towards some target distribution. We explain this by analyzing our approach as a partially-observed Markov decision process. With a partially decentralized approach, the agent is able to act more presciently, both by preventing the formation of undesirable structures and by better stabilizing target structures as compared to a fully decentralized approach.

cond-mat.stat-mech↗

Probing the Theoretical and Computational Limits of Dissipative Design

Self-assembly, the process by which interacting components form well-defined and often intricate structures, is typically thought of as a spontaneous process arising from equilibrium dynamics. When a system is driven by external \emph{nonequilibrium} forces, states statistically inaccessible to the equilibrium dynamics can arise, a process sometimes termed direct self-assembly. However, if we fix a given target state and a set of external control variables, it is not well-understood i) how to design a protocol to drive the system towards the desired state nor ii) the energetic cost of persistently perturbing the stationary distribution. Here we derive a bound that relates the proximity to the chosen target with the dissipation associated with the external drive, showing that high-dimensional external control can guide systems towards target distribution but with an inevitable entropic cost. Remarkably, the bound holds arbitrarily far from equilibrium. Secondly, we investigate the performance of deep reinforcement learning algorithms and provide evidence for the realizability of complex protocols that stabilize otherwise inaccessible states of matter.

cond-mat.stat-mech↗