arXiv subjects

Simran Arora

Publications and source records attributed to Simran Arora.

At least 37 records · Page 2Linked to original sources

arXiv2025

Cartridges: Lightweight and general-purpose long context representations via self-study

Large language models are often used to answer queries grounded in large text corpora (e.g. codebases, legal documents, or chat histories) by placing the entire corpus in the context window and leveraging in-context learning (ICL). Although current models support contexts of 100K-1M tokens, this setup is costly to serve because the memory consumption of the KV cache scales with input length. We explore an alternative: training a smaller KV cache offline on each corpus. At inference time, we load this trained KV cache, which we call a Cartridge, and decode a response. Critically, the cost of training a Cartridge can be amortized across all the queries referencing the same corpus. However, we find that the naive approach of training the Cartridge with next-token prediction on the corpus is not competitive with ICL. Instead, we propose self-study, a training recipe in which we generate synthetic conversations about the corpus and train the Cartridge with a context-distillation objective. We find that Cartridges trained with self-study replicate the functionality of ICL, while being significantly cheaper to serve. On challenging long-context benchmarks, Cartridges trained with self-study match ICL performance while using 38.6x less memory and enabling 26.4x higher throughput. Self-study also extends the model's effective context length (e.g. from 128k to 484k tokens on MTOB) and surprisingly, leads to Cartridges that can be composed at inference time without retraining.

cs.CL↗

arXiv2025

Towards Learning High-Precision Least Squares Algorithms with Sequence Models

This paper investigates whether sequence models can learn to perform numerical algorithms, e.g. gradient descent, on the fundamental problem of least squares. Our goal is to inherit two properties of standard algorithms from numerical analysis: (1) machine precision, i.e. we want to obtain solutions that are accurate to near floating point error, and (2) numerical generality, i.e. we want them to apply broadly across problem instances. We find that prior approaches using Transformers fail to meet these criteria, and identify limitations present in existing architectures and training procedures. First, we show that softmax Transformers struggle to perform high-precision multiplications, which prevents them from precisely learning numerical algorithms. Second, we identify an alternate class of architectures, comprised entirely of polynomials, that can efficiently represent high-precision gradient descent iterates. Finally, we investigate precision bottlenecks during training and address them via a high-precision training recipe that reduces stochastic gradient noise. Our recipe enables us to train two polynomial architectures, gated convolutions and linear attention, to perform gradient descent iterates on least squares problems. For the first time, we demonstrate the ability to train to near machine precision. Applied iteratively, our models obtain 100,000x lower MSE than standard Transformers trained end-to-end and they incur a 10,000x smaller generalization gap on out-of-distribution problems. We make progress towards end-to-end learning of numerical algorithms for least squares.

cs.LG↗

arXiv2025

Language Models Enable Simple Systems for Generating Structured Views of Heterogeneous Data Lakes

A long standing goal of the data management community is to develop general, automated systems that ingest semi-structured documents and output queryable tables without human effort or domain specific customization. Given the sheer variety of potential documents, state-of-the art systems make simplifying assumptions and use domain specific training. In this work, we ask whether we can maintain generality by using large language models (LLMs). LLMs, which are pretrained on broad data, can perform diverse downstream tasks simply conditioned on natural language task descriptions. We propose and evaluate EVAPORATE, a simple, prototype system powered by LLMs. We identify two fundamentally different strategies for implementing this system: prompt the LLM to directly extract values from documents or prompt the LLM to synthesize code that performs the extraction. Our evaluations show a cost-quality tradeoff between these two approaches. Code synthesis is cheap, but far less accurate than directly processing each document with the LLM. To improve quality while maintaining low cost, we propose an extended code synthesis implementation, EVAPORATE-CODE+, which achieves better quality than direct extraction. Our key insight is to generate many candidate functions and ensemble their extractions using weak supervision. EVAPORATE-CODE+ not only outperforms the state-of-the art systems, but does so using a sublinear pass over the documents with the LLM. This equates to a 110x reduction in the number of tokens the LLM needs to process, averaged across 16 real-world evaluation settings of 10k documents each.

cs.CL↗

arXiv2025

Simple linear attention language models balance the recall-throughput tradeoff

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.

cs.CL↗

arXiv2025

LoLCATs: On Low-Rank Linearizing of Large Language Models

Recent works show we can linearize large language models (LLMs) -- swapping the quadratic attentions of popular Transformer-based LLMs with subquadratic analogs, such as linear attention -- avoiding the expensive pretraining costs. However, linearizing LLMs often significantly degrades model quality, still requires training over billions of tokens, and remains limited to smaller 1.3B to 7B LLMs. We thus propose Low-rank Linear Conversion via Attention Transfer (LoLCATs), a simple two-step method that improves LLM linearizing quality with orders of magnitudes less memory and compute. We base these steps on two findings. First, we can replace an LLM's softmax attentions with closely-approximating linear attentions, simply by training the linear attentions to match their softmax counterparts with an output MSE loss ("attention transfer"). Then, this enables adjusting for approximation errors and recovering LLM quality simply with low-rank adaptation (LoRA). LoLCATs significantly improves linearizing quality, training efficiency, and scalability. We significantly reduce the linearizing quality gap and produce state-of-the-art subquadratic LLMs from Llama 3 8B and Mistral 7B v0.1, leading to 20+ points of improvement on 5-shot MMLU. Furthermore, LoLCATs does so with only 0.2% of past methods' model parameters and 0.4% of their training tokens. Finally, we apply LoLCATs to create the first linearized 70B and 405B LLMs (50x larger than prior work). When compared with prior approaches under the same compute budgets, LoLCATs significantly improves linearizing quality, closing the gap between linearized and original Llama 3.1 70B and 405B LLMs by 77.8% and 78.1% on 5-shot MMLU.

cs.LG↗

arXiv2025

KernelBench: Can LLMs Write Efficient GPU Kernels?

Efficient GPU kernels are crucial for building performant machine learning architectures, but writing them is a time-consuming challenge that requires significant expertise; therefore, we explore using language models (LMs) to automate kernel generation. We introduce KernelBench, an open-source framework for evaluating LMs' ability to write fast and correct kernels on a suite of 250 carefully selected PyTorch ML workloads. KernelBench represents a real-world engineering environment and making progress on the introduced benchmark directly translates to faster practical kernels. We introduce a new evaluation metric fast_p, which measures the percentage of generated kernels that are functionally correct and offer a speedup greater than an adjustable threshold p over baseline. Our experiments across various state-of-the-art models and test-time methods show that frontier reasoning models perform the best out of the box but still fall short overall, matching the PyTorch baseline in less than 20% of the cases. While we show that results can improve by leveraging execution and profiling feedback during iterative refinement, KernelBench remains a challenging benchmark, with its difficulty increasing as we raise speedup threshold p.

cs.LG↗

arXiv2025

TeV Scale Resonant Leptogenesis with triplet Fermion in Connection to Muon $g-2$

We propose an extension of the minimal scotogenic model with a triplet fermion and a singlet scalar. An imposed $Z_{4}\times Z_{2}$ symmetry allows only diagonal Yukawa couplings among different generations of SM leptons and right-handed singlet neutrinos. The Yukawa coupling of the triplet fermion with the inert doublet positively contributes to the muon anomalous magnetic moment. The imposed $Z_{4}\times Z_{2}$ symmetry forbids the conventional leptogenesis from the lightest right-handed neutrino decay. A net lepton asymmetry can be generated in the muonic sector from $N_{2}$ and triplet fermion decay through resonant leptogenesis scenario. The Yukawa coupling of triplet plays significant role both in leptogenesis and in the anomalous magnetic moment of the muon. We show a viable parameter space for TeV scale leptogenesis while explaining the Fermi lab results. The inert scalar is the dark matter candidate in this model. The Muon $(g-2)$ and dark matter both favor the same parameter space for mass of the dark matter and the triplet fermion.

hep-ph↗

arXiv2025

Cosmological dynamics of interacting dark energy and dark matter in $f(Q)$ gravity

In this work, we explore the behavior of interacting dark energy and dark matter within a model of $f(Q)$ gravity, employing a standard framework of dynamical system analysis. We consider the power-law $f(Q)$ model incorporating with two different forms of interacting dark energy and dark matter: $3αHρ_m$ and $\fracα{3H}ρ_m ρ_{DE}$. The evolution of $Ω_m, Ω_r, Ω_{DE}, q$, and $ω$ for different values of the model parameter $n$ and the interaction parameter $α$ has been examined. Our results show that the universe was dominated by matter in the early stages and will be dominated by dark energy in later stages. Using the observational data, the fixed points are found to be stable and can be represented the de Sitter and quintessence acceleration solutions. We discover that the dynamical profiles of the universe in $f(Q)$ dark energy models are influenced by both the interaction term and the relevant model parameters.

gr-qc↗

arXiv2024

Interacting Models of Dark Energy and Dark Matter in Einstein scalar Gauss Bonnet Gravity

We study the dynamics of the interacting models between the Gauss-Bonnet (GB) coupled scalar field and the dark matter fluid in a homogeneous and isotropic background. A key feature of GB coupling models is the varying speed of gravitational waves (GWs). We utilize recent constraints on the GW speed and conduct our analysis in two primary scenarios: model-dependent and model-independent. In the model-dependent scenario, where determining the GW speed requires a specific GB coupling functional form, we choose an exponential GB coupling. We adopt a dynamical system analysis to obtain the necessary constraints on the model parameters that describe different phases of the universe and produce a stable late-time accelerating solution following the GW constraint, and find that to satisfy all these constraints, fine-tuning of the free parameters involved in the models is often needed. In the model-independent scenario, the GW speed is fixed to one, and we construct the autonomous system to identify the late-time stable accelerating critical points. Furthermore, we adopt a Bayesian inference method using late-time observational data sets, including 31 data points from cosmic chronometer data (Hubble data) and 1701 data points from Pantheon+ and find that all the observational constraints can be satisfied without fine-tuning. In addition, we also utilize simulated binned Roman and LSST data to study the evolution of the universe in the model-independent scenario. We find that the model shows significant deviation at higher redshifts from $Λ$CDM and fits the current data much better than $Λ$CDM within the error bars.

gr-qc↗

arXiv2024

Optimistic Verifiable Training by Controlling Hardware Nondeterminism

The increasing compute demands of AI systems have led to the emergence of services that train models on behalf of clients lacking necessary resources. However, ensuring correctness of training and guarding against potential training-time attacks, such as data poisoning and backdoors, poses challenges. Existing works on verifiable training largely fall into two classes: proof-based systems, which are difficult to scale, and ``optimistic'' methods that consider a third-party auditor who can replicate the training process and contest the trainer. A key challenge with the latter is that nondeterminism between GPU types during training prevents exact replication of the training process, resulting in schemes that are non-robust. We propose a method that combines training in a higher precision than the target, rounding after intermediate computations, and sharing rounding decisions based on an adaptive thresholding procedure, to successfully control for nondeterminism. Across three different NVIDIA GPUs (A40, Titan XP, RTX 2080 Ti), we achieve exact training replication at FP32 precision for both full-training and fine-tuning of ResNet-50 (23M) and GPT-2 (117M) models. Our verifiable training scheme significantly decreases the storage and time costs compared to proof-based systems, and is publicly released at https://github.com/meghabyte/verifiable-training.

cs.CR↗

arXiv2024

Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT

Retrieval pipelines-an integral component of many machine learning systems-perform poorly in domains where documents are long (e.g., 10K tokens or more) and where identifying the relevant document requires synthesizing information across the entire text. Developing long-context retrieval encoders suitable for these domains raises three challenges: (1) how to evaluate long-context retrieval performance, (2) how to pretrain a base language model to represent both short contexts (corresponding to queries) and long contexts (corresponding to documents), and (3) how to fine-tune this model for retrieval under the batch size limitations imposed by GPU memory constraints. To address these challenges, we first introduce LoCoV1, a novel 12 task benchmark constructed to measure long-context retrieval where chunking is not possible or not effective. We next present the M2-BERT retrieval encoder, an 80M parameter state-space encoder model built from the Monarch Mixer architecture, capable of scaling to documents up to 32K tokens long. We describe a pretraining data mixture which allows this encoder to process both short and long context sequences, and a finetuning approach that adapts this base model to retrieval with only single-sample batches. Finally, we validate the M2-BERT retrieval encoder on LoCoV1, finding that it outperforms competitive Transformer-based models by at least 23.3 points, despite containing upwards of 90x fewer parameters.

cs.IR↗

arXiv2024

ThunderKittens: Simple, Fast, and Adorable AI Kernels

The challenge of mapping AI architectures to GPU hardware is creating a critical bottleneck in AI progress. Despite substantial efforts, hand-written custom kernels fail to meet their theoretical performance thresholds, even on well-established operations like linear attention. The diverse hardware capabilities of GPUs might suggest that we need a wide variety of techniques to achieve high performance. However, our work explores whether a small number of key abstractions can drastically simplify the process. We present ThunderKittens (TK), a framework for writing performant AI kernels while remaining easy to use and maintain. Our abstractions map to the three levels of the GPU hierarchy: (1) at the warp-level, we provide 16x16 matrix tiles as basic data structures and PyTorch-like parallel compute operations over tiles, (2) at the thread-block level, we provide a template for overlapping asynchronous operations across parallel warps, and (3) at the grid-level, we provide support to help hide the block launch and tear-down, and memory costs. We show the value of TK by providing kernels that match or outperform prior kernels for a range of AI operations. We match CuBLAS and FlashAttention-3 on GEMM and attention inference performance and outperform the strongest baselines by $10-40\%$ on attention backwards, $8\times$ on state space models, and $14\times$ on linear attention.

cs.LG↗

arXiv2024

p-(001)NiO/n-(0001)ZnO heterostructures grown by pulsed laser deposition technique

NiO/ZnO heterostructures are grown on c-sapphire substrates using pulsed laser deposition (PLD) technique. X-ray diffraction study shows that the ZnO layer epitaxially grows along [0001]-direction on (0001)sapphire surface as expected. While, the epitaxial NiO film is found to be deposited along [001]-direction on the (0001)ZnO surface. Moreover, the presence of three (001)NiO domains laterally rotated by 30° with respect to each other, has also been observed in our NiO films. The study reveals the continuous nature of the NiO film, which also possesses a very smooth surface morphology. In a sharp contrast, ZnO films are found to grow along [0001]-direction when deposited on (111)NiO layers. These films also show columnar morphology. (001)NiO/(0001)ZnO layers exhibit the rectifying current-voltage characteristics that suggests the existence of p-n junction in these devices. However, the behavior could not be observed in (0001)ZnO/(111)NiO heterojunctions. The reason could be the columnar morphology of the ZnO layer. Such a morphology can facilitate the propagation of the metal ions from the contact pads to the underlying NiO layer and suppress the p-n junction effect.

cond-mat.mtrl-sci↗

arXiv2024

Reconstruction of the singularity-free $f(\mathcal{R})$ gravity via Raychaudhuri equations

We study the bounce cosmology to construct a singularity-free $f(\mathcal{R})$ model using the reconstruction technique. The formulation of the $f(\mathcal{R})$ model is based on the Raychaudhari equation, a key element employed in reconstructed models to eliminate singularities. We explore the feasibility of obtaining stable gravitational Lagrangians, adhering to the conditions $f_{\mathcal{R}}>0$ and $f_{\mathcal{R}\mathcal{R}}>0$. Consequently, both models demonstrate stability, effectively avoiding the Dolgov-Kawasaki instability. Our assessment extends to testing the reconstructed model using energy conditions and the effective equation-of-state (EoS). Our findings indicate that the reconstructed super-bounce model facilitates the examination of a singularity-free accelerating universe for both phantom and non-phantom phases. However, in the case of the reconstructed oscillatory bounce model, two scenarios are considered with $ω=-1/3$ and $ω=-2/3$. While the model proves suitable for studying a singular-free accelerating universe in the $ω=-1/3$ case, it fails to demonstrate such behavior under energy conditions for the $ω=-2/3$ scenario. The reconstructed models accommodate early-time bouncing behavior and late-

gr-qc↗

arXiv2024

Just read twice: closing the recall gap for recurrent language models

Recurrent large language models that compete with Transformers in language modeling perplexity are emerging at a rapid rate (e.g., Mamba, RWKV). Excitingly, these architectures use a constant amount of memory during inference. However, due to the limited memory, recurrent LMs cannot recall and use all the information in long contexts leading to brittle in-context learning (ICL) quality. A key challenge for efficient LMs is selecting what information to store versus discard. In this work, we observe the order in which information is shown to the LM impacts the selection difficulty. To formalize this, we show that the hardness of information recall reduces to the hardness of a problem called set disjointness (SD), a quintessential problem in communication complexity that requires a streaming algorithm (e.g., recurrent model) to decide whether inputted sets are disjoint. We empirically and theoretically show that the recurrent memory required to solve SD changes with set order, i.e., whether the smaller set appears first in-context. Our analysis suggests, to mitigate the reliance on data order, we can put information in the right order in-context or process prompts non-causally. Towards that end, we propose: (1) JRT-Prompt, where context gets repeated multiple times in the prompt, effectively showing the model all data orders. This gives $11.0 \pm 1.3$ points of improvement, averaged across $16$ recurrent LMs and the $6$ ICL tasks, with $11.9\times$ higher throughput than FlashAttention-2 for generation prefill (length $32$k, batch size $16$, NVidia H100). We then propose (2) JRT-RNN, which uses non-causal prefix-linear-attention to process prompts and provides $99\%$ of Transformer quality at $360$M params., $30$B tokens and $96\%$ at $1.3$B params., $50$B tokens on average across the tasks, with $19.2\times$ higher throughput for prefill than FA2.

cs.CL↗

arXiv2024

Energy conditions in the $f(R,L,T)$ theory of gravity

We construct the energy conditions for the recently proposed $f(R,L,T)$ gravity theory, for which $f$ is a generic function of the Ricci scalar $R$, matter lagrangian density $L$ and trace of the energy-momentum tensor $T$. We analyse two different forms for the $f(R,L,T)$ function within the framework of the Friedmann-Lemâitre-Robertson-Walker universe. We constrain the model parameters from the energy conditions. This approach allows us to assess the feasibility of specific forms of the $f(R,L,T)$ gravity.

gr-qc↗

arXiv2024

Revisiting kink-like parametrization and constraints using OHD/Pantheon+/BAO samples

We reexamine the kink-like parameterization of the deceleration parameter to derive constraints on the transition redshift from cosmic deceleration to acceleration. This is achieved using observational Hubble data, Type Ia Supernovae Pantheon+ samples and Baryon acoustic oscillations. In this parametrization, the value of the initial $q$ parameter is $q_{i}$, the final value is $q_f$, the present value is denoted by $q_{0}$, and the transition duration is given by $α$. We perform our calculations using the Monte Carlo Markov Chain method, utilizing the emcee package. Under the assumption of a flat geometry, we constrain the range of possible values for three scenarios: when $q_{f}$ is unrestricted, when $q_{f}$ is equal to $-1$, and when $α$ is $1/3$. This is done assuming that $q_{i}=1/2$. Here, we achieve that the $SN$ data fixes the free parameters tightly as in the flat $Λ$CDM for unrestricted $q_{f}$. In addition, if we fix $q_{f}=-1$, the model behaves well as the $Λ$CDM for the combined dataset. We also acquire the current value of the deceleration parameter, which is consistent with the latest results that assume the $Λ$CDM model. Furthermore, we observe a deviation from the standard $Λ$CDM model in the current model based on the evolution of $j(z)$, and it is evident that the universe transitions from deceleration to acceleration and will eventually reach the $Λ$CDM model in the near future.

astro-ph.CO↗

arXiv2024

RELIC: Investigating Large Language Model Responses using Self-Consistency

Large Language Models (LLMs) are notorious for blending fact with fiction and generating non-factual content, known as hallucinations. To address this challenge, we propose an interactive system that helps users gain insight into the reliability of the generated text. Our approach is based on the idea that the self-consistency of multiple samples generated by the same LLM relates to its confidence in individual claims in the generated texts. Using this idea, we design RELIC, an interactive system that enables users to investigate and verify semantic-level variations in multiple long-form responses. This allows users to recognize potentially inaccurate information in the generated text and make necessary corrections. From a user study with ten participants, we demonstrate that our approach helps users better verify the reliability of the generated text. We further summarize the design implications and lessons learned from this research for future studies of reliable human-LLM interactions.

cs.HC↗