EasyDeepLearn

How do FLOPs scale for a forward pass through an MLP and a transformer?

hard

Answer

  • MLP layer(in, out) on batch B: 2 * B * in * out FLOPs.
  • Transformer forward (per token, per layer) ≈ 2  ⋅  nlayers  ⋅  d  ⋅  dffn2\; \cdot \;n_{\mathrm{layers}}\; \cdot \;d\; \cdot \;d_{\mathrm{ffn}} (MLP)  +  2  ⋅  nlayers  ⋅  d  ⋅  (d  +  dhead)(\mathrm{MLP})\; + \;2\; \cdot \;n_{\mathrm{layers}}\; \cdot \;d\; \cdot \;(d\; + \;d_{\mathrm{head}}) (attention  linear  projections)  +  2  ⋅  nlayers  ⋅  seq  ⋅  d(\mathrm{attention}\;\mathrm{linear}\;\mathrm{projections})\; + \;2\; \cdot \;n_{\mathrm{layers}}\; \cdot \;\mathrm{seq}\; \cdot \;d (attention softmax + matmuls).
  • For long context, attention dominates because it's O(seq2  ⋅  d)O(\mathrm{seq}^{2}\; \cdot \;d).
  • Approximation: forward FLOPs ≈ 2 * N (params) per token — thus training a model with N params on T tokens takes ~6 * N * T FLOPs (forward + backward + optimizer).
Check yourself — multiple choice
  • FLOPs are independent of sequence length
  • Forward ≈ 2N per token; training ≈ 6·N·T; attention is O(seq2⋅d)O(\mathrm{seq}^{2} \cdot d)
  • Only training FLOPs matter
  • Attention has fewer FLOPs than the MLP

Rule of thumb: 6·N·T FLOPs to train N-param model on T tokens.

#fundamentals#architectures

Practise Deep Learning

214 interview questions in this topic.

Read the lessonFlashcardsTake the quiz

Related questions