\documentclass[border=10pt]{standalone}
\usepackage[cjk]{supertensor}

\stsetrole{q}{stTeal}
\stsetrole{k}{stOrange}
\stsetrole{s}{stCoral}

\stdim{T}{6}
\stdim{dh}{3}

\begin{document}
\begin{tikzpicture}
\ststage{SA}{每头打分：沿 $d_h$ 收缩}
\strow{rowA}{T}
  \ststack[role=q, bracket=true]{Q}{}{T}{dh}{3}
  \stglyph{mA}{$\times$}
  \ststack[role=k, bracket=true]{KT}{}{dh}{T}{3}
  \stglyph{eA}{$=$}
  \ststack[role=s]{S}{}{T}{T}{3}
\strowend
\stcaption{Q}{$\mathbf Q^{(i)}$}{$h\times T\times d_h$}
\stcaption{KT}{$\mathbf K^{(i)\top}$}{$h\times d_h\times T$}
\stcaption{S}{$\mathbf S^{(i)}$}{$h\times T\times T$}
\sttopformula{F}{$S^{(i)}=Q^{(i)}K^{(i)\top}$}
\stbbox{all}
\stmeaningbox{mb}{120mm}{all}
  {$T$ 时间；$d_h$ 头维；$h$ 头数画成 stack 深度}
  {$Q/K$ 是 value；$S$ 是 score，不是 mask}
  {沿 $d_h$ 收缩；$K^\top$ 换面，收缩边等长}
\end{tikzpicture}
\end{document}