ImageNet is the new MNIST
Chris Ying

Research SWE @ Google Braing.co/brain
on behalf of many people across Google

Goal: “Interactive ML supercomputing”●●●Hardware○Cloud TPUs○

TPU pods

Software○TensorFlow Datasets, Layers, and Estimator APIs (open-source)○

XLA compiler (open-source) with TPU backend

Research○

Understanding of generalization gap○

Large-batch training advances

Motivation(classical workflow)Motivation(classical workflow)x 10

Motivation(what's happening now)x 10

Motivation(our vision of the future)

x 1000ImageNet is the new MNISTMNIST: 60,000 B&W images

ImageNet: 1,281,167 color images

Motivating results
ResNet-50-v2 on ImageNet# of TPU devices

Batch size

Time to 90 epochs

Accuracy

1256

23 hours 22 minutes

76.6%

41024

5 hours 48 minutes

76.3%

164096

1 hour 30 minutes

76.5%

32819245 minutes

76.1%

6416384

22 minutes

75.0%

Only change between different runs is batch size (linearly scale LR) and hardware,
no model changes or hyperparameter re-tuning!

Cloud TPU

180 TFLOPS of computation, 64 GB of HBM memory, 2400 GB/s mem BWCloud TPU

TPUv2 Chip
HBM
8 GB

corevectorunit

●
●
●
●
●
●

45 TFLOPS
16 GB of HBM
600 GB/s mem BW
Vector unit: float32
Scalar unit: float32
Matrix unit (MXU):
float32 input/output,
reduced precisionmultiplication

scalarunit

MXU
128x128

corevector
unit

scalarunit

MXU
128x128HBM
8 GB

TPUv2 Chip
HBM
8 GB

Matrix Unit
●
●
●
●

●

core
vectorunit

scalarunit

corevector
unit

scalarunit

128x128
systolic array
16 GB of HBM
600 GB/s mem BW
float32 results*
Scalar unit: 32b float
MXU: 32b floataccumulation
but* reduced precision multiplication
reduced precision
for multipliers
45 TFLOPS

MXU
128x128MXU
128x128

HBM8 GBinputs

X33
X32

X23

X31

X22

X13

X21

X12

Computing y = Wx

Matrix Unit (MXU)

X11weights

Matrix Unit Systolic Array

W11

W12

W13

W21

W22

W23

W31

W32

W33

accumulation

Toy example: 3x3 systolicarray
W = 3x3 matrix
batch_size(x) = 3

Matrix Unit Systolic Array

Matrix Unit (MXU)

weights

inputs

X33
X32

X23

X31

X22

X13

X21

X12

W11
X11

W12

W13

W21

W22

W23

W31W32

W33

accumulation

Computing y = Wx
with W = 3x3, batch_size(x) = 3

Matrix Unit Systolic Array
X33

Matrix Unit (MXU)

weights

inputsComputing y = Wx
X32

X23

X31

X22

X13

W11
X21

W12X12+
W11X11

W13

W21
X11

W22

W23

W31

W32

W33

accumulation

with W = 3x3, batch_size(x) = 3

Matrix Unit Systolic Array
inputsComputing y = Wx
with W = 3x3, batch_size(x) = 3

Matrix Unit (MXU)

weights

X33
X32

X23

W11X31

W12X22+
W11X21

W13X13+...

W21
X21

W22X12+
W21X11

W23W31X11W32

W33

accumulation

Matrix Unit Systolic Array
inputsComputing y = Wx
with W = 3x3, batch_size(x) = 3

outputs

Matrix Unit (MXU)

weights

X33

W11W12X32+
W11X31

W13X23+
...

W21
X31

W22X22+
W21X21

W23X13+
...

W31
X21

W32X12+
W31X11

W33

accumulation

Y11 = W11X11 + W12X12 + W13X13

Matrix Unit Systolic Array
inputsComputing y = Wx
with W = 3x3, batch_size(x) = 3

Matrix Unit (MXU)

weights

outputs
W11W12

W13X33+...

Y21 = W11X21 + W12X22 + W13X23W21W22X32+W21X31

W23X23+...

Y12 = W21X11 + W22X12 + W23X13

W31X31W32X22+
W31X21

W33X13+...

accumulation

Y11 = W11X11 + W12X12 + W13X13

Matrix Unit Systolic Array
inputsComputing y = Wx
with W = 3x3, batch_size(x) = 3

Matrix Unit (MXU)

weights

outputs
W11W12W13

Y31 = W11X31 + W12X32 + W13X33

Y21 = W11X21 + W12X22 + W13X23

W21W22W23X33+...

Y22 = W21X21 + W22X22 + W23X23

Y12 = W21X11 + W22X12 + W23X13

W31W32X32+W31X31W33X23+...

Y13 = W31X11 + W32X12 + W33X13accumulation

Y11 = W11X11 + W

Matrix Unit Systolic Array
inputsComputing y = Wx
with W = 3x3, batch_size(x) = 3

Matrix Unit (MXU)

weights

outputs
W11W12

W13

W21

W22W23

W31

W32

W33X33+...

accumulation

Y31 = W11X11 + W12X12 + W13X13

Y21 = W11X21 + W

Y32 = W21X31 + W22X32 + W23X33

Y22 = W21X21 + W22X22 + W23X23

Y12 = W21X11 + W

Y23 = W31X21 + W32X22 + W33X23

Y13 = W31X11 + W32X12 + W33X13

Matrix Unit Systolic Arrayinputs

Computing y = Wx
with W = 3x3, batch_size(x) = 3

Matrix Unit (MXU)

weights

outputs
W11W12W13W21W22W23

W31W32W33

accumulation

Y31 = W11X11 + W

Y33 = W31X31 + W32X32 + W33X33

Y32 = W21X31 + W22X32 + W23X33

Y22 = W21X21 + W

Y23 = W31X21 + W32X22 + W33X23

Y13 = W31X11 + W

Cloud TPU Pod
64 Cloud TPUs in 2-D toroidal mesh11.5 petaFLOPS
4 terabytes of HBM memory

Accelerated Linear Algebra (XLA)●●●JIT / AOT compiler for linear algebra
Targets multiple backends, e.g. CPUs, GPUs, and TPUs
Compiler, runtime, and accelerator-specific optimizer

The life of a neural network:

model.py

TF Estimator code

TF GraphAccelerated Linear Algebra (XLA)●●●JIT / AOT compiler for linear algebra
Targets multiple backends, e.g. CPUs, GPUs, and TPUs
Compiler, runtime, and accelerator-specific optimizer

The life of a neural network:

XLA

model.py

TF Estimator code

TF Graph

Target-independent
optimizations

XLA

Target-specificcode generationLarge batch training●Understanding generalization gap (2016 N. Keskar et. al., 2017 E. Hoffer et.
al.)●

Relationship of batch size and noise scale (2018 S. Smith et. al.)●Learning rate scaling and schedule (2017 P. Goyal et. al.)●

New optimizers○

K-FAC*: approximate Fisher information matrix (2015 J. Martens)○

Neumann*: approximate inverse Hessian (2018 S. Krishnan et. al.)○

LARS: per-layer learning rate (2018 Y. You et. al.)* stick around after this talk to hear more about these!

Experiments
ResNet-50 training on ImageNet

validation accuracy

76.6%

45 minbatch size

# TPUs25611024

44096

168192

32

16384

64

hours to 90 epochs

Experiments

Experiments# of TPU devices

Batch size

Time to 90 epochs

Accuracy

32819244.9 minutes

76.1%

64819229.8 minutes

75.7%

64

16384

22.3 minutes

75.0%

646553617.5 minutes

65.4%

64

8192 → 16384[1]29.5 minutes

76.1%

Only change between different runs is batch size (linearly scale LR) and hardware,
no model changes or hyperparameter re-tuning!

[1] Don't Decay the Learning Rate, Increase theBatch Size (2018 S. Smith et. al)More than just ImageNet
Transformer model from "Attention
is All You Need" (2017 A. Vaswani et.
al.)

WMT’14 English-German translation
task

Adam optimizer - same learning rateschedule across configurations

batch size
(i/o tokens)# TPUsTime to
PPL=4.8

16k / 16k

1

17.9 hours

32k / 32k

4

3.5 hours

256k / 256k

161.1 hours

1M / 1M64

0.5 hours

Implications●Faster training enables neural architecture search○

Reinforcement learning architectures beat existing models
in accuracy and cost [1][1] Learning Transferable Architectures for
Scalable Image Recognition (2017 B. Zoph et. al)Implications●Faster training enables neural architecture search○

Reinforcement learning architectures beat existing models
in accuracy and cost [1]●What's the "new ImageNet"?
○

Full ImageNet (14M), Open Images (9M), YouTube-8M○

Performance increases logarithmically with data [2][1] Learning Transferable Architectures for
Scalable Image Recognition (2017 B. Zoph et. al)[2] Revisiting Unreasonable Effectiveness of
Data in Deep Learning Era (2017 C. Sun et. al)Pieter-janBrennanSamJonathanSameer

Zak

Quoc

BjarkeNoam

NaveenThank you!
chrisying@google.com
Chrisg.co/braing.co/tpusignup