Derive: Getting Gradients From a Computer

(a) For the call count, the claim is $T(L)=(L/2){\log }_{2}L$. At $L=2$ the right side is $1\cdot 1=1$, which matches the base case. Assume the claim at $L/2$. Then $$T(L)=2\left(\frac{L}{4}{\log }_2\frac{L}{2}\right)+\frac{L}{2}=\frac{L}{2}\left({\log }_{2}L-1\right)+\frac{L}{2}=\frac{L}{2}{\log }_{2}L.$$ The first term uses $2\cdot (L/4)=L/2$, and the second cancels the $L/2$ that the induction step adds. For the slot count, the claim is $S(L)={\log }_{2}L$. At $L=2$ it gives 1, matching the base case. Assuming it at $L/2$ gives $S(L)={\log }_2(L/2)+1={\log }_{2}L-1+1={\log }_{2}L$.

(b) Divide the call count by the store-all call count: $$R(L)=\frac{T(L)}{L-1}=\frac{L{\log }_{2}L}{2(L-1)}.$$ Divide the store-all slot count by the recursive-halving slot count: $$M(L)=\frac{L}{S(L)}=\frac{L}{{\log }_{2}L}.$$

(c) Write $L/(L-1)=1+1/(L-1)$. Then $$R(L)=\frac{{\log }_{2}L}{2}\left(1+\frac{1}{L-1}\right),$$ so the ratio is above $({\log }_{2}L)/2$ by a term that shrinks as $L$ grows. The values are:

| $L$ | $T(L)$ | $S(L)$ | $L-1$ | $R(L)$ | $M(L)$ | |---|---|---|---|---|---| | 8 | 12 | 3 | 7 | 1.71 | 2.67 | | 64 | 192 | 6 | 63 | 3.05 | 10.67 | | 1024 | 5120 | 10 | 1023 | 5.00 | 102.4 |

At $L=8$ the schedule does 12 calls in 3 slots against 7 calls in 8 slots. At $L=1024$ it does 5120 calls in 10 slots against 1023 calls in 1024 slots. The work ratio grows like half the logarithm of the depth, while the memory ratio grows like the depth divided by its logarithm.

(d) Section 18.4 establishes that a trace needs one slot per node the backward sweep will read, which prices storing everything at $L$ slots. Section 18.5 introduces the schedule and the recomputation that buy the memory back. The recurrence assumes that a block is deterministic: a recomputed block must return the value it returned the first time. A block that draws a fresh random mask or increments a global counter breaks that condition, and the gradient becomes wrong with no diagnostic. The count also assumes that blocks cost the same, so it measures calls rather than seconds. Both assumptions belong to the code, not to the mathematics.

(a) The three outputs are $F_1={\theta }_1{\theta }_2$, $F_2={\theta }_1+{\theta }_2$ and $F_3={\theta }_1^2$. Their partial derivatives at $(3,2)$ are $$J=\left[\begin{array}{cc}\partial F_1/\partial {\theta }_1& \partial F_1/\partial {\theta }_2\\ \partial F_2/\partial {\theta }_1& \partial F_2/\partial {\theta }_2\\ \partial F_3/\partial {\theta }_1& \partial F_3/\partial {\theta }_2\end{array}\right]=\left[\begin{array}{cc}{\theta }_2& {\theta }_1\\ 1& 1\\ 2{\theta }_1& 0\end{array}\right]=\left[\begin{array}{cc}2& 3\\ 1& 1\\ 6& 0\end{array}\right].$$

(b) Forward mode multiplies $J$ on the right by a direction. The two coordinate directions give the two columns: $$J\left[\begin{array}{c}1\\ 0\end{array}\right]=\left[\begin{array}{c}2\\ 1\\ 6\end{array}\right],J\left[\begin{array}{c}0\\ 1\end{array}\right]=\left[\begin{array}{c}3\\ 1\\ 0\end{array}\right].$$ Two passes return the whole matrix, because there are two inputs.

(c) Reverse mode multiplies $J$ on the left by a seed. The three coordinate seeds give the three rows: $$\left[\begin{array}{ccc}1& 0& 0\end{array}\right]J=\left[\begin{array}{cc}2& 3\end{array}\right],\left[\begin{array}{ccc}0& 1& 0\end{array}\right]J=\left[\begin{array}{cc}1& 1\end{array}\right],\left[\begin{array}{ccc}0& 0& 1\end{array}\right]J=\left[\begin{array}{cc}6& 0\end{array}\right].$$ Three passes return the whole matrix, because there are three outputs.

(d) The claim that reverse mode is the cheap direction holds for a scalar loss, where $m=1$. It does not hold in general. Forward mode costs $n$ passes for the full Jacobian, and reverse mode costs $m$ passes. Reverse mode is cheaper exactly when $m<n$; forward mode is cheaper when $n<m$; and the two tie when $m=n$. Here $m=3$ and $n=2$, so forward mode needs 2 passes and reverse mode needs 3. A function with many outputs and few inputs, such as a layer that returns a whole feature map, is the case where forward mode wins. Section 18.2 derives the forward count, and section 18.3 derives the reverse count; this example puts the two counts side by side and shows that the shape of the Jacobian decides.

(a) Expand the two cubes: $$(2+\Delta {)}^3=8+12\Delta +6{\Delta }^2+{\Delta }^3,(2-\Delta {)}^3=8-12\Delta +6{\Delta }^2-{\Delta }^3.$$ Subtracting gives $24\Delta +2{\Delta }^3$. Dividing by $2\Delta$ gives $12+{\Delta }^2$. The exact slope is 12, so the truncation error is ${\Delta }^2$.

At $x=2$ the function value is 8, so each evaluation carries absolute error at most $8\epsilon$. The subtraction can bring in $2\cdot 8\epsilon =16\epsilon$, and dividing by $2\Delta$ gives the roundoff bound $16\epsilon /(2\Delta )=8\epsilon /\Delta$. The two errors add in the model: $$E(\Delta )={\Delta }^2+\frac{8\epsilon }{\Delta }.$$

(b) Differentiate with respect to $\Delta$: $$E^{\prime }(\Delta )=2\Delta -\frac{8\epsilon }{{\Delta }^2}.$$ Setting $E^{\prime }=0$ gives $2{\Delta }^3=8\epsilon$, or ${\Delta }^3=4\epsilon$. The second derivative $2+16\epsilon /{\Delta }^3$ is positive, so this stationary point is a minimum. Therefore $${\Delta }^{\star }=(4\epsilon {)}^{1/3}.$$ At this step ${\Delta }^3=4\epsilon$, so $8\epsilon /\Delta =2{\Delta }^2$. The two terms are in the ratio $1:2$, and $$E({\Delta }^{\star })={\Delta }^2+2{\Delta }^2=3{\Delta }^2=3(4\epsilon {)}^{2/3}.$$

(c) With $\epsilon =2^{-53}$ the combination $4\epsilon =2^{-51}$, so $${\Delta }^{\star }=2^{-17}=7.62939453125\times 10^{-6},$$ which is $7.63\times 10^{-6}$ to three significant figures. The minimum error is $$E({\Delta }^{\star })=3\cdot 2^{-34}=1.7462298274\times 10^{-10},$$ which is $1.75\times 10^{-10}$. The ratio is $$\frac{E({\Delta }^{\star })}{\epsilon }=3\cdot 2^{19}=1572864,$$ which is $1.57\times 10^6$.

(d) The best centred quotient is about $1.57$ million times less accurate than one rounded evaluation of $f$. That gap is the edge case: the numeric route is exact for nothing, and its best accuracy is set by the arithmetic rather than by the function. A smaller step makes the subtraction worse, because the roundoff term $8\epsilon /\Delta$ grows; a larger step makes the truncation worse, because ${\Delta }^2$ grows. At ${\Delta }^{\star }$ the two are balanced. A gradient check at a fixed step therefore reports agreement to the digits that step can reach, and those digits are far fewer than the machine carries. Section 18.1 calls that agreement evidence rather than proof; this derivation gives the size of the evidence.