Publications

Topics:
  1. E. Amrani, L. Karlinsky, A. M. Bronstein, Sample- and parameter-efficient auto-regressive image models, Proc. CVPR, 2025 details

    Sample- and parameter-efficient auto-regressive image models

    E. Amrani, L. Karlinsky, A. M. Bronstein
    Proc. CVPR, 2025

    We introduce XTRA, a vision model pre-trained with a novel auto-regressive objective that significantly enhances both sample and parameter efficiency compared to previous auto-regressive image models. Unlike contrastive or masked image modeling methods, which have not been demonstrated as having consistent scaling behavior on unbalanced internet data, auto-regressive vision models exhibit scalable and promising performance as model and dataset size increase. In contrast to standard auto-regressive models, XTRA employs a Block Causal Mask, where each Block represents k × k tokens rather than relying on a standard causal mask. By reconstructing pixel values block by block, XTRA captures higher-level structural patterns over larger image regions. Predicting on blocks allows the model to learn relationships across broader areas of pixels, enabling more abstract and semantically meaningful representations than traditional next-token prediction. This simple modification yields two key results. First, XTRA is sample-efficient. Despite being trained on 152× fewer samples (13.1M vs. 2B), XTRA ViT-H/14 surpasses the top-1 average accuracy of the previous state-of-the-art auto-regressive model across 15 diverse image recognition benchmarks. Second, XTRA is parameter-efficient. Compared to auto-regressive models trained on ImageNet-1k, XTRA ViT-B/16 outperforms in linear and attentive probing tasks, using 7-16× fewer parameters (85M vs. 1.36B/0.63B).

    D. Freedman, E. Rozenberg, A. M. Bronstein, A theoretical framework for an efficient normalizing flow-based solution to the Schrödinger equation, Proc. AAAI, 2025 details

    A theoretical framework for an efficient normalizing flow-based solution to the Schrödinger equation

    D. Freedman, E. Rozenberg, A. M. Bronstein
    Proc. AAAI, 2025

    A central problem in quantum mechanics involves solving the Electronic Schrödinger Equation for a molecule or material. The Variational Monte Carlo approach to this problem approximates a particular variational objective via sampling, and then optimizes this approximated objective over a chosen parameterized family of wavefunctions, known as the ansatz. Recently neural networks have been used as the ansatz, with accompanying success. However, sampling from such wavefunctions has required the use of a Markov Chain Monte Carlo approach, which is inherently inefficient. In this work, we propose a solution to this problem via an ansatz which is cheap to sample from, yet satisfies the requisite quantum mechanical properties. We prove that a normalizing flow using the following two essential ingredients satisfies our requirements: (a) a base distribution which is constructed from Determinantal Point Processes; (b) flow layers which are equivariant to a particular subgroup of the permutation group. We then show how to construct both continuous and discrete normalizing flows which satisfy the requisite equivariance. We further demonstrate the manner in which the non-smooth nature (“cusps”) of the wavefunction may be captured, and how the framework may be generalized to provide induction across multiple molecules. The resulting theoretical framework entails an efficient approach to solving the Electronic Schrödinger Equation.

    A. Maddipatla, N. Bojan Sellam, M. Bojan, S. Vedula, P. Schanda, A. Marx, A. M. Bronstein, Inverse problems with experiment-guided AlphaFold, arXiv:2502.09372, 2025 details

    Inverse problems with experiment-guided AlphaFold

    A. Maddipatla, N. Bojan Sellam, M. Bojan, S. Vedula, P. Schanda, A. Marx, A. M. Bronstein
    arXiv:2502.09372, 2025

    Proteins exist as a dynamic ensemble of multiple conformations, and these motions are often crucial for their functions. However, current structure prediction methods predominantly yield a single conformation, overlooking the conformational heterogeneity revealed by diverse experimental modalities. Here, we present a framework for building experiment-grounded protein structure generative models that infer conformational ensembles consistent with measured experimental data. The key idea is to treat state-of-the-art protein structure predictors (e.g., AlphaFold3) as sequence-conditioned structural priors, and cast ensemble modeling as posterior inference of protein structures given experimental measurements. Through extensive real-data experiments, we demonstrate the generality of our method to incorporate a variety of experimental measurements. In particular, our framework uncovers previously unmodeled conformational heterogeneity from crystallographic densities, and generates high-accuracy NMR ensembles orders of magnitude faster than the status quo. Notably, we demonstrate that our ensembles outperform AlphaFold3 and sometimes better fit experimental data than publicly deposited structures to the Protein Data Bank (PDB). We believe that this approach will unlock building predictive models that fully embrace experimentally observed conformational diversity.

    Y. Davidson, A. Philipp, S. Chakraborty, A. M. Bronstein, R. Gershoni-Poranne, How local is "local"? Deep learning reveals locality of the induced magnetic field of polycyclic aromatic hydrocarbons, chemrxiv 10.26434/chemrxiv-2025-pqmcc, 2025 details

    How local is "local"? Deep learning reveals locality of the induced magnetic field of polycyclic aromatic hydrocarbons

    Y. Davidson, A. Philipp, S. Chakraborty, A. M. Bronstein, R. Gershoni-Poranne
    chemrxiv 10.26434/chemrxiv-2025-pqmcc, 2025

    We investigate the locality of magnetic response in polycyclic aromatic molecules using a novel deep-learning approach. Our method employs graph neural networks (GNNs) with a graph-of-rings representation to predict Nucleus-Independent Chemical Shifts in the space around the molecule. We train a series of models, each time reducing the size of the largest molecules used in training. The accuracy of prediction remains high (MAE < 0.5 ppm), even when training the model only on molecules with up to 4 rings, thus providing strong evidence for the locality of magnetic response. To overcome the known problem of generalization of GNNs, we implement a k-hop expansion strategy and succeed in achieving accurate predictions for molecules with up to 15 rings (almost 4 times the size of the largest training example). Our findings have implications for understanding the magnetic response in complex molecules and demonstrate a promising approach to overcoming GNN scalability limitations. Furthermore, the trained models enable rapid characterization, without the need for more expensive DFT calculations.