← Search

Azalia Mirhoseini

15 accepted papers

2026

Agent JIT Compilation for Latency-Optimizing Computer-Use Agent Planning and Scheduling

ICML 2026poster

Computer-use agents (CUA) automate tasks specified with natural language such as "order the cheapest item from Taco Bell", by generating sequences of calls to tools such as click, type, and scroll on a browser. Current CUA implementations follow a sequential fetch-screenshot-execute loop where each …

Cited by 0SourceScholar
2026

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

ICLR 2026poster

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-10M tokens, this se…

Cited by 0SourcecodeScholar
2025

An Architecture Search Framework for Inference-Time Techniques

ICML 2025poster

Inference-time techniques, such as repeated sampling or iterative revisions, are emerging as powerful ways to enhance large-language models (LLMs) at test time. However, best practices for developing systems that combine these techniques remain underdeveloped due to our limited understanding of the…

Cited by 0SourcePDFScholar
2025

Exploring Diffusion Transformer Designs via Grafting

NeurIPS 2025oral

Designing model architectures requires decisions such as selecting operators (e.g., attention, convolution) and configurations (e.g., depth, width). However, evaluating the impact of these decisions on model quality requires costly pretraining, limiting architectural investigation. Inspired by how n…

Cited by 0SourcecodeScholar
2025

How Do Large Language Monkeys Get Their Power (Laws)?

ICML 2025oral

Recent research across mathematical problem solving, proof assistant programming and multimodal jailbreaking documents a striking finding: when (multimodal) language model tackle a suite of tasks with multiple attempts per task -- succeeding if any attempt is correct -- then the negative log of the…

Cited by 0SourcePDFScholar
2025

KernelBench: Can LLMs Write Efficient GPU Kernels?

ICML 2025poster

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…

Cited by 1SourcePDFScholar
2025

RoboMonkey: Scaling Test-Time Sampling and Verification for Vision-Language-Action Models

CoRL 2025poster

Vision-Language-Action (VLA) models, pre-trained on large-scale imitation learning datasets, have demonstrated remarkable capabilities in visuomotor control. However, these models exhibit diverse failure modes in unstructured real-world environments, limiting the widespread adoption of VLAs in robot…

Cited by 0SourceScholar
2025

SPRINT: Enabling Interleaved Planning and Parallelized Execution in Reasoning Models

NeurIPS 2025poster

Large reasoning models (LRMs) excel at complex reasoning tasks but typically generate lengthy sequential chains-of-thought, resulting in long inference times before arriving at the final answer. To address this challenge, we introduce SPRINT, a novel post-training and inference-time framework design…

Cited by 0SourceScholar
2025

Weaver: Shrinking the Generation-Verification Gap by Scaling Compute for Verification

NeurIPS 2025poster

Verifiers can improve language model (LM) capabilities by providing feedback or selecting the best response from a pool of generated candidates. Currently, high-quality verifiers are either unscalable (e.g., humans) or limited in utility (e.g., tools like Lean for formal proofs). While LM judges and…

Cited by 0SourceScholar
2023

Embroid: Unsupervised Prediction Smoothing Can Improve Few-Shot Classification

NeurIPS 2023poster

Recent work has shown that language models' (LMs) prompt-based learning capabilities make them well suited for automating data labeling in domains where manual annotation is expensive. The challenge is that while writing an initial prompt is cheap, improving a prompt is costly---practitioners often…

Cited by 6SourcePDFScholar
2021

Representing Long-Range Context for Graph Neural Networks with Global Attention

NeurIPS 2021poster

Graph neural networks are powerful architectures for structured datasets. However, current methods struggle to represent long-range dependencies. Scaling the depth or width of GNNs is insufficient to broaden receptive fields as larger GNNs encounter optimization instabilities such as vanishing gradi…

2020

Transferable Graph Optimizers for ML Compilers

NeurIPS 2020oral

Most compilers for machine learning (ML) frameworks need to solve many correlated optimization problems to generate efficient machine code. Current ML compilers rely on heuristics based algorithms to solve these optimization problems one at a time. However, this approach is not only hard to maintain…

Cited by 63SourcePDFScholar
2019

Deep Mixture of Experts via Shallow Embedding

UAI 2019poster

Larger networks generally have greater representational power at the cost of increased computational complexity. Sparsifying such networks has been an active area of research but has been generally limited to static regularization or dynamic approaches using reinforcement learning. We explore a mixt…

Cited by 136SourcePDFScholar
2018

A Hierarchical Model for Device Placement

ICLR 2018poster

We introduce a hierarchical model for efficient placement of computational graphs onto hardware devices, especially in heterogeneous environments with a mixture of CPUs, GPUs, and other computational devices. Our method learns to assign graph operations to groups and to allocate those groups to avai…

Cited by 210SourcePDFScholar
2017

Device Placement Optimization with Reinforcement Learning

ICML 2017poster

The past few years have witnessed a growth in size and computational requirements for training and inference with neural networks. Currently, a common approach to address these requirements is to use a heterogeneous distributed environment with a mixture of hardware devices such as CPUs and GPUs. Im…

Cited by 556SourcePDFScholar