Robust and data-efficient generalization of self-supervised machine learning for diagnostic imaging.

Journal: Nature biomedical engineering
Published Date:

Abstract

Machine-learning models for medical tasks can match or surpass the performance of clinical experts. However, in settings differing from those of the training dataset, the performance of a model can deteriorate substantially. Here we report a representation-learning strategy for machine-learning models applied to medical-imaging tasks that mitigates such 'out of distribution' performance problem and that improves model robustness and training efficiency. The strategy, which we named REMEDIS (for 'Robust and Efficient Medical Imaging with Self-supervision'), combines large-scale supervised transfer learning on natural images and intermediate contrastive self-supervised learning on medical images and requires minimal task-specific customization. We show the utility of REMEDIS in a range of diagnostic-imaging tasks covering six imaging domains and 15 test datasets, and by simulating three realistic out-of-distribution scenarios. REMEDIS improved in-distribution diagnostic accuracies up to 11.5% with respect to strong supervised baseline models, and in out-of-distribution settings required only 1-33% of the data for retraining to match the performance of supervised models retrained using all available data. REMEDIS may accelerate the development lifecycle of machine-learning models for medical imaging.

Authors

  • Shekoofeh Azizi
  • Laura Culp
    Google Research, Mountain View, CA, USA.
  • Jan Freyberg
    Google Research, Mountain View, CA, USA.
  • Basil Mustafa
    Google Research, Mountain View, CA, USA.
  • Sebastien Baur
    Synthetic Biology Group, Microbiology Department, Institut Pasteur, Paris, France.
  • Simon Kornblith
    Google Research, Mountain View, CA, USA.
  • Ting Chen
    CAS Key Laboratory of Tropical Marine Bio-resources and Ecology (LMB), Guangdong Provincial Key Laboratory of Applied Marine Biology (LAMB), South China Sea Institute of Oceanology, Chinese Academy of Sciences, Guangzhou 510301, China. chan1010@scsio.ac.cn.
  • Nenad Tomasev
    DeepMind, London, EC4A 3TW, UK.
  • Jovana Mitrović
    DeepMind, London, UK.
  • Patricia Strachan
    Google Research, Mountain View, CA, USA.
  • S Sara Mahdavi
    Google Research, Mountain View, CA, USA.
  • Ellery Wulczyn
    Google Health, Palo Alto, CA USA.
  • Boris Babenko
    Google Health, Google, Mountain View, CA, USA.
  • Megan Walker
    Google Research, Mountain View, CA, USA.
  • Aaron Loh
    Division of Neurosurgery, Department of Surgery, University Health Network and University of Toronto, Toronto, ON, Canada.
  • Po-Hsuan Cameron Chen
    Google Health, Palo Alto, CA USA.
  • Yuan Liu
    Department of General Surgery, Wuxi People's Hospital Affiliated to Nanjing Medical University, Wuxi, China.
  • Pinal Bavishi
    Google Health, Google, Mountain View, CA, USA.
  • Scott Mayer McKinney
    Google Health, Palo Alto, CA, USA. scottmayer@google.com.
  • Jim Winkens
    Google Health, London, UK.
  • Abhijit Guha Roy
    Department of Electrical Engineering, Indian Institute of Technology Kharagpur, West Bengal, India.
  • Zach Beaver
    Google Research, Mountain View, CA, USA.
  • Fiona Ryan
    Georgia Institute of Technology, Computer Science, Atlanta, GA, USA.
  • Justin Krogue
    Google Research, Mountain View, CA, USA.
  • Mozziyar Etemadi
    From the School of Electrical and Computer Engineering, Georgia Institute of Technology, Atlanta (O.T.I., M.B.P., A.Q.J., A.D., A.O.B.); Division of Cardiology (S.D., T.D.M., L.K.) and Department of Bioengineering and Therapeutic Sciences (S.R.), University of California, San Francisco; and Department of Anesthesiology and Department of Biomedical Engineering, Northwestern University, Chicago, IL (M.E., J.A.H.).
  • Umesh Telang
    Google Research, Mountain View, CA, USA.
  • Yun Liu
    Google Health, Palo Alto, CA USA.
  • Lily Peng
    Google Inc, Mountain View, California.
  • Greg S Corrado
    Google Health, Palo Alto, CA USA.
  • Dale R Webster
    Google Inc, Mountain View, California.
  • David Fleet
    Google Research, Mountain View, CA, USA.
  • Geoffrey Hinton
    1] Google, 1600 Amphitheatre Parkway, Mountain View, California 94043, USA. [2] Department of Computer Science, University of Toronto, 6 King's College Road, Toronto, Ontario M5S 3G4, Canada.
  • Neil Houlsby
    Google Research, Mountain View, CA, USA.
  • Alan Karthikesalingam
    Department of Outcomes Research, St George's Vascular Institute, London, SW17 0QT, United Kingdom.
  • Mohammad Norouzi
    From Google AI Healthcare, Google Research, Mountain View, California (Drs Liu, Kohlberger, Norouzi, Dahl, Peng, Hipp, and Stumpe); and Laboratory Department, Naval Medical Center, San Diego, California (Drs Smith, Mohtashamian, and Olson).
  • Vivek Natarajan
    Google, Mountain View, CA, USA.