Lecture 10 · Transformers and ACT · November 18
Transformers and ACT From observations to a sequence of actions
How can a robot combine several observations and predict a coordinated movement?
Camera features Attention Action chunk Mathematics for Robotics · ROB-GY 6013 · NYU Tandon · Jing Zhang · Fall 2026
Zhao et al. (2023), §IV-B–C, Appendix C
Concept map
Transformers and ACT The mathematical route
Attention selects information. Position identifies order. A policy turns the resulting representation into actions.
Linear maps Scores → weights Value mixture Action chunk Vaswani et al. (2017), §3.2.1, Eq. 1
Definition
Tokens and projections A token is a feature vector
$$X=[x_1\ \cdots\ x_T]^T\in\mathbb R^{T\times D}$$
$x_i$ column feature vector in ℝᴰ $T,D$ tokens, features per token Rows identify tokens; columns identify features. Vaswani et al. (2017), §3.2.1, Eq. 1
Derivation
Tokens and projections Column weights and row-stored tokens
$$q_i=W_Qx_i\quad\Rightarrow\quad Q=XW_Q^T$$
$W_Q$ learned dₖ × D map $Q$ T × dₖ query matrix Vaswani et al. (2017), §3.2.1, Eq. 1
Definition
Tokens and projections Three learned views of the same input
$$Q=XW_Q^T,\quad K=XW_K^T,\quad V=XW_V^T$$
$Q,K$ T × dₖ: queries and keys $V$ T × dᵥ: values A query asks. A key describes what can be matched. A value supplies the information to combine.
Vaswani et al. (2017), §3.2.1, Eq. 1
Worked example
Tokens and projections A projection computed by hand
$$X=\begin{bmatrix}1&2&-1\\0&1&3\end{bmatrix},\quad W_Q=\begin{bmatrix}1&0&1\\0&2&-1\end{bmatrix}$$
$X$ two tokens with three features $W_Q$ two output coordinates $Q=XW_Q^T=\begin{bmatrix}0&5\\3&-1\end{bmatrix}$
Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Think, explain, compare
Tokens and projections A new token, the same learned map
$x=(2,-1,1)^T$, $W_Q=\begin{bmatrix}1&0&1\\0&2&-1\end{bmatrix}$. Find $q=W_Qx$.
A $(3,-3)^T$✓ B $(1,-1)^T$C A three-coordinate vector
Solution 1 / 2
Apply the two output rows $x$: input column; $W_Q$: 2 × 3 map; $q$: output column.
$q_1=1(2)+0(-1)+1(1)=3$. $q_2=0(2)+2(-1)-1(1)=-3$. $q=(3,-3)^T$ has two coordinates. Solution 2 / 2
Check the alternatives $q$: projected token.
A is correct. B drops coefficients from the map. C keeps the input width, but the two weight rows define two outputs. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Definition
Scaled dot-product attention A score compares a query with a key
$$s_{ij}=\frac{q_i^Tk_j}{\sqrt{d_k}}$$
$i,j$ query index, key index $q_i,k_j$ dₖ-dimensional columns $s_{ij}$ scalar compatibility score The key direction changes; the query stays fixed. Vaswani et al. (2017), §3.2.1, Eq. 1
Think, explain, compare
Scaled dot-product attention Alignment and length both affect the score
$q=(1,0)^T$, $k_1=(1,0)^T$, $k_2=(2,0)^T$, $d_k=2$. Which key has the larger score?
A Key 1B Key 2✓ C They must tie
Solution 1 / 1
Evaluate the dot products $q$: query; $k_1,k_2$: keys; $d_k=2$.
$q^Tk_1=1$ and $q^Tk_2=2$. $s_1=1/\sqrt2<s_2=2/\sqrt2$. B is correct. Equal directions alone do not imply equal dot products. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Derivation
Scaled dot-product attention All query-key pairs in one matrix
$$S=\frac{QK^T}{\sqrt{d_k}}\in\mathbb R^{T_q\times T_k}$$
$Q$ Tq × dₖ query rows $K$ Tk × dₖ key rows $S_{ij}$ query i versus key j Each row compares one query with all keys. Vaswani et al. (2017), §3.2.1, Eq. 1
Definition
Scaled dot-product attention Weights combine values
$$P_{ij}=\frac{e^{S_{ij}}}{\sum_\ell e^{S_{i\ell}}},\qquad y_i=\sum_jP_{ij}v_j$$
$P_{ij}$ weight of key j for query i $v_j,y_i$ dᵥ-dimensional value and output Vaswani et al. (2017), §3.2.1, Eq. 1
Worked example
Scaled dot-product attention Equal scores give an ordinary average
$$s=(0,0),\quad v_1=(2,0)^T,\quad v_2=(0,4)^T$$
$s;v_1,v_2$ one query’s scores; two value vectors $p=(1/2,1/2)$, so $y=(1,2)^T$.
Equal weights place the output halfway between the value vectors. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Think, explain, compare
Scaled dot-product attention Unequal scores change the mixture
$s=(0,\log 3)$, $v_1=(2,0)^T$, $v_2=(0,4)^T$. Find the output.
A $(1,2)^T$B $(1/2,3)^T$✓ C $(0,4)^T$
Solution 1 / 2
Convert scores to weights $s$: score row; $p$: its softmax weights.
$e^{s}=(1,3)$; the sum is $4$. $p=(1/4,3/4)$. Solution 2 / 2
Mix the value vectors $v_1=(2,0)^T$, $v_2=(0,4)^T$.
$y=\frac14v_1+\frac34v_2=(1/2,3)^T$. B is correct. A treats scores as equal; C replaces soft weighting with a hard maximum. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Motivation
Scaled dot-product attention Why divide by the square root?
Wider query-key vectors can produce more variable dot products. Very unequal scores make softmax nearly one-hot.
Scaling a fixed score gap changes the softmax weights. Vaswani et al. (2017), §3.2.1, Eq. 1
Derivation
Scaled dot-product attention Variance of one product
$$\mathbb E[q_rk_r]=0,\quad\operatorname{Var}(q_rk_r)=\mathbb E[q_r^2]\mathbb E[k_r^2]=1$$
$q_r,k_r$ independent, mean 0, variance 1 Vaswani et al. (2017), §3.2.1, Eq. 1
Derivation
Scaled dot-product attention Variance of the sum
$$\operatorname{Var}(q^Tk)=d_k,\qquad\operatorname{Var}\!\left(\frac{q^Tk}{\sqrt{d_k}}\right)=1$$
$d_k$ number of independent coordinate products Vaswani et al. (2017), §3.2.1, Eq. 1
Think, explain, compare
Scaled dot-product attention The scale is determined by the assumptions
Independent zero-mean, unit-variance coordinates; $d_k=16$. What is the standard deviation of the unscaled dot product?
A 16B 4✓ C 1
Solution 1 / 1
Separate variance from standard deviation $u=q^Tk$; $d_k=16$.
$\operatorname{Var}(u)=16$. $\operatorname{SD}(u)=\sqrt{16}=4$. B is correct. A confuses variance with standard deviation; C describes the scaled score. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Derivation
Stable attention and masks The largest logit can be subtracted
$$\frac{e^{s_j-c}}{\sum_\ell e^{s_\ell-c}}=\frac{e^{-c}e^{s_j}}{e^{-c}\sum_\ell e^{s_\ell}}=p_j$$
$s_j$ logit: score before softmax $c$ same shift for every entry in the row Blanchard, Higham & Higham (2021), §4
Worked example
Stable attention and masks A large-logit example
$$s=(1000,1000+\log3),\quad s-\max s=(-\log3,0)$$
$s$ one score row $e^{s-\max s}=(1/3,1)$; normalize to $(1/4,3/4)$.
Blanchard, Higham & Higham (2021), §4 · Course example
Think, explain, compare
Stable attention and masks A shift or a temperature?
A score row is $s=(0,\log3)$. Which change preserves its softmax weights?
A Subtract 5 from both entries✓ B Multiply both entries by 2C Add 5 only to the second entry
Solution 1 / 1
Identify the invariance $s$: logits; softmax normalizes their exponentials.
$\operatorname{softmax}(s-5\mathbf1)=\operatorname{softmax}(s)$. $\operatorname{softmax}(2s)=(1/10,9/10)$. A is correct. B changes the score gap; C changes one exponential relative to the other. Blanchard, Higham & Higham (2021), §4 · Course example
Definition
Stable attention and masks A causal mask forbids future keys
$$M_{ij}=\begin{cases}0&j\le i\\-\infty&j>i\end{cases},\qquad P=\operatorname{softmax}_{\rm row}(S+M)$$
$i,j$ query position, key position $M$ additive causal mask Rows are queries; columns are keys. The upper triangle is forbidden. Vaswani et al. (2017), §3.2.1, Eq. 1
Definition
Stable attention and masks Padding and causality answer different questions
$$S:[B,H,T_q,T_k],\quad V_k:[B,1,1,T_k],\quad C:[1,1,T_q,T_k]$$
$V_k,C$ Boolean key-valid and causal masks; True means allowed Broadcast $V_k\land C$ across heads and queries. Query validity $V_q:[B,1,T_q,1]$ is checked separately.
Vaswani et al. (2017), §3.2.1, Eq. 1
Think, explain, compare
Stable attention and masks Mask before normalizing
Scores are $(0,\log3)$; the second key is forbidden. Which output weights are correct?
A $(1,0)$✓ B $(1/4,0)$C $(1/4,3/4)$
Solution 1 / 2
Normalize only over permitted keys $s$: score row; key 2 is unavailable.
$s+M=(0,-\infty)$. $\exp(s+M)=(1,0)$ and its sum is $1$. $p=(1,0)$. A is correct. Solution 2 / 2
Why the other operations differ $p$: normalized attention weights.
B zeroes a probability after normalization and leaves row sum 1/4. C ignores the mask. Renormalizing allowed probabilities could work, but masking logits is direct. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Boundary case
Stable attention and masks Every valid query needs an allowed key
If every logit is $-\infty$, softmax has denominator zero. Subtracting the row maximum also produces $-\infty-(-\infty)$.
Vaswani et al. (2017), §3.2.1, Eq. 1
Think, explain, compare
Stable attention and masks Diagnose the NaN without changing the task
A padded query has no allowed key. Stable softmax returns NaN. Which repair preserves the declared information rule?
A Use uniform weights over all keysB Mark the query invalid and output zero✓ C Remove all masks for the batch
Solution 1 / 1
Separate invalid queries from valid predictions $M$: mask; padded queries are excluded from the loss.
Detect a row with no permitted key before normalization. For that invalid row, use zero weights/output and exclude its loss. B is correct. A and C read forbidden entries. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Integrated worked example
Stable attention and masks One attention query, end to end
$$q=(\sqrt2,0)^T,\quad K=\begin{bmatrix}0&1\\\log3&0\end{bmatrix},\quad V=\begin{bmatrix}2&0\\0&4\end{bmatrix}$$
$d_k=2$ two keys, both permitted Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Worked example
Stable attention and masks The output follows the same data path
$$s=(0,\log3)\Rightarrow p=(1/4,3/4)\Rightarrow y=(1/2,3)^T$$
$s$ scaled scores $p$ row weights $y$ mixture of (2,0)ᵀ and (0,4)ᵀ Weights (1/4, 3/4) place the output at (1/2, 3). Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Think, explain, compare
Stable attention and masks Transfer: change the permitted information
$q=(\sqrt2,0)^T$, $k_1=(0,1)^T$, $k_2=(\log3,0)^T$; $v_1=(2,0)^T$, $v_2=(0,4)^T$. Only key 1 is allowed.
A $(1/2,3)^T$B $(2,0)^T$✓ C $(0,0)^T$
Solution 1 / 1
Use the same scores with a different mask $s$: scaled scores; $p$: allowed-key weights.
$s=(0,\log3)$, but $s+M=(0,-\infty)$. $p=(1,0)$, so $y=1v_1+0v_2=(2,0)^T$. B is correct. A ignores masking; C discards a valid key. Vaswani et al. (2017), §3.2.1, Eq. 1 · Course example
Motivation
Multi-head attention One mixture can hide useful differences
Several heads learn different comparisons and preserve several value mixtures before combining them.
Tokens Different comparisons Different mixtures Output map Vaswani et al. (2017), §3.2.2
Definition
Multi-head attention A feature axis can be split into heads
$$Q:[B,T,D]\to[B,T,H,d_h],\qquad D=Hd_h$$
$B,T$ batch size, tokens $H,d_h$ heads, features per head Q = (X @ Wq.T).reshape(B,T,H,dh)
A packed feature row splits into two head blocks. Vaswani et al. (2017), §3.2.2
Derivation
Multi-head attention Transpose puts matrix axes last
$$[B,T,H,d_h]\to[B,H,T,d_h],\quad S:[B,H,T,T]$$
$T,d_h$ last two axes used by matrix multiplication Vaswani et al. (2017), §3.2.2
Think, explain, compare
Multi-head attention Trace a tensor through the heads
$B=2$, $T=3$, $D=8$, $H=2$. Find query and score shapes after splitting and transposing.
A Q [2,2,3,4]; S [2,2,3,3]✓ B Q [2,3,2,4]; S [2,3,2,2]C Q [2,2,3,8]; S [2,2,8,8]
Solution 1 / 2
Name each axis $B$: batch; $T$: tokens; $H$: heads; $d_h=D/H$.
$d_h=8/2=4$. $Q:[B,H,T,d_h]=[2,2,3,4]$. $S:[B,H,T,T]=[2,2,3,3]$. Solution 2 / 2
Check the alternatives $S$: query-by-key scores in each head.
A is correct. B omits the token/head transpose. C fails to split the width and confuses features with token counts. Vaswani et al. (2017), §3.2.2 · Course example
Definition
Multi-head attention Concatenate, then mix across heads
$$C=\operatorname{Concat}(Y^{(1)},\ldots,Y^{(H)}),\quad Y=CW_O^T$$
$C$ T × Hdᵥ concatenated outputs $W_O$ D × Hdᵥ learned output map C = Y.transpose(1,2).reshape(B,T,H*dv)out = C @ Wo.T
Vaswani et al. (2017), §3.2.2
Worked example
Multi-head attention Two heads, one explicit output map
$$c=(1,2,3,-1)^T,\quad W_O=\begin{bmatrix}1&0&1&0\\0&1&0&1\end{bmatrix}$$
$c$ head 1: (1,2), head 2: (3,−1) $W_Oc=(1+3,2-1)^T=(4,1)^T$.
Vaswani et al. (2017), §3.2.2 · Course example
Think, explain, compare
Multi-head attention Transfer: a different output map
$c=(1,2,3,-1)^T$, $W_O=\begin{bmatrix}0&0&1&0\\1&0&0&-1\end{bmatrix}$. Find $W_Oc$.
A $(3,2)^T$✓ B $(4,1)^T$C $(1,2,3,-1)^T$
Solution 1 / 1
Read the new map $c$: concatenated heads; $W_O$: two output rows.
First output: $0+0+3+0=3$. Second output: $1+0+0-(-1)=2$. A is correct. B reuses the earlier map; C stops before projection. Vaswani et al. (2017), §3.2.2 · Course example
Implementation
Multi-head attention The forward pass follows the equations
Q,K,V: [B,H,T,dₕ]S = Q @ K.transpose(-2,-1) / sqrt(dₕ)S = masked_fill(S, forbidden, -inf)P = softmax(S, dim=-1)Y = P @ V
Vaswani et al. (2017), §3.2.2
Think, explain, compare
Multi-head attention Correct shapes can hide a wrong normalization
Scores have shape [B,H,Tq,Tk]. A programmer uses softmax(S,dim=-2). What is normalized?
A Keys for each queryB Queries for each key✓ C Features of each token
Solution 1 / 1
Read the axis labels $S_{b,h,i,j}$ compares query i with key j.
The second-last axis is i, the query index. dim=−2 makes columns sum to one across queries. B describes the bug. Correct attention uses dim=−1 for key normalization. Vaswani et al. (2017), §3.2.2 · Course example
Think, explain, compare
Multi-head attention Transfer: broadcast a key-validity mask
Scores have shape [2,3,4,5]: batch, heads, queries, keys. Key validity has shape [2,5]. Which reshape shares the key rule across every head and query?
A [2,1,1,5]✓ B [2,5,1,1]C [1,1,2,5]
Solution 1 / 1
Align the batch and key axes $B=2$, $H=3$, $T_q=4$, $T_k=5$.
Use singleton head and query axes: [2,1,1,5]. Broadcast those singleton axes to [2,3,4,5]. A is correct. B places validity on the head axis; C moves batch into the query axis. Vaswani et al. (2017), §3.2.2 · Course example
Derivation
Position and order Content-only attention follows a permutation
$$\operatorname{Attn}(\Pi X)=\Pi\operatorname{Attn}(X)$$
$\Pi$ token permutation; no positional terms or fixed mask Vaswani et al. (2017), §3.5
Think, explain, compare
Position and order What does content alone fail to identify?
The same two feature vectors are presented in opposite orders to unmasked attention with no position features.
A Dot products recover their original orderB Outputs move with the input permutation✓ C Every output becomes zero
Solution 1 / 1
Apply the stated symmetry $\Pi$: swap of token positions.
$X\mapsto\Pi X$ implies $Y\mapsto\Pi Y$. B is correct: relationships are preserved without an original-position label. A invents missing information; C does not follow from equivariance. Vaswani et al. (2017), §3.5 · Course example
Definition
Position and order Sinusoids give each position a pattern
$$\mathrm{PE}_{p,2r}=\sin(p\omega_r),\quad\mathrm{PE}_{p,2r+1}=\cos(p\omega_r)$$
$p$ token position $\omega_r=10000^{-2r/D}$ frequency of coordinate pair r Add the pattern to the token features before projection.
One sine/cosine pair across token positions. Vaswani et al. (2017), §3.5
Derivation
Position and order A shift follows angle-addition identities
$$\begin{bmatrix}\sin((p+\Delta)\omega)\\\cos((p+\Delta)\omega)\end{bmatrix}=\begin{bmatrix}\cos(\Delta\omega)&\sin(\Delta\omega)\\-\sin(\Delta\omega)&\cos(\Delta\omega)\end{bmatrix}\begin{bmatrix}\sin(p\omega)\\\cos(p\omega)\end{bmatrix}$$
$p,\Delta,\omega$ position, offset, one frequency Vaswani et al. (2017), §3.5
Think, explain, compare
Position and order Transfer: one sinusoidal pair at a new position
One coordinate pair is $(\sin(p\omega),\cos(p\omega))$. Take $\omega=\pi/2$ and move from p=0 to p=1.
A (1,0)✓ B (0,1)C (−1,0)
Solution 1 / 1
Evaluate in the stated coordinate order $\omega=\pi/2$; the pair is sine first, cosine second.
At $p=0$: $(\sin0,\cos0)=(0,1)$. At $p=1$: $(\sin(\pi/2),\cos(\pi/2))=(1,0)$. A is correct. B keeps the old position; C uses the wrong sine sign. Vaswani et al. (2017), §3.5 · Course example
Definition
Position and order Rotary embeddings act inside the comparison
$$\widetilde q_p=R(p\omega)q_p,\quad\widetilde k_s=R(s\omega)k_s$$
$R(\theta)$ 2D rotation by θ $p,s$ query and key positions Rotate each coordinate pair at its own frequency. Values need not rotate.
A query pair rotates with position while keeping its length. Su et al., RoFormer (2021), §3.2
Derivation
Position and order Relative positions appear in the dot product
$$\widetilde q_p^T\widetilde k_s=q_p^TR(p\omega)^TR(s\omega)k_s=q_p^TR((s-p)\omega)k_s$$
$s-p$ key position relative to the query Su et al., RoFormer (2021), §3.2
Worked example
Position and order A quarter-turn example
$$q=k=(1,0)^T,\quad p=0,\quad s=1,\quad\omega=\pi/2$$
$p,s$ positions $\omega$ teaching example’s angle per position $\widetilde q=(1,0)^T$, $\widetilde k=(0,1)^T$, so their dot product is $0$.
Su et al., RoFormer (2021), §3.2 · Course example
Think, explain, compare
Position and order Transfer: shift both positions
$q=k=(1,0)^T$, $\omega=\pi/2$. Change $(p,s)$ from $(0,1)$ to $(2,3)$. Find the raw dot product.
A 0✓ B 1C −1
Solution 1 / 1
Use the relative angle $p,s$: query/key positions; content unchanged.
$(s-p)\omega=(3-2)\pi/2=\pi/2$. $q^TR(\pi/2)k=0$. A is correct. B uses zero offset; C uses a half-turn. Su et al., RoFormer (2021), §3.2 · Course example
Definition
Transformer blocks The pre-LN block has two residual paths
$$U=X+\operatorname{MHA}(\operatorname{LN}(X)),\quad Y=U+\operatorname{FFN}(\operatorname{LN}(U))$$
$X,U,Y$ T × D representations $LN$ LayerNorm within each token X LN → MHA + X LN → FFN + U Y Xiong et al. (2020), §2, Fig. 1
Recall
Transformer blocks LayerNorm uses the token feature axis
$$\widehat x_r=\frac{x_r-\mu}{\sqrt{v+\epsilon}},\quad\mu=\frac1D\sum_rx_r,\quad v=\frac1D\sum_r(x_r-\mu)^2$$
$x$ one token’s D features $\epsilon>0$ stability constant Output: $\gamma\odot\widehat x+\beta$, with learned D-coordinate scale and shift.
Xiong et al. (2020), §2, Fig. 1
Think, explain, compare
Transformer blocks Transfer: normalize one token
One token is $x=(1,3)$. Use $\gamma=(1,1)$, $\beta=(0,0)$ and $\epsilon=0$ for this nonconstant numerical example.
A (−1,1)✓ B (−1/√2, 1/√2)C (1,3)
Solution 1 / 2
Use the feature mean and mean squared deviation $D=2$ features; no batch or token averaging.
$\mu=(1+3)/2=2$; $x-\mu=(-1,1)$. $v=[(-1)^2+1^2]/2=1$. $\widehat x=(-1,1)/\sqrt1=(-1,1)$. A is correct. Solution 2 / 2
Diagnose the alternative statistics $\gamma=1$ and $\beta=0$ preserve the normalized values.
B uses the sample-variance divisor D−1; LayerNorm uses D. C omits centering and scaling. In real implementations keep positive epsilon. Xiong et al. (2020), §2, Fig. 1 · Course example
Definition
Transformer blocks The feed-forward branch transforms each token
$$f(x)=W_2\phi(W_1x+b_1)+b_2$$
$W_1,W_2$ F × D and D × F maps $\phi$ coordinatewise nonlinearity $b_1,b_2$ F- and D-coordinate biases Vaswani et al. (2017), §3.3
Think, explain, compare
Transformer blocks A residual connection constrains the shape
$X$ has shape [2,5,12], while the attention branch returns [2,5,8]. Can they be directly added?
A Yes, ordinary broadcasting repairs itB No, the branch must return width 12✓ C Yes, because both have five tokens
Solution 1 / 1
Compare corresponding coordinates $X$: residual input; branch: learned update.
Feature widths 12 and 8 differ and neither is one. Ordinary broadcasting cannot align these feature coordinates. B is correct: use the output projection to return width 12. Xiong et al. (2020), §2, Fig. 1 · Course example
Think, explain, compare
Transformer blocks Transfer: a missing path changes the architecture
$Y=X+FFN(MHA(X))$ is described as the pre-LN block. What is missing?
A Only position informationB Both normalizations and the intermediate residual✓ C A final softmax only
Solution 1 / 1
Compare the computation graph $U$: intermediate residual representation.
$U=X+MHA(LN(X))$. $Y=U+FFN(LN(U))$. B is correct. Position features or final softmax do not restore these missing paths. Xiong et al. (2020), §2, Fig. 1 · Course example
Derivation
Cost and the KV cache Pairwise comparison creates quadratic work
$$QK^T:O(T^2d_h),\quad PV:O(T^2d_h),\quad Hd_h=D$$
$T,D$ tokens, total feature width $H$ heads Mixing: $O(T^2D)$. Projections also cost $O(TD^2)$.
Query-key pair counts grow quadratically with token length. Vaswani et al. (2017), §3.2.2
Worked example
Cost and the KV cache Memory follows the score shape
$$B=1,\ H=8,\ T=1024:\quad BHT^2=8{,}388{,}608$$
$B,H,T$ batch, heads, tokens At 2 bytes per entry: 16 MiB for one materialized score tensor.
Vaswani et al. (2017), §3.2.2 · Course example
Think, explain, compare
Cost and the KV cache Transfer: a longer sequence
Keep B, H and datatype fixed. Double T. How does the materialized score-tensor memory change?
A ×2B ×4✓ C Unchanged
Solution 1 / 1
Count both token axes $S$ has shape [B,H,T,T].
$BH(2T)^2=4BHT^2$. B is correct. A counts only the extra queries and forgets the extra keys. C confuses fixed feature width with a fixed pair count. Vaswani et al. (2017), §3.2.2 · Course example
Definition
Cost and the KV cache Unchanged causal prefixes permit reuse
Earlier causal hidden states cannot read a newly appended token. Their keys and values can be cached per layer.
An unchanged prefix stays cached while one new row is appended. Shazeer (2019), §2, incremental attention
Derivation
Cost and the KV cache One new query reads all cached keys
$$q_t^TK_{1:t}^T:\ [1,d_h][d_h,t]=[1,t]$$
$q_t$ new query column $K_{1:t}$ t cached key rows Per-step mixing: $O(tD)$. Cache storage: $O(tD)$ per layer and batch item.
K = concat([Kold, knew], axis=-2)V = concat([Vold, vnew], axis=-2)
Shazeer (2019), §2, incremental attention
Think, explain, compare
Cost and the KV cache Transfer: can this memory be reused?
Camera tokens are replaced by a new frame. Can their old cross-attention keys and values serve as exact representations of the new frame?
A Yes, when shapes matchB No, recompute the changed memory✓ C Yes, if action horizon is fixed
Solution 1 / 1
Trace projections back to their inputs $E$: encoded camera tokens; K,V are their projections.
$K=EW_K^T$, $V=EW_V^T$. A new frame can change E and both projections. B is correct. Matching shape or horizon does not imply matching values. Shazeer (2019), §2, incremental attention · Course example
Worked example
Cost and the KV cache Cache validity and cumulative decoding cost
$$\sum_{t=1}^Tt=\frac{T(T+1)}2,\quad\sum_{t=1}^Tt^2=\frac{T(T+1)(2T+1)}6$$
$t,T$ current and final prefix length Cached mixing $O(T^2D)$; naïve repeated full-prefix mixing $O(T^3D)$.
Shazeer (2019), §2, incremental attention · Course example
Think, explain, compare
Cost and the KV cache Transfer: count incremental comparisons
Generate four tokens from an empty prefix. Count query-key comparisons for one head, including each current token. Compare cached decoding with full-prefix recomputation.
A Cached 10; recomputed 30✓ B Cached 4; recomputed 16C Cached 16; recomputed 16
Solution 1 / 1
Add work over all four decoding steps $t=1,2,3,4$; one new query uses t keys.
Cached: $1+2+3+4=10$. Full-prefix reruns: $1^2+2^2+3^2+4^2=30$. A is correct. B drops past keys and repeated passes; C counts only one final full-sequence pass. Shazeer (2019), §2, incremental attention · Course example
Motivation
Action chunking A coordinated movement spans several decisions
A policy can predict a short future action sequence from the current observation, then gather new observations as execution continues.
Each row is a chunk prediction; columns are target action times. Zhao et al. (2023), §IV-A
Definition
Action chunking Prediction time and target time are different indices
$$\widehat A_t\in\mathbb R^{K\times d_a},\quad\widehat A_t[j,:]=\hat a_{t+j|t}^T,\quad0\le j<K$$
$t$ observation/prediction time $K,d_a$ chunk length, action dimension Zhao et al. (2023), §IV-A
Think, explain, compare
Action chunking Chunk length is not execution length
A policy predicts K=6 future actions and executes E=2 before reading a new observation. Which targets from the current chunk are executed?
A Times t through t+5B Times t and t+1✓ C Only time t+5
Solution 1 / 1
Separate prediction and execution $K=6$: predicted rows; $E=2$: executed prefix length.
The chunk covers target times t,…,t+5. Its first two rows target t and t+1. B is correct. A confuses K with E; C chooses the last row rather than the prefix. Zhao et al. (2023), §IV-A · Course example
Definition
Temporal ensembling Overlapping chunks predict the same target
$$\bar a_s=\frac{\sum_{j=0}^{J-1}w_j\hat a_{s|r_j}}{\sum_{j=0}^{J-1}w_j},\qquad w_j=e^{-mj}$$
$r_0<\cdots<r_{J-1}$ prediction times, oldest first $s,m$ common target time, nonnegative weight parameter Combine one target-time column across overlapping chunks. Zhao et al. (2023), §IV-A
Worked example
Temporal ensembling One target, two predictions
$$\hat a_{s|r_0}=2,\quad\hat a_{s|r_1}=4,\quad m=\log2$$
$r_0<r_1$ oldest then newest prediction $s$ same target time $w=(1,1/2)$, so $\bar a_s=(2+2)/(3/2)=8/3$.
Zhao et al. (2023), §IV-A · Course example
Think, explain, compare
Temporal ensembling Transfer: three estimates for one target
Oldest-to-newest predictions for target s: 1, 3, 5. Let $m=\log2$. Find the ensemble.
A 3B 15/7✓ C 27/7
Solution 1 / 2
Weight, sum, normalize $j=0,1,2$: oldest-to-newest; all predictions target s.
$w=(1,1/2,1/4)$, with sum $7/4$. Weighted sum $=1+3/2+5/4=15/4$. $\bar a_s=(15/4)/(7/4)=15/7$. Solution 2 / 2
Interpret the result $\bar a_s$: ensemble action.
B is correct. A uses equal weights. C reverses the weights and favors the newest prediction. Zhao et al. (2023), §IV-A · Course example
Interpretation
Temporal ensembling Averages need compatible action representations
Average predictions for the same target time in the same coordinate system. Smoothing alone does not enforce collision, dynamics or joint-limit constraints.
Zhao et al. (2023), §IV-A
Definition
ACT: observations to actions ACT combines observation tokens with action queries
Action Chunking with Transformers (ACT): encode image features, joint state and a latent variable; decode K ordered action queries into K motor targets.
Images + joints + z Encoded memory Action query positions Decoder K × dₐ targets Zhao et al. (2023), §IV-B–C, Appendix C
Worked example
ACT: observations to actions A concrete observation-token count
$$4(15\cdot20)+2=1202$$
$4$ camera views in the ACT paper $15\times20$ feature grid per view $+2$ joint-state and latent tokens Each token has 512 features in the reported model. Spatial positions accompany image features.
Zhao et al. (2023), §IV-B–C, Appendix C · Course example
Definition
ACT: observations to actions Action queries use cross-attention
$$Q_a\in\mathbb R^{K\times d_k},\quad K_o\in\mathbb R^{T_o\times d_k},\quad S_a\in\mathbb R^{K\times T_o}$$
$K,T_o$ action queries, observation tokens Zhao et al. (2023), §IV-B–C, Appendix C
Think, explain, compare
ACT: observations to actions Transfer: encoder tokens and action slots
A model has 20 observation tokens and predicts a chunk of 6 actions. Each action query reads every observation token. What is one head’s cross-attention score shape?
A 20 × 20B 6 × 20✓ C 6 × 6
Solution 1 / 1
Identify query and key collections $K=6$: action queries; $T_o=20$: observation keys.
$Q_a:[6,d_k]$, $K_o^T:[d_k,20]$. $Q_aK_o^T:[6,20]$. B is correct. A and C replace cross-attention with self-attention on one collection. Zhao et al. (2023), §IV-B–C, Appendix C · Course example
Motivation
ACT: latent-variable training Demonstrations can differ at the same observation
A latent variable z can represent variation among demonstrated action chunks. Training infers z using the demonstration; deployment follows a separate rule.
Demo chunk + joints Latent distribution Observation + latent Predicted chunk Zhao et al. (2023), §IV-B–C, Appendix C
Definition
ACT: latent-variable training The encoder predicts a diagonal Gaussian
$$q_\phi(z\mid A_t,r_t)=\mathcal N(\mu,\operatorname{diag}(\sigma^2)),\quad p(z)=\mathcal N(0,I)$$
$\mu,\sigma$ latent mean and positive standard deviation $r_t$ joint observation; Aₜ is the demonstration Kingma & Welling (2014), §2, Appendix B
Derivation
ACT: latent-variable training Reparameterization exposes a differentiable path
$$\epsilon\sim\mathcal N(0,I),\qquad z=\mu+\sigma\odot\epsilon$$
$\epsilon$ parameter-independent random draw $\mu,\sigma$ encoder outputs $z$ latent sample Changing the mean shifts a Gaussian without changing its width. Kingma & Welling (2014), §2, Appendix B
Think, explain, compare
ACT: latent-variable training Log-variance is not the standard deviation
One latent coordinate has $\mu=1$, $\ell=\log\sigma^2=\log4$, and $\epsilon=-1$. Find z.
A −3B −1✓ C 1
Solution 1 / 1
Convert the parameter before sampling $\ell$: log-variance; $\sigma$: standard deviation.
$\sigma^2=e^\ell=4$, so $\sigma=e^{\ell/2}=2$. $z=\mu+\sigma\epsilon=1+2(-1)=-1$. B is correct. A uses variance as scale; C ignores the sampled noise. Kingma & Welling (2014), §2, Appendix B · Course example
Derivation
ACT: latent-variable training A bound makes the latent objective tractable
$$\begin{aligned}\log p(A,z\mid o)&=\log p(A\mid o)+\log p(z\mid A,o)\\\mathbb E_q\log\frac{p(A,z\mid o)}{q(z)}&=\log p(A\mid o)-D_{KL}(q\Vert p(z\mid A,o))\end{aligned}$$
$p,q$ model and encoder densities; θ suppressed for space Rules: Bayes’ factorization, expectation, then the KL definition.
Kingma & Welling (2014), §2, Appendix B
Derivation
ACT: latent-variable training Nonnegative KL gives the evidence lower bound
$$\begin{aligned}\mathcal L&=\mathbb E_q\log\frac{p(A\mid o,z)p(z)}{q(z)}\\&=\mathbb E_q\log p(A\mid o,z)-\mathbb E_q\log\frac{q(z)}{p(z)}\\&=\mathbb E_q\log p(A\mid o,z)-D_{KL}(q\Vert p)\end{aligned}$$
$\mathcal L$ ELBO; p(z) is the prior Nonnegative posterior-gap KL gives $\log p(A\mid o)\ge\mathcal L$.
Kingma & Welling (2014), §2, Appendix B
Think, explain, compare
ACT: latent-variable training Transfer: which divergence creates the gap?
$\log p_\theta(A\mid o)=\mathcal L_{ELBO}+G$. Which G makes this an exact identity?
A $D_{KL}(q\Vert p_\theta(z\mid A,o))$✓ B $-D_{KL}(q\Vert p(z))$C The action prediction error alone
Solution 1 / 1
Distinguish two different distributions $q$: encoder; posterior: conditioned on the observed target.
$G=D_{KL}(q\Vert\text{true posterior})\ge0$. The prior KL is already inside $\mathcal L_{ELBO}$ with a minus sign. A is correct. B confuses regularization with the posterior approximation gap. Kingma & Welling (2014), §2, Appendix B · Course example
Definition
ACT: latent-variable training ACT balances reconstruction and latent regularization
$$\mathcal L=\frac1{Kd_a}\sum_{j,r}|\widehat A_{jr}-A_{jr}|+\beta D_{KL}(q_\phi\Vert\mathcal N(0,I))$$
$A,\widehat A$ target and predicted action chunks $\beta$ chosen regularization weight Zhao et al. (2023), §IV-B–C, Appendix C
Derivation
ACT: latent-variable training Gaussian KL follows from a log-density ratio
$$\log\frac{q_j(z_j)}{p_j(z_j)}=-\log\sigma_j-\frac{(z_j-\mu_j)^2}{2\sigma_j^2}+\frac{z_j^2}{2}$$
$q_j,p_j$ N(μⱼ,σⱼ²) and N(0,1) densities Rule: cancel Gaussian log-density constants. Then $\mathbb E_q[(z_j-\mu_j)^2]=\sigma_j^2$ and $\mathbb E_q[z_j^2]=\mu_j^2+\sigma_j^2$.
Kingma & Welling (2014), §2, Appendix B
Derivation
ACT: latent-variable training The diagonal-Gaussian KL is explicit
$$\begin{aligned}\mathbb E_q\log(q_j/p_j)&=-\log\sigma_j-\frac12+\frac12(\mu_j^2+\sigma_j^2)\\D_{KL}(q\Vert p)&=\frac12\sum_j(\mu_j^2+\sigma_j^2-1-\log\sigma_j^2)\end{aligned}$$
$j$ independent Gaussian latent coordinate Rules: substitute the two moments; add coordinate log ratios.
Kingma & Welling (2014), §2, Appendix B
Worked example
ACT: latent-variable training A latent penalty calculated by hand
$$\mu=(1,0)^T,\quad\sigma^2=(1,4)^T$$
$\mu,\sigma^2$ two encoder-coordinate means and variances $D_{KL}=\frac12[1+(4-1-\log4)]=2-\log2\approx1.3069$.
Kingma & Welling (2014), §2, Appendix B · Course example
Think, explain, compare
ACT: latent-variable training Transfer: combine the two loss terms
Four action errors are $1,-1,0,2$. A one-dimensional encoder has $\mu=1$, $\sigma^2=1$. Use mean L1 and $\beta=2$. Find total loss.
A 1B 2✓ C 5
Solution 1 / 2
Reconstruct with the stated reduction $K d_a=4$ action coordinates; errors are prediction minus target.
$\mathcal L_{rec}=(|1|+|-1|+|0|+|2|)/4=1$. $D_{KL}=\frac12(1+1-1-\log1)=1/2$. Solution 2 / 2
Combine without changing the reduction $\beta=2$; mean reconstruction error equals 1.
$\mathcal L=1+2(1/2)=2$. B is correct. A drops the KL term; C uses summed rather than mean L1. Zhao et al. (2023), §IV-B–C, Appendix C · Course example
Definition
ACT: deployment The deployed ACT policy fixes the latent to zero
$$\widehat A_t=f_\theta(o_t,z=0)$$
$o_t$ current image and joint observations $z=0$ mean of the standard Gaussian prior The demonstration-conditioned encoder is discarded at inference.
Zhao et al. (2023), §IV-B–C, Appendix C
Think, explain, compare
ACT: deployment Training and inference use different information
Which input route matches the reported ACT inference procedure?
A Encode the unknown future target chunkB Current observations plus z=0✓ C Draw z from the demonstration posterior
Solution 1 / 1
Identify what is available now $o_t$: current observations; the future demonstration is unavailable.
Discard the demonstration-conditioned latent encoder. Set z to the standard-prior mean, zero, and run the decoder. B is correct. A and C require future target information. Zhao et al. (2023), §IV-B–C, Appendix C · Course example
Integrated synthesis
ACT: deployment One complete robot-policy calculation
Observe → encode → attend → predict a chunk → select same-time predictions → execute → observe again.
Replanning updates future chunks as fresh observations arrive. Zhao et al. (2023), §IV-B–C, Appendix C
Think, explain, compare
ACT: deployment Transfer: trace the live policy loop
At target time s, two chunks predict actions 2 and 4, oldest first. Use $m=\log2$. A fresh image arrives after execution. What is the executed target and the next computation?
A Execute 8/3; encode the fresh observation✓ B Execute 3; change βC Execute 4; reuse all old image features
Solution 1 / 2
Compute the same-time ensemble $r_0<r_1$; both predictions target s.
$w=(1,1/2)$; normalize to $(2/3,1/3)$. $\bar a_s=(2/3)2+(1/3)4=8/3$. Solution 2 / 2
Close the feedback loop $o_{s+1}$: fresh observations; inference latent remains zero.
Encode $o_{s+1}$ and predict $\widehat A_{s+1}=f_\theta(o_{s+1},0)$. A is correct. B changes the ensemble and a training coefficient. C discards the stated weighting and fresh evidence. Zhao et al. (2023), §IV-B–C, Appendix C · Course example
Think, explain, compare
ACT: deployment Transfer: a changed scene after prediction
A chunk predicts six actions. After one executed action, the object moves. Which response uses the new observation?
A Keep executing all old targets without recomputingB Encode the new observation and predict a new chunk✓ C Change the KL coefficient at inference
Solution 1 / 1
Locate the input that changed $o$: observation; learned model parameters remain fixed.
Object motion changes the current observation. Re-encoding it can change attention memory and future predicted actions. B uses that new evidence. A is open-loop; C changes a training parameter rather than reading the scene. Zhao et al. (2023), §IV-B–C, Appendix C · Course example
Bridge
Discrete tokens and continuous actions Byte-pair encoding is a text representation choice
Byte-pair encoding (BPE) builds a subword vocabulary by repeatedly merging frequent adjacent symbol pairs. A tokenizer maps text to discrete token IDs.
Text symbols BPE token IDs Embedding vectors Attention Sennrich, Haddow & Birch (2016), §3.2
Definition
Discrete tokens and continuous actions Top-p restricts a categorical sampling set
$$\sum_{j=1}^{k-1}p_{(j)}<\rho\le\sum_{j=1}^kp_{(j)}$$
$p_{(1)}\ge p_{(2)}\ge\cdots$ sorted category probabilities $\rho\in(0,1]$ top-p threshold Keep the smallest prefix reaching ρ, renormalize, then sample one category.
Holtzman et al. (2020), §3.1
Think, explain, compare
Discrete tokens and continuous actions Transfer: the retained set depends on mass
Sorted token probabilities are $(0.5,0.3,0.2)$ and top-p threshold $\rho=0.7$. Which distribution is sampled?
A $(0.5,0.3,0.2)$B $(0.625,0.375,0)$✓ C $(1,0,0)$
Solution 1 / 1
Find the smallest sufficient prefix $\rho=0.7$: cumulative-mass threshold.
$0.5<0.7$, while $0.5+0.3=0.8\ge0.7$. Retain two tokens and divide by $0.8$: $(5/8,3/8,0)$. B is correct. A omits truncation; C stops before reaching the threshold. Holtzman et al. (2020), §3.1 · Course example
Synthesis
Transformers and ACT One representation, several output models
Attention builds a conditional representation. ACT maps it to an action chunk. Diffusion and flow matching will use conditioning to guide iterative generation.
Observations Attention representation ACT action chunk Diffusion / flow output Zhao et al. (2023), §IV-B–C, Appendix C