Search for: All records

Large-scale pre-training has fundamentally changed how machine learning research is done today: large foundation models are trained once, and then can be used by anyone in the community (including those without data or compute resources to train a model from scratch) to adapt and fine-tune to specific tasks. Applying this same framework to reinforcement learning (RL) is appealing because it offers compelling avenues for addressing core challenges in RL, including sample efficiency and robustness. However, there remains a fundamental challenge to pre-train large models in the context of RL: actions have long-term dependencies, so training a foundation model that reasons across time is important. Recent advances in generative AI have provided new tools for modeling highly complex distributions. In this paper, we build a probabilistic model to predict which states an agent will visit in the temporally distant future (i.e., an occupancy measure) using flow matching. As large datasets are often constructed by many distinct users performing distinct tasks, we include in our model a latent variable capturing the user intention. This intention increases the expressivity of our model, and enables adaptation with generalized policy improvement. We call our proposed method intention-conditioned flow occupancy models (InFOM). Comparing with alternative methods for pre-training, our experiments on 36 state-based and 4 image-based benchmark tasks demonstrate that the proposed method achieves 1.8x median improvement in returns and increases success rates by 36%. 

The vast majority of visible matter in our universe comes from protons and neutrons (the nucleons). Nucleon interactions are fundamental to how the universe developed after the Big Bang and govern all nuclear phenomena. The subtle balance in how two nucleons interact shapes the universe's hydrogen content that is central to our existence. Our objective is to compute the interaction strength while varying the parameters of nature to understand how delicate this balance is. We developed a new code using sophisticated physics algorithms and a highly optimized library for simulations on CPU-GPU parallel architectures. It has excellent weak scaling and impressive linear scaling for a fixed problem size with increasing number of nodes up to El Capitan's full ∼11,000 nodes. On Alps, El Capitan, Frontier, Jupiter, and Perlmutter supercomputers we achieve a maximum disruptive speed-up of ∼240 times the previous state-of-the-art, signaling a new era of supercomputing. 

The infinite-layer compound ${\mathrm{CaCoO}}_2$ provides a compelling platform to explore how local orbital distortions evolve into collective electronic order. Using resonant soft x-ray scattering (RSXS) and its time-resolved counterpart (tr-RSXS), we reveal a cooperative coupling between two intertwined orders at the same in-plane wave vector, $q\approx (1/4,1/4,0)$ , comprising a quasi-two-dimensional orbital order associated with the distorted ${\mathrm{Co}}_{\left(1\right)}$ sites and a quasi-three-dimensional trimerlike charge order stabilized through Co-O-Ca hybridization. The charge-sensitive component exhibits a uniform relaxation time of $\sim 400\mathrm{fs}$ ( $391\pm 5\mathrm{fs}$ ) across multiple Co resonances, evidencing a collective charge response. A comparable timescale observed in the orbital order associated with the highly distorted ${\mathrm{Co}}_{\left(1\right)}$ sites indicates that this collective behavior naturally originates from the same orbital framework that mediates both charge and orbital degrees of freedom. Together, these findings establish ${\mathrm{CaCoO}}_2$ as a model system in which Jahn-Teller distortions, orbital hybridization, and geometric frustration cooperate to form molecular trimer units that dictate the macroscopic electronic behavior. 

Large language models (LLMs) are notoriously memory-intensive during training, particularly with the popular AdamW optimizer. This memory burden necessitates using more or higher-end GPUs or reducing batch sizes, limiting training scalability and throughput. To address this, various memory-efficient optimizers have been proposed to reduce optimizer memory usage. However, they face critical challenges: (i) reliance on costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial optimizer memory overhead to maintain competitive performance. In this work, we identify that AdamW's learning rate adaptation rule can be effectively coarsened as a structured learning rate update. Based on this insight, we propose Approximated Gradient Scaling for Memory-Efficient LLM Optimization (APOLLO), which approximates learning rate scaling using an auxiliary low-rank optimizer state based on pure random projection. This structured learning rate update rule makes APOLLO highly tolerant to further memory reductions while delivering comparable pre-training performance. Even its rank-1 variant, APOLLO-Mini, achieves superior pre-training performance compared to AdamW with SGD-level memory costs. Extensive experiments demonstrate that the APOLLO series performs on-par with or better than AdamW, while achieving greater memory savings by nearly eliminating the optimization states of AdamW. These savings provide significant system-level benefits: (1) Enhanced Throughput: 3x throughput on an 8xA100-80GB setup compared to AdamW by supporting 4x larger batch sizes. (2) Improved Model Scalability: Pre-training LLaMA-13B with naive DDP on A100-80GB GPUs without system-level optimizations. (3) Low-End GPU Friendly Pre-training: Pre-training LLaMA-7B on a single GPU using less than 12 GB of memory with weight quantization.