Research / Work

I list some arcs in my history as a researcher, including some notable papers from each era. I also add some commentary about each research arc in my life.

Gemini era

Some of the notable blog/launches that I was involved in (apart from conventional Gemini launches).

  1. Advanced version of Gemini with Deep Think officially achieves gold-medal standard at the International Mathematical Olympiad

  2. Gemini achieves gold-medal level at the International Collegiate Programming World Finals

  3. Accelerating Mathematical and Scientific Discovery with Gemini Deep Think

Pioneering LLM work (pre-2023)

This section lists pioneering LLM contributions and some of my favorite works from this era/arc. Some notable memories from this era is co-leading the modeling for PaLM-2, open sourcing UL2 20B, Flan-2 series and the Emergent abilities paper with my good friend Jason Wei. This was really a magical and really fun era in Google Brain.

  1. PaLM 2 Technical Report

  2. Transcending Scaling Laws with 0.1% Extra Compute

  3. UL2: Unifying Language Learning Paradigms

  4. Scaling Instruction-Finetuned Language Models (Flan-2 paper)

  5. PaLM: Scaling Language Models with Pathways

  6. Emergent Abilities of Large Language Models

  7. PaLI-X: On Scaling a Multilingual Vision and Language Model

  8. Scaling Vision Transformers to 22 Billion Parameters

Generative Retrieval / DSI

I accidentally invented a new paradigm that is widely used today in industry as a side project. This is known as generative retrieval and we do this by directly predicting document identifiers with LLMs. (it’s super cool ngl). The pioneering work DSI was a naive proof of concept that something like this even works. Today, the work has evolved in many ways by the community/industry, i.e., by using complex semantic ids. I don’t work on this anymore but my last involvement was a collaboration in the context of recommender systems (Youtube). Last I heard this is pretty much used broadly in industry today.

  1. Transformer Memory as a Differentiable Search Index

  2. DSI++: Updating Transformer Memory with New Documents

  3. Recommender Systems with Generative Retrieval

Pre-2023 architecture / scaling research

I had a lot of fun at Google working on a bunch of architecture research, scaling research and I think at one point of time I dedicated like 2 whole years just playing around with the Transformer architecture and making it more efficient. Synthesizer was one of my first works since joining Google in 2019. Over this time, I spent a lot of time understanding attention mechanisms.

  1. Synthesizer: Rethinking Self-Attention in Transformer Models

  2. Long Range Arena: A Benchmark for Efficient Transformers

  3. Charformer: Fast Character Transformers via Gradient-based Subword Tokenization

  4. Do Transformer Modifications Transfer Across Implementations and Applications?

  5. Efficient Transformers: A Survey

  6. Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers

  7. Scaling Laws vs Model Architectures: How does Inductive Bias Influence Scaling?