Publications

Recent probing studies reveal that large language models exhibit linear subspaces that separate true from false statements, yet the mechanism behind their emergence is unclear. We introduce a transparent, one-layer transformer toy model that reproduces such truth subspaces end-to-end and exposes one concrete route by which they can arise. We study one simple setting in which truth encoding can emerge: a data distribution where factual statements co-occur with other factual statements (and vice-versa), encouraging the model to learn this distinction in order to lower the LM loss on future tokens. We corroborate this pattern with experiments in pretrained language models. Finally, in the toy setting we observe a two-phase learning dynamic: networks first memorize individual factual associations in a few steps, then---over a longer horizon---learn to linearly separate true from false, which in turn lowers language-modeling loss. Together, these results provide both a mechanistic demonstration and an empirical motivation for how and why linear truth representations can emerge in language models.

Graph Neural Networks (GNNs) typically scale with the number of graph edges, making them well suited for sparse graphs but less efficient on dense graphs, such as point clouds or molecular interactions. A common remedy is to sparsify the graph via similarity thresholding or distance pruning, but this forces an arbitrary choice of a single interaction scale and discards crucial information from other scales. To overcome this limitation, we introduce a multi-view graph-tuple framework. Instead of a single graph, our graph-tuple framework partitions the graph into disjoint subgraphs, capturing primary local interactions and weaker, long-range connections. We then learn multi-view representations from the graph-tuple via a heterogeneous message-passing architecture inspired by the theory of non-commuting operators, which we formally prove is strictly more expressive and guarantees a lower oracle risk compared to single-graph message-passing models. We instantiate our framework on two scientific domains: molecular property prediction from feature-scarce Coulomb matrices and cosmological parameter inference from geometric point clouds. On both applications, our multi-view graph-tuple models demonstrate better performance than single-graph baselines, highlighting the power and versatility of our multi-view approach.

We introduce a numerical method for computing spectral densities, and apply it to the evaluation of the local density of states (LDOS) of sparse Hamiltonians derived from tight-binding models. The approach, which we call the high-order delta-Chebyshev method, can be viewed as a variant of the popular regularized Chebyshev kernel polynomial method (KPM), but it uses a high-order accurate approximation of the δ-function to achieve rapid convergence to the thermodynamic limit for smooth spectral densities. The costly computational steps are identical to those for KPM, with high-order accuracy achieved by an inexpensive post-processing procedure. We apply the algorithm to tight-binding models of graphene and twisted bilayer graphene, demonstrating high-order convergence to the LDOS at non-singular points.

Markov state models (MSMs) are widely employed to analyze the kinetics of complex systems. But despite their effectiveness in many applications, MSMs are prone to systematic or statistical errors, often exacerbated by suboptimal hyperparameter choice. In this article, we attempt to understand how these choices affect the error of estimates of mean first-passage times and committors, key quantities in chemical rate theory. We first evaluate the performance of the recently introduced “stopped-process estimator” [Strahan, J. Long-time-scale predictions from short-trajectory data: A benchmark analysis of the trp-cage miniprotein. J. Chem. Theory Comput. 2021, 17, 2948–2963. 10.1021/acs.jctc.0c00933.] that attempts to reduce error caused by choosing a too-large lag time. We then study the effect of statistical errors on Markov state model construction using the condition number, which measures an MSM’s sensitivity to perturbation. This analysis helps give an insight into which factors cause an MSM to be more or less sensitive to statistical error. Our work highlights the importance of choosing a good sampling measure, the measure from which the initial points are drawn, and has implications for recent work applying a variational principle for evaluating the committor.

We present a family of integral equation-based solvers for the heat equation, reaction-diffusion systems, the unsteady Stokes equation and the incompressible Navier-Stokes equations in two space dimensions. Our emphasis is on the development of methods that can efficiently follow complex solution features in space-time by refinement and coarsening at each time step on an adaptive quadtree. For simplicity, we focus on problems posed in a square domain with periodic boundary conditions. The performance and robustness of the methods are illustrated with several numerical examples.

The randomized SVD is a method to compute an inexpensive, yet accurate, low-rank approximation of a matrix. The algorithm assumes access to the matrix through matrix-vector products (matvecs). Therefore, when we would like to apply the randomized SVD to a matrix function, f(A), one needs to approximate matvecs with f(A) using some other algorithm, which is typically treated as a black-box. Chen and Hallman (SIMAX 2023) argued that, in the common setting where matvecs with f(A) are approximated using Krylov subspace methods (KSMs), a more efficient low-rank approximation is possible if we open this black-box. They present an alternative approach that significantly outperforms the naive combination of KSMs with the randomized SVD, although the method lacked theoretical justification. In this work, we take a closer look at the method, and provide strong and intuitive error bounds that justify its excellent performance for low-rank approximation of matrix functions.

We present a spectrally accurate fast algorithm for evaluating the solution to the scalar wave equation in free space driven by a large collection of point sources in a bounded domain. With $M$ sources temporally discretized by $N_t$ time steps of size $\Delta t$, a naive potential evaluation at $M$ targets on the same time grid requires $\mathcal O(M^2 N_t)$ work. Our scheme requires $\mathcal{O}\left((M + N^3\log N)N_t\right)$ work, where $N$ scales as $\mathcal O(1/\Delta t)$, i.e., the maximum signal frequency. This is achieved by using the recently-proposed windowed Fourier projection (WFP) method to split the potential into a local part, evaluated directly, plus a smooth history part approximated by an $N^3$-point equispaced discretization of the Fourier transform, where each Fourier coefficient obeys a simple recursion relation. The growing oscillations in the spectral representation (which would be present with a naive use of the Fourier transform) are controlled by spatially truncating the hyperbolic Green's function itself. Thus, the method avoids the need for absorbing boundary conditions. We demonstrate the performance of our algorithm with up to a million sources and targets at 6-digit accuracy. We believe it can serve as a key component in addressing time-domain wave equation scattering problems.

Recent advances in mechanistic interpretability have revealed that large language models (LLMs) develop internal representations corresponding not only to concrete entities but also distinct, human-understandable abstract concepts and behaviour. Moreover, these hidden features can be directly manipulated to steer model behaviour. However, it remains an open question whether this phenomenon is unique to models trained on inherently structured data (ie. language, images) or if it is a general property of foundation models. In this work, we investigate the internal representations of a large physics-focused foundation model. Inspired by recent work identifying single directions in activation space for complex behaviours in LLMs, we extract activation vectors from the model during forward passes over simulation datasets for different physical regimes. We then compute "delta" representations between the two regimes. These delta tensors act as concept directions in activation space, encoding specific physical features. By injecting these concept directions back into the model during inference, we can steer its predictions, demonstrating causal control over physical behaviours, such as inducing or removing some particular physical feature from a simulation. These results suggest that scientific foundation models learn generalised representations of physical principles. They do not merely rely on superficial correlations and patterns in the simulations. Our findings open new avenues for understanding and controlling scientific foundation models and has implications for AI-enabled scientific discovery

Stellarators are a prospective class of fusion-based power plants that confine a hot plasma with three-dimensional magnetic fields. Typically framed as a PDE-constrained optimization problem, stellarator design is a time-consuming process that can take hours to solve on a computing cluster. Developing fast methods for designing stellarators is crucial for advancing fusion research. Given the recent development of large datasets of optimized stellarators, machine learning approaches have emerged as a potential candidate. Motivated by this, we present an open inverse problem to the machine learning community: to rapidly generate high-quality stellarator designs which have a set of desirable characteristics. As a case study in the problem space, we train a conditional diffusion model on data from the QUASR database to generate quasisymmetric stellarator designs with desirable characteristics (aspect ratio and mean rotational transform). The diffusion model is applied to design stellarators with characteristics not seen during training. We provide evaluation protocols and show that many of the generated stellarators exhibit solid performance: less than 5% deviation from quasisymmetry and the target characteristics. The modest deviation from quasisymmetry highlights an opportunity to reach the sub 1% target. Beyond the case study, we share multiple promising avenues for generative modeling to advance stellarator design.

Foundation models have transformed machine learning for language and vision, but achieving comparable impact in physical simulation remains a challenge. Data heterogeneity and unstable long-term dynamics inhibit learning from sufficiently diverse dynamics, while varying resolutions and dimensionalities challenge efficient training on modern hardware. Through empirical and theoretical analysis, we incorporate new approaches to mitigate these obstacles, including a harmonic-analysis-based stabilization method, load-balanced distributed 2D and 3D training strategies, and compute-adaptive tokenization. Using these tools, we develop Walrus, a transformer-based foundation model developed primarily for fluid-like continuum dynamics. Walrus is pretrained on nineteen diverse scenarios spanning astrophysics, geoscience, rheology, plasma physics, acoustics, and classical fluids. Experiments show that Walrus outperforms prior foundation models on both short and long term prediction horizons on downstream tasks and across the breadth of pretraining data, while ablation studies confirm the value of our contributions to forecast stability, training throughput, and transfer performance over conventional approaches. Code and weights are released for community use.