One Set of Weights, Fourteen Medical Tasks: Inside Google's Multimodal Med-PaLM M

Host Nathan Labenz opens by noting how fast Google's medical AI work has moved. Med-PaLM reached passing level on the licensing exam, Med-PaLM 2 reached expert-level question answering, and now Med-PaLM M adds imaging and genomics. Natarajan says the pace reflects building on the shoulders of giants: the Transformer in 2017, the PaLM line of models, and the infrastructure for training at scale.

He adds that his team has had years of successes and failures in medical AI. It also includes an interdisciplinary group that understands medicine well enough to frame the right problems. In his words it feels like "surfing the wave", being in the right place at the right time.

The multimodal push follows from how clinicians work. Doctors read clinical text, imaging, lab results, vital signs and sensor data together, and any AI that wants to understand a patient needs the same context. The team chose the hardest version of the problem: a generalist biomedical agent that uses one set of parameters to encode many modalities and solve many tasks. DeepMind's Gato and Google's PaLM-E showed the idea in other domains, but medicine's data poses its own challenges.

Tu lists the data they included: dermatology, pathology, chest X-ray, mammography and genomics, alongside clinical text. Every non-text modality is converted into an image that the model's vision encoder can digest.

Labenz asks how genomic data could be shown as an image. Tu explains that collaborators from Google's genomics team had built DeepVariant, a state-of-the-art variant-calling model. It converts sequencing data into a 3D tensor that an Inception-style vision network can read, and the Med-PaLM team reshaped that tensor to fit its vision transformer encoder.

Natarajan gives the history. Around 2018-2019 the Google Health genomics group cast variant calling as image classification because computer vision was the strength of deep learning then. He says the system won the Precision FDA challenge and was a key component in a record-setting sequencing effort from Stanford researchers.

Med-PaLM M does not reuse any DeepVariant weights. Everything, including the genomics capability, was trained from scratch inside a generative setup where classification outputs are simply language such as a suspected variant or a dermatology diagnosis. Language becomes the common grounding. The only inherited weights came from PaLM-E, a vision-language model originally designed for robotics.

PaLM-E had no medical fine-tuning at all. The researchers used it out of the box as a baseline and it performed poorly. They read this as evidence that medical data and specialization are essential for such systems.

Natarajan says earlier generalist systems made the project look doable, but the devil is in the details. Specialist tools, such as a mammography screening system already in clinical trials, set a high bar. The team had two questions: whether a generalist can match narrow models that took years of tuning, and whether training on many modalities produces positive transfer and combinatorial generalization.

Tu notes they did not fine-tune with multiple images in one example, mainly for compute reasons, even though comparing longitudinal change matters in practice. Instead they used text tags standing in for images in few-shot prompts. The paper shows the model still generalizes to cases with several chest X-ray views or prior images.

Natarajan generalizes the point. Building large multimodal models is about assembling a reasonable bag of tricks and combining them so they do not eat too much compute. There is no secret new knowledge, and he says the team behind GPT-4 would likely say the same. Labenz adds that this accumulated know-how may be a better place to look for moats than anything else.

The paper has roughly 50 authors. Natarajan says there was no top-down mandate at Google: researchers saw the problem was tractable and came together. He insists that program managers, product managers and the people handling data access, licensing and legal review deserve authorship, since doing medical AI the right way requires them.

The team also leaned on the authors of PaLM and PaLM-E for advice and on in-house clinicians for fast feedback. Radiologist Chuck helped craft the detailed chest X-ray prompt, which mirrors how a radiologist inspects regions and summarizes findings. Only about eight or ten people were deep ML researchers, and roles were fluid across infrastructure and modeling.

Asked about Pathways, Natarajan describes it as a large-scale orchestration system for training multitask, multimodal models across thousands of accelerators, typically TPUs. JAX and the XLA compiler map code onto the hardware, and Pax and Flax provide model and data parallelism. Most researchers on the project could stay at the application level, and could ask a Pathways expert when something odd appeared, such as loss spikes or non-determinism.

Tu says the team first trained on each task separately, since they did not know whether mixing tasks would hold performance. That meant experiments on prompt design for each dataset. Chest X-ray report generation used detailed instructions from Chuck, while simpler tasks such as mammography classification did not.

The hard part was deciding task proportions. Datasets vary widely in size, and if the model sees mostly chest X-rays it might become biased toward grayscale images and hurt genomics. Even on mature infrastructure, combining everything produced gradient instability that required some clever stabilization tricks.

Natarajan says the core code is not large. The bulk of the work was experiments that transfer to the larger models, plus the human-evaluation infrastructure, interfaces and clinician studies. Tu says he ran around 300 experiments over four months.

Labenz points to the smooth loss curves in the GPT-4 report and asks whether they hold here. Natarajan says they hold for text-only models. In multimodal systems, where separately optimized components are fused, the predictability fades. In PaLM-E and Med-PaLM M the language side was scaled far beyond the vision side.

Tu describes the three parts: vision encoder, language model, and an alignment layer that projects image tokens into language space. The vision transformer was pretrained on natural images, whose distribution differs greatly from medical ones. To study scaling properly, the team would need to adapt the vision encoder to medicine first, which requires much more medical data.

This explains why the 562B model barely beat the 84B model on some tasks. On classification, for example saying a skin image shows eczema, the output space is tiny and the difficulty lies in seeing patches and signs. Image size is another constraint: all inputs share one shape, so chest X-ray resolution is lower than in specialist report-generation models.

Natarajan says these knobs can be tuned in the next iteration. Examples are adaptive image resolution, varying the number of tokens per modality, and more capacity and data. That is why the team calls Med-PaLM M a proof of concept.

The dataset totals roughly one million samples. Natarajan estimates that comes to tens of millions of text tokens, with each image becoming 256 tokens. That is small next to the trillion-ish tokens used to pretrain PaLM. Compute was about 1 to 10 percent of base training, on TPU v4 chips.

Iteration speed depends on chip count. Tu says the smallest 12B model finishes training in a couple of hours, but with a tenth of the chips that becomes two weeks. The 562B model takes weeks. Labenz and Natarajan also discuss data and model parallelism, which together form a 3D mesh, and note that scale partly offsets hardware gains.

Labenz asks about adapters and bridges. Tu says they tried freezing the language model, freezing the vision encoder, and tuning only the alignment layer. Frozen setups did work but converged more slowly. For Med-PaLM 1 the team used prompt tuning, while Med-PaLM 2 and Med-PaLM M use full end-to-end fine-tuning, which gave the best results given a reasonable amount of data.

Natarajan adds that the same held when comparing against LoRA, and that LoRA's compute savings may not hold up once convergence speed is counted. The next step is a more modular design with small modality-specific encoders and a shared, newer language model.

Future versions aim to widen the task bench to more imagery and data such as proteomics. The team says approaches range from tool use to adapters to end-to-end fine-tuning, and the question they ask first is which problem they are solving.

Asked whether the aim is simply an AI doctor, Natarajan says that is no longer science fiction, though the field is early in validation and real-world trials. He wants systems to be effective teammates that remove administrative burden from clinicians.

He also sees models doing things clinicians are not trained for, such as interpreting genomic data, whose measurement cost is falling faster than Moore's law. Training on diverse biomedical data could also reveal new biomarkers, gene variants or therapies, which makes the system something like an AI research assistant.

Natarajan says he and Tu debated the term emergent, given the wider argument over the "Mirage" paper. For him it simply means something unpredicted. The model was trained on chest X-rays with 14 common labels such as enlarged heart, and tuberculosis was not among them.

When shown an X-ray from a different center and asked to describe findings, the model characterized the cavity lesion typical of tuberculosis and located it correctly. It also missed some lesions. Natarajan adds that the model needed zero tuberculosis examples because a language description was enough, and the 8B model could not do it while the 84B and 562B models could.

Labenz describes the projected image as 256 tokens that, in effect, whisper to the language model what is seen where. Tu is less surprised, noting the image is still a chest X-ray and thus in distribution. He would be more excited to show a fungal infection unlike anything in training. Natarajan characterizes the current result as interpolation, with early extrapolation evidence promising.

Asked about system-level clinical trials, Natarajan says regulation is the barrier. Software as a medical device is regulated on the assumption that each system does a few tasks with known limits, which lets you put numbers on effectiveness. A generally capable system allows endless interactions, so the old framework cannot be retrofitted. He says talks with the FDA and the EU are ongoing, and a controlled trial of health outcomes could be a reasonable path.

For now, Med-PaLM 2 is rolled out to trusted testers for feedback. On guardrails, Tu says the team tested with adversarial questions and must avoid biased answers and amplified health inequality. Natarajan argues that medicine's constrained domain makes safety and alignment more tractable than for an open chatbot.

He says it is acceptable for society to deliberate rather than rush. Asked about the timeline for democratizing care, he recalls growing up where seeing a doctor could mean walking 30 or 40 miles and giving up a day's wages. He sees access, cost and quality as the industry's biggest problems, and all are addressable by AI. Tu, from neuroscience, hopes to use the language model's knowledge of PubMed to find biomarkers for Alzheimer's or Parkinson's.