Atri Rudra

Large language models (LLMs) store factual knowledge in their parameters, yet the underlying mechanism remains unclear. Recent work has found that factual knowledge resides in LLM MLPs as key-value mappings and has provided closed-form fact-storing MLP constructions. However, existing constructive and mechanistic interpretability approaches fail to capture three empirically observed properties of MLPs in LLMs: optimal fact-storage capacity scaling, ability to handle arbitrary embedding geometries, and usability within Transformer blocks. Towards understanding fact-storage in LLMs, we analyze the decoding margin of MLPs, whereas prior work only studies MLP fact storage. Our analysis allows us to develop the first closed-form MLP construction that realizes all three properties. Under isotropic embeddings, our construction achieves optimal storage capacity scaling, reducing the capacity gap to trained MLPs by 8-16x over prior constructions. Further, our theoretical margin scaling bounds precisely characterize how embedding geometry governs empirical margin (R^2 > 0.95) and illustrate our construction’s ability to store fact-sets under arbitrary embedding geometries. Moreover, we show that MLPs can be used within Transformer blocks for factual recall tasks at optimal capacity scaling, reducing the capacity gap to trained MLPs by up to 100x over prior constructions. Finally, as a proof-of-concept, we show that fact-storing MLPs enable modular fact editing by swapping a Transformer's MLP with a new one.

Sequence models face stark tradeoffs between recall quality and memory efficiency. The ability to use information over long sequences (i.e. recall) is critical for sequence modeling tasks ranging from information extraction to reasoning. Prior work has shown that in theory, *linear* attention models with sufficient recurrent state sizes can expand the Pareto frontier of the recall-memory tradeoff space beyond alternative architectures such as softmax attention and state space models. However, it is difficult to scale the linear attention state size due to hardware bottlenecks. I/O-aware algorithms store the linear attention states in thread registers, however state sizes beyond $\approx 3$ megabytes exhaust register memory and trigger expensive register spills. In this work, we introduce Cylon, a hardware-aware modification to linear attention architectures which partitions the recurrent state across the registers of multiple GPU processors and asynchronously combines the partitions. When applying Cylon to popular architectures, such as Hedgehog, Mamba-2, and Gated Deltanet, we unlock up to $6\times$ higher forward pass throughput and $10\times$ higher backward pass throughput compared to prior linear attention algorithms on Hopper and Blackwell GPUs. Moreover, Cylon preserves recall quality compared to non-partitioned baselines. Finally, Cylon makes large states available to model designers by unlocking sizes ($\geq 131$ MB) which are not achievable by the existing linear attention kernels.