I received my Ph.D. in EECS from MIT CSAIL, advised by Prof. Antonio Torralba. I study how to make model training, from pretraining through post-training, more efficient and stable: hardware-aware architecture design, attention mechanisms, reinforcement learning, and optimization methods. During my Ph.D., I built and trained models on vision, language, 3D, tactile, physics-simulation, and robotics data. This breadth enables me to train generalist foundation models that can acquire and combine diverse skills.
More recently, I have become interested in the physical and engineering constraints underlying computation, particularly how models and hardware can be designed together.
EDA, silicon, power electronics, spintronic devices, quantum optics, and lithography. I want to bring AI to these fields, including the hardware that AI itself is trained on and to the development of broad physical sciences.
AC: Model Architecture Compiler under Hardware ConstraintsOpen source · 2026AC sits between model scaling laws and roofline throughput estimation. Given a hardware platform, parameter count, training tokens, serving workload, and an optional base model, it searches for multi-objective Pareto-improving model architectural designs.
pip install archc
A compiler for model architecture design under hardware constraints. Given a hardware target (H100, B200, GB200 NVL72, TPU v5e/v5p, Trainium 2/3), parameter count, training tokens, a serving workload, and optionally a base model, AC returns Pareto-improving architectures or architecture deltas across loss, training cost, memory, and serving latency, and names the binding hardware constraint (HBM bandwidth, KV footprint, tile fit, all-reduce or all-to-all pressure) with the price of relaxing it. It connects scaling-law quality estimates to roofline compute, memory, and communication models, so a choice like MLA vs. GQA becomes a constrained multi-objective decision rather than folklore. In a pre-registered study against five published architecture decisions (DeepSeek-V3 MLA and first-K-dense, Llama-3 GQA-8, Mistral SWA, wide expert parallelism on NVL72), AC scored 4 hits and 1 partial; absolute loss and MFU remain uncalibrated priors, and every miss is published.
Efficient RL for Post-Training LLM and Autoregressive Generative ModelsICML 2026 · ReportsThree ways to make RL post-training cheaper and more stable: forward-pass-only evolution strategies (ESES), training with the pretraining loss (AWM), and densified reward over autoregressive rollouts (VARL).
Evolution strategies estimate the policy gradient from forward passes alone, which buys rollout diversity and training stability at the cost of sample efficiency. ESES makes ES cheaper by sampling low-rank (LoRA) perturbations that can be batch-served, and maps how population size, noise scale, and perturbation rank trade off in the small-population regime (Qwen2.5-3B). Because ES never backpropagates, perturbations can stay at the base model's precision, so weights can be quantized (INT8, FP8, INT4) without the training–inference mismatch that destabilizes PPO and GRPO.
Shows that DDPO is an implicit form of score/flow matching with noisy targets, which raises variance and slows convergence. AWM instead trains with the pretraining score/flow-matching loss, reweighting each sample by its advantage. Up to 34× faster than Flow-GRPO on SD3.5 Medium and FLUX.
In autoregressive block-diffusion models, one rollout costs autoregressive steps × denoising steps. VARL rolls out many short block sequences, scores them as a group, then extends a random subset and rescores, densifying reward along the sequence for a fraction of the cost of full rollouts.
Projective Attention: A Geometric Design Space for AttentionBlog series · 5 posts · 2026What if softmax is just one choice? Attention read as a generalized vector activation, with its own design choices: activation function, gauge control, and decoding.
Main finding. Under random, shuffled, and compute-matched controls, the gains that held up came from separating key direction from key norm, and every variant made stable drifted back toward standard attention.
Reads each key as a hyperplane normal, so a head becomes a one-hidden-layer network whose hidden units come from the context: logits are signed incidences, the attention row is an activation code, and values decode it. This places 18 published variants (ReLU, sigmoid, signed, sparse, routed, QK-norm, XSA, …) on three axes (activation, hyperplane gauge, decoding) and generates new ones (Hyperplane-Projection Attention, Projection Feature, Projective Clipping, two-sided decoding), each tested against random, shuffled, rotated, and compute-matched versions of itself.
Pretraining & architecture. Softmax is not sacred, but normalized activation is: every replacement had to rebuild part of what SDPA already provides (row normalization, bounded output, scale control) before it trained stably. That gives a checklist for judging an attention proposal before spending compute on it, and points to extensions that open one degree of freedom of SDPA at a time instead of replacing it.
Training stability. The most reliable gains came from gauge control. A raw dot product ties a key’s direction to its logit scale, the same entanglement QK-norm removes to stop attention-logit growth. The open proposal is to keep the norm as a separate confidence or temperature signal rather than discard it.
Interpretability. Hyperplane-Projection Attention separates addressing from content (the QK and OV roles) by subtracting the key-aligned “address shadow” from what a head reads out. With normalized keys it improved synthetic retrieval in 16 of 16 settings, and head-split variants let some heads keep the address while others drop it.
Research methodology. Controls overturned several early wins: the first seed was misleadingly optimistic, and random or shuffled branches matched the Projection Feature and Projective Clipping gains. A branch-budget measure (branch norm relative to the attention output) separates branches that are decorative from ones that overwhelm the readout.
Scope: small models (2-layer, synthetic retrieval and character-level LM). The findings are about mechanism and method; quality at scale is untested.
Natively Multimodal Token Embeddings and Multimodal Generative ModelsICCV 2025 · ICLR 2024One embedding shared across input modalities, and a video generator conditioned on language, motion, and appearance at once.
An autodecoder learns a single cross-modal latent manifold, so full-hand tactile pressure maps and manipulation video map to one shared embedding: either modality can be predicted from the other, and the latent space exposes structure common to both. Comes with a new paired visual–tactile dataset of human manipulation. The recipe is modality-agnostic, which is what lets the learned embedding serve as an extended token embedding for a foundation model.
Sole student author; trained from scratch, end to end: architecture, data, training, and evaluation for a video generator that jointly conditions on language, bounding-box trajectories, and reference frames. 40% FVD improvement over text-conditioned baselines.
Built a unified physics simulation framework for soft, articulated, and rigid body dynamics.
Designed anisotropic Young's modulus learning across physics regimes.
51% error reduction over team baseline.
