Doctoral student (Ph.D.)
Berkeley Artificial Intelligence Research (BAIR)
Department of Electrical Engineering and Computer Science (EECS)
UC Berkeley
Email: ajayj at berkeley dot edu
Biography
From SIMD to 
I'm a PhD student at UC Berkeley, affiliated with Berkeley Artificial Intelligence Research.
I work on machine learning and performance optimization, with a special interest in generative modeling.
At Berkeley, I am advised by Pieter Abbeel and Joseph Gonzalez and am supported by an NSF GRFP Fellowship.
In Spring 2019, I graduated from MIT with an S.B. in Computer Science, where I was fortunate to be advised by Saman Amarasinghe and affiliated with the MIT COMMIT lab. I've also had the good fortune of working with Raquel Urtasun at Uber ATG.
If you go to MIT, you've likely received an email from me at some point about a machine learning reading group or talk -- from 2017-2018,
I was the co-president of MIT MIC, and volunteered my time by ordering a tremendous amount of pizza to
entice students to read papers together. Now, I sit on the board of the non-profit
Machine Intelligence Community.
In the past, I organized HackMIT and Blueprint.
News
New! January 2020: Paper on memory-efficient DL accepted to MLSys 2020
September 2019: Paper on behavior forecasting accepted to CoRL 2019
June 2019: Presented a 10 minute talk at ISCA 2019 Workshop on Machine Learning for Systems
June 2019: Started my PhD at UC Berkeley
June 2018: Headed to Toronto for a research internship at Uber Advanced Technologies Group
June 2017: Joining Facebook's Applied Machine Learning team for the summer
Publications * indicates equal contribution
Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
Paras Jain*, Ajay Jain*, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, Ion Stoica, Joseph E. Gonzalez
Proceedings of the 3rd Conference on Machine Learning and Systems (MLSys), Mar 2020
Use up to 5x less memory when training DNNs by recomputing activations
Discrete Residual Flow for Probabilistic Pedestrian Behavior Prediction
Ajay Jain*, Sergio Casas Romero*, Renjie Liao*, Yuwen Xiong*, Song Feng, Sean Segal, Raquel Urtasun
Proceedings of the 3rd Conference on Robot Learning (CoRL), Oct 2019
Multimodal, long-range behavior forecasts by predicting state marginals
Revec: Program Rejuvenation through Revectorization
Charith Mendis*, Ajay Jain*, Paras Jain, Saman Amarasinghe
Proceedings of the 28th International Conference on Compiler Construction (CC), Feb 2019
Achieve performance portability for hand-vectorized programs, with up to 1.88x speedup
Autonomy for Surface Ship Interception
C. Mirabito, D.N. Subramani, T. Lolla, J. P.J. Haley, A. Jain, P.F.J. Lermusiaux, C. Li, D. Yue, Y. Liu, F. Hover, N. Pulsone, J. Edwards, K. Railey, and G. Shaw
Proceedings of the 60th OCEANS Conference, Oceans MTS/IEEE Aberdeen, June 2017
Time-optimal path planning for underwater robot
Preprints
The Case for GPU Multitenancy
Paras Jain, Xiangxi Mo, Ajay Jain, Alexey Tumanov, Joseph E. Gonzalez, Ion Stoica
arXiv:1910.02653 preprint, Jan 2019
Workshop papers
Learning Automatic Schedulers with Projective Reparameterization
Ajay Jain, Saman Amarasinghe
The 46th International Symposium on Computer Architecture (ISCA)
Workshop on Machine Learning for Systems, Jun 2019, Talk
Supervised learning of schedulers, with correctness constraints
Using effective dimension to analyze feature transformations in deep neural networks
Kavya Ravichandran, Ajay Jain, Alexander Rakhlin
The 36th International Conference on Machine Learning (ICML)
Workshop on Identifying and Understanding Deep Learning Phenomena, Jun 2019
Dynamic Space-Time scheduling for GPU inference
Paras Jain, Xiangxi Mo, Ajay Jain, Harikaran Subbaraj, Rehan Sohail Durrani, Alexey Tumanov, Joseph E. Gonzalez, and Ion Stoica
The 32nd Annual Conference on Neural Information Processing Systems (NeurIPS)
Workshop on Systems for Machine Learning, Dec 2018
Demonstrate 2.5x-4.9x speedups for deep learning inference workloads via GPU multitenancy