Google's First Tensor Processing Unit: Architecture
thechipletter.substack.com
Google's First Tensor Processing Unit: Architecture
1–10 of 197 posts
Re: Google's First Tensor Processing Unit: Architecture
#2I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
Re: Google's First Tensor Processing Unit: Architecture
#3> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
Re: Google's First Tensor Processing Unit: Architecture
#4> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
Re: Google's First Tensor Processing Unit: Architecture
#5> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
Re: Google's First Tensor Processing Unit: Architecture
#6> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
If nothing else, the term "tensor" is shorter than "vectors and matrices," and then has the added benefit of representing n-dimensional arrays.
And, strictly speaking, a vector can be considered a 1xn (or nx1) matrix, so Matrix Processing Unit would have been fine.
Re: Google's First Tensor Processing Unit: Architecture
#7> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
You're right: fundamentally ML is about vector and matrix operations (1D and 2D). So then why are most ML programs 3D, 4D, and in a transformer sometimes up to 6D (?!)
One reasonable guess is that the third dimension is time. Actually not. It turns out that time is pretty rare in ML, and it's only (relatively) recently that it's been introduced into e.g. video models.
Another guess is that it's to represent "time" as in, think of how transformers work: they generate a token, then another given the previous, then a third given the first two, etc. That's a certain way of describing "time". But it turns out that transformers don't do this as a 3D or 4D dimension. It only needs to be 2D, because tokens are 1D -- if you're representing tokens over time, you get a 2D output. So even with a cutting edge model like transformers, you still only need plain old 2D matrix operations. The attention layer creates a mask, which ends up being 2D.
So then why do models get to 3D and above? Usually batching. You get a certain efficiency boost when you pack a bunch of operations together. And if you pack a bunch of 2D operations together, that third dimension is the batch dimension.
For images, you typically end up with 4D, with the convension N,C,H,W, which stands for "Batch, Channel, Height, Width". It can also be N,H,W,C, which is the same thing but it's packed in memory as red green blue, red green blue, etc instead of all the red pixels first, then all the green pixels, then all the blue pixels. This matters in various subtle ways.
I have no idea why the batch dimension is called N, but it's probably "number of images".
"Vector" wouldn't quite cover all of this, and although "tensor" is confusing, it's fine. It's the ham sandwich of naming conventions: flexible, satisfying to some, and you can make them in a bunch of different varieties.
Under the hood, TPUs actually flatten 3D tensors down into 2D matrix multiplications. I was surprised by this, but it makes total sense. The native size for a TPU is 8x128 -- you can think of it a bit like the native width of a CPU, except it's 2D. So if you have a 3x4x256 tensor, it actually gets flattened out to 12x256, then the XLA black box magic figures out how to split that across a certain number of 8x128 vector registers. Note they're called "vector registers" rather than "tensor registers", which is interesting. See https://cloud.google.com/tpu/docs/performance-guide
Re: Google's First Tensor Processing Unit: Architecture
#8> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
It's branding (see: TensorFlow); also, pretty much anything (linear) you would do with an arbitrarily ranked tensor can be expressed in terms of vector ops and matmuls
Re: Google's First Tensor Processing Unit: Architecture
#9> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
Re: Google's First Tensor Processing Unit: Architecture
#10> However, although tensors describe the relationship between arbitrary higher-dimensional arrays, in practice the TPU hardware that we will consider is designed to perform calculations associated with one and two-dimensional arrays. Or, more specifically, vector and matrix operations. I still don’t understand why the term “tensor” is used if it’s only vectors and matrices.
I was confused as hell for a long time when I first got into ML, until I figured out how to think about tensors in a visual way. You're right: fundamentally ML is about vector and matrix operations (1D and 2D). So then why are most ML programs 3D, 4D, and in a transformer sometimes up to 6D (?!) One reasonable guess is that the third dimension is time. Actually not. It turns out that time is pretty rare in ML, and it…