MetaOptimize: A Framework for Optimizing Step Sizes and Other Meta-parameters — part 7/10

Cleaned-text code points [80289, 92232). Bibliography omitted here; source.md includes it.

By taking the derivative of ( 52 ) with respect to 𝒉 t \boldsymbol{h}_{t} , we obtain d ​ 𝒉 t + 1 d ​ 𝒉 t = γ ⁡ [ I − [ 𝜶 t ] ​ ∇ 2 f t ​ ( 𝒘 t ) 0 0 0 0 I − [ 𝜶 t ] ​ ∇ 2 f t ​ ( 𝒘 t ) 0 0 0 0 ⋱ 0 0 0 0 I − [ 𝜶 t ] ​ ∇ 2 f t ​ ( 𝒘 t ) ] ​ ← 1 ​ st ← 2 ​ nd ⋮ ← m ​ th . \frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{h}_{t}}=\gamma\,\left[\begin{array}[]{c|c|c|c}\begin{array}[]{c}I-\left[\boldsymbol{\alpha}_{t}\right]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t})\\ \\ \end{array}&0&0&0\\ \hline\cr 0&\begin{array}[]{c}\\ I-\left[\boldsymbol{\alpha}_{t}\right]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t})\\ \\ \end{array}&0&0\\ \hline\cr 0&0&\begin{array}[]{c}\\ \ddots\\ \\ \end{array}&0\\ \hline\cr 0&0&0&\begin{array}[]{c}\\ I-\left[\boldsymbol{\alpha}_{t}\right]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t})\end{array}\end{array}\right]\,\,\begin{array}[]{l}\begin{array}[]{c}{\color[rgb]{0.75,0.75,0.75}\leftarrow 1\textrm{st}}\\ \\ \end{array}\\ \begin{array}[]{c}\\ {\color[rgb]{0.75,0.75,0.75}\leftarrow 2\textrm{nd}}\\ \\ \end{array}\\ \begin{array}[]{c}\\ \,\,{\color[rgb]{0.75,0.75,0.75}\vdots}\\ \\ \end{array}\\ \begin{array}[]{c}\\ {\color[rgb]{0.75,0.75,0.75}\leftarrow m\textrm{th}}\end{array}\end{array}. (54) In the above equation, note that d ​ B t / d ​ 𝒉 t = 0 \mathrm{d}\,B_{t}/\mathrm{d}\,\boldsymbol{h}_{t}=0 due to ( 47 ). Let β t ​ [ i ] \beta_{t}[i] and w t ​ [ j ] w_{t}[j] denote the i i th and j j th entries of 𝜷 t \boldsymbol{\beta}_{t} and 𝒘 t \boldsymbol{w}_{t} , for i = 1 , … , m i=1,\ldots,m and j = 1 , … , n j=1,\ldots,n , respectively. It then follows from ( 52 ) and ( 47 ) that d ​ 𝒉 t + 1 d ​ 𝜷 t = − γ ⁡ [ [ d ​ 𝜶 t d ​ β t ​ [ 1 ] ] ​ ∇ 2 f t ​ ℋ t [ 1 ] + [ ∇ f t ] ​ ∂ σ ′ ​ ( 𝜷 t ) ∂ β t ​ [ 1 ] ​ B t [ 1 ] ⋯ [ d ​ 𝜶 t d ​ β t ​ [ m ] ] ​ ∇ 2 f t ​ ℋ t [ 1 ] + [ ∇ f t ] ​ ∂ σ ′ ​ ( 𝜷 t ) ∂ β t ​ [ m ] ​ B t [ 1 ] ⋮ ⋱ ⋮ [ d ​ 𝜶 t d ​ β t ​ [ 1 ] ] ​ ∇ 2 f t ​ ℋ t [ m ] + [ ∇ f t ] ​ ∂ σ ′ ​ ( 𝜷 t ) ∂ β t ​ [ 1 ] ​ B t [ m ] ⋯ [ d ​ 𝜶 t d ​ β t ​ [ m ] ] ​ ∇ 2 f t ​ ℋ t [ m ] + [ ∇ f t ] ​ ∂ σ ′ ​ ( 𝜷 t ) ∂ β t ​ [ m ] ​ B t [ m ] ] , \frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}=-\gamma\left[\begin{array}[]{c|c|c}\hskip-7.0pt\begin{array}[]{c}\\ {\Big[\frac{\mathrm{d}\,\boldsymbol{\alpha}_{t}}{\mathrm{d}\,\beta_{t}[1]}\Big]\nabla^{2}f_{t}\,\mathcal{H}_{t}^{[1]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\,\sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[1]}\,B_{t}^{[1]}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\cdots}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\Big[\frac{\mathrm{d}\,\boldsymbol{\alpha}_{t}}{\mathrm{d}\,\beta_{t}[m]}\Big]\nabla^{2}f_{t}\,\mathcal{H}_{t}^{[1]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\,\sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[{m}]}\,B_{t}^{[1]}}\\ \\ \end{array}\\ \hline\cr\hskip-7.0pt\begin{array}[]{c}\\ {\vdots}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\ddots}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\vdots}\\ \\ \end{array}\\ \hline\cr\hskip-7.0pt\begin{array}[]{c}\\ {\Big[\frac{\mathrm{d}\,\boldsymbol{\alpha}_{t}}{\mathrm{d}\,\beta_{t}[1]}\Big]\nabla^{2}f_{t}\,\mathcal{H}_{t}^{[m]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\,\sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[1]}\,B_{t}^{[m]}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\cdots}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\Big[\frac{\mathrm{d}\,\boldsymbol{\alpha}_{t}}{\mathrm{d}\,\beta_{t}[m]}\Big]\nabla^{2}f_{t}\,\mathcal{H}_{t}^{[m]}\,+\,\left[\nabla f_{t}\right]\,\frac{\partial\,\sigma^{\prime}(\boldsymbol{\beta}_{t})}{\partial\,\beta_{t}[{m}]}\,B_{t}^{[m]}}\\ \\ \end{array}\end{array}\right], (55) where ∂ ∂ β \frac{\partial}{\partial\beta} stands for the entry-wise partial derivative of a matrix with respect to a scalar variable β \beta . In the same vein, ( 53 ) and ( 47 ) imply that d ​ 𝒉 t + 1 d ​ 𝒘 t = − γ ⁡ [ [ 𝜶 t ] ​ d ⁡ ( ∇ 2 f t ​ ( 𝒘 t ) ​ ℋ t [ 1 ] ) d ​ 𝒘 t + [ σ ′ ​ ( 𝜷 t ) ​ B t [ 1 ] ] ​ ∇ 2 f t ​ ( 𝒘 t ) ⋮ [ 𝜶 t ] ​ d ⁡ ( ∇ 2 f t ​ ( 𝒘 t ) ​ ℋ t [ m ] ) d ​ 𝒘 t + [ σ ′ ​ ( 𝜷 t ) ​ B t [ m ] ] ​ ∇ 2 f t ​ ( 𝒘 t ) ] . \frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{w}_{t}}=-\gamma\left[\begin{array}[]{c}\hskip-7.0pt\begin{array}[]{c}\\ {\left[\boldsymbol{\alpha}_{t}\right]\frac{\mathrm{d}\,(\nabla^{2}f_{t}(\boldsymbol{w}_{t})\,\mathcal{H}_{t}^{[1]})}{\mathrm{d}\,\boldsymbol{w}_{t}}\,+\,\big[\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[1]}\big]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t})}\\ \\ \end{array}\\ \hline\cr\hskip-7.0pt\begin{array}[]{c}\\ {\vdots}\\ \\ \end{array}\\ \hline\cr\hskip-7.0pt\begin{array}[]{c}\\ {\left[\boldsymbol{\alpha}_{t}\right]\frac{\mathrm{d}\,(\nabla^{2}f_{t}(\boldsymbol{w}_{t})\,\mathcal{H}_{t}^{[{m}]})}{\mathrm{d}\,\boldsymbol{w}_{t}}\,+\,\big[\sigma^{\prime}(\boldsymbol{\beta}_{t})\,B_{t}^{[m]}\big]\,\nabla^{2}f_{t}(\boldsymbol{w}_{t})}\\ \\ \end{array}\end{array}\right]. (56) Finally, G t base G^{\textrm{base}}_{t} is obtained by plugging ( 54 ), ( 55 ), and ( 56 ) into ( 50 ).

In the special case that β \boldsymbol{\beta} is a scalar (equivalently m = 1 m=1 ), and furthermore α = σ ⁡ ( β ) = e β \alpha=\sigma(\beta)=e^{\beta} , matrix G t base G^{\textrm{base}}_{t} would be simplified to G t base (scalar) = [ 1 − η ​ 𝒉 t T ​ ∇ 2 f t ​ ( 𝒘 t ) − η ∇ f t ( 𝒘 t ) T − α ∇ f t ( 𝒘 t ) I − α ​ ∇ 2 f t ​ ( 𝒘 t ) 0 − γ α ∇ 2 f t ( 𝒘 t ) 𝒉 t − B t α ∇ f t ( 𝒘 t ) − γ ​ α ​ d ⁡ ( ∇ 2 f t ​ ( 𝒘 t ) ​ 𝒉 t ) d ​ 𝒘 t − B t ​ α ​ ∇ 2 f t ​ ( 𝒘 t ) γ ⁡ ( I − α ​ ∇ 2 f t ​ ( 𝒘 t ) ) ] . G_{t}^{\textrm{base (scalar)}}=\left[\begin{array}[]{cccc}\vskip 6.0pt1&-\eta\,\boldsymbol{h}_{t}^{T}\nabla^{2}f_{t}(\boldsymbol{w}_{t})&-\eta\,\nabla f_{t}(\boldsymbol{w}_{t})^{T}\\ \vskip 6.0pt-\alpha\nabla f_{t}(\boldsymbol{w}_{t})&I-\alpha\nabla^{2}f_{t}(\boldsymbol{w}_{t})&0\\ -\gamma\alpha\nabla^{2}f_{t}(\boldsymbol{w}_{t})\boldsymbol{h}_{t}-B_{t}\,\alpha\nabla f_{t}(\boldsymbol{w}_{t})&\quad-\gamma\alpha\frac{\mathrm{d}\,\left(\nabla^{2}f_{t}(\boldsymbol{w}_{t})\boldsymbol{h}_{t}\right)}{d\boldsymbol{w}_{t}}-B_{t}\,\alpha\nabla^{2}f_{t}(\boldsymbol{w}_{t})&\gamma\big(I-\alpha\nabla^{2}f_{t}(\boldsymbol{w}_{t})\big)\end{array}\right].

#### A.3.1 Base AdamW

The base update according to the AdamW algorithm ( Loizou et al., 2021 ) is as follows, 𝒎 t + 1 = ρ ​ 𝒎 t + ∇ f t ​ ( 𝒘 t ) , 𝒗 t + 1 = λ 𝒗 t + ∇ f t ( 𝒘 t ) 2 , μ t = ( 1 − ρ 1 − ρ t ) / 1 − λ 1 − λ t , 𝒘 t + 1 = 𝒘 t − 𝜶 t ​ μ t ​ 𝒎 t 𝒗 t − κ ​ 𝜶 t ​ 𝒘 t , \begin{split}\boldsymbol{m}_{t+1}&=\rho\,\boldsymbol{m}_{t}\,+\,\nabla f_{t}(\boldsymbol{w}_{t}),\\ \boldsymbol{v}_{t+1}&=\lambda\,\boldsymbol{v}_{t}\,+\,\nabla f_{t}(\boldsymbol{w}_{t})^{2},\\ \mu_{t}&=\left(\frac{1-\rho}{{\color[rgb]{0.75,0.75,0.75}1-\rho^{t}}}\right)\,/\,\sqrt{\frac{1-\lambda}{1-\lambda^{t}}},\\ \boldsymbol{w}_{t+1}&=\boldsymbol{w}_{t}\,-\,\boldsymbol{\alpha}_{t}\mu_{t}\frac{\boldsymbol{m}_{t}}{\sqrt{\boldsymbol{v}_{t}}}-\kappa\boldsymbol{\alpha}_{t}\boldsymbol{w}_{t},\end{split} (57) where 𝒎 t \boldsymbol{m}_{t} is the momentum vector, 𝒗 t \boldsymbol{v}_{t} is the trace of gradient square used for normalization, and κ > 0 \kappa>0 is a weight-decay parameter. Therefore the base algorithm needs to keep track of 𝒘 t , 𝒎 t , 𝒗 t \boldsymbol{w}_{t},\boldsymbol{m}_{t},\boldsymbol{v}_{t} , i.e., 𝒙 t = [ 𝒘 t 𝒎 t 𝒗 t ] . \boldsymbol{x}_{t}=\left[\begin{array}[]{c}\boldsymbol{w}_{t}\\ \hfill\boldsymbol{m}_{t}\\ \boldsymbol{v}_{t}\end{array}\right]. (58) It then follows from ( 44 ) that G t base = [ d ​ 𝒙 t + 1 d ​ 𝜷 t 0 d ​ 𝒙 t + 1 d ​ 𝒙 t 0 d ​ 𝒉 t + 1 d ​ 𝜷 t 0 d ​ 𝒉 t + 1 d ​ 𝒙 t d ​ 𝒉 t + 1 d ​ 𝒉 t ] = [ d ​ 𝒘 t + 1 d ​ 𝜷 t 0 d ​ 𝒘 t + 1 d ​ 𝒘 t d ​ 𝒘 t + 1 d ​ 𝒎 t d ​ 𝒘 t + 1 d ​ 𝒗 t 0 d ​ 𝒎 t + 1 d ​ 𝜷 t 0 d ​ 𝒎 t + 1 d ​ 𝒘 t d ​ 𝒎 t + 1 d ​ 𝒎 t d ​ 𝒎 t + 1 d ​ 𝒗 t 0 d ​ 𝒗 t + 1 d ​ 𝜷 t 0 d ​ 𝒗 t + 1 d ​ 𝒘 t d ​ 𝒗 t + 1 d ​ 𝒎 t d ​ 𝒗 t + 1 d ​ 𝒗 t 0 d ​ 𝒉 t + 1 d ​ 𝜷 t 0 d ​ 𝒉 t + 1 d ​ 𝒘 t d ​ 𝒉 t + 1 d ​ 𝒎 t d ​ 𝒉 t + 1 d ​ 𝒗 t d ​ 𝒉 t + 1 d ​ 𝒉 t ] = [ − μ t ​ [ 𝒎 t 𝒗 t + κ ​ 𝒘 t ] ​ σ ′ ​ ( 𝜷 t ) 0 I − κ ⁡ [ 𝜶 t ] − μ t ​ [ 𝜶 t 𝒗 t ] μ t 2 ​ [ 𝜶 t ​ 𝒎 t 𝒗 t 1.5 ] 0 0 0 ∇ 2 f t ρ ​ I 0 0 0 0 2 ​ [ ∇ f t ] ​ ∇ 2 f t 0 λ ​ I 0 d ​ 𝒉 t + 1 d ​ 𝜷 t 0 d ​ 𝒉 t + 1 d ​ 𝒘 t d ​ 𝒉 t + 1 d ​ 𝒎 t d ​ 𝒉 t + 1 d ​ 𝒗 t d ​ 𝒉 t + 1 d ​ 𝒉 t ] \begin{split}G^{\textrm{base}}_{t}&=\left[\begin{array}[]{cccc}\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{x}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}}\\ \\ \end{array}&0&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{x}_{t+1}}{\mathrm{d}\,\boldsymbol{x}_{t}}}\\ \\ \end{array}&0\\ \hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}}\\ \\ \end{array}&0&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{x}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{h}_{t}}}\\ \\ \end{array}\end{array}\right]\\ &=\left[\begin{array}[]{cc|ccc|c}\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{w}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}}\\ \\ \end{array}&0&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{w}_{t+1}}{\mathrm{d}\,\boldsymbol{w}_{t}}}\\ \\ \end{array}&\frac{\mathrm{d}\,\boldsymbol{w}_{t+1}}{\mathrm{d}\,\boldsymbol{m}_{t}}&\frac{\mathrm{d}\,\boldsymbol{w}_{t+1}}{\mathrm{d}\,\boldsymbol{v}_{t}}&0\\ \frac{\mathrm{d}\,\boldsymbol{m}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}&0&\frac{\mathrm{d}\,\boldsymbol{m}_{t+1}}{\mathrm{d}\,\boldsymbol{w}_{t}}&\frac{\mathrm{d}\,\boldsymbol{m}_{t+1}}{\mathrm{d}\,\boldsymbol{m}_{t}}&\frac{\mathrm{d}\,\boldsymbol{m}_{t+1}}{\mathrm{d}\,\boldsymbol{v}_{t}}&0\\ \frac{\mathrm{d}\,\boldsymbol{v}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}&0&\frac{\mathrm{d}\,\boldsymbol{v}_{t+1}}{\mathrm{d}\,\boldsymbol{w}_{t}}&\frac{\mathrm{d}\,\boldsymbol{v}_{t+1}}{\mathrm{d}\,\boldsymbol{m}_{t}}&\frac{\mathrm{d}\,\boldsymbol{v}_{t+1}}{\mathrm{d}\,\boldsymbol{v}_{t}}&0\\ \hline\cr\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}}\\ \\ \end{array}&0&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{w}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{m}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{v}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{h}_{t}}}\\ \\ \end{array}\end{array}\right]\\ &=\left[\begin{array}[]{cc|ccc|c}-\mu_{t}\Big[\frac{\boldsymbol{m}_{t}}{\sqrt{\boldsymbol{v}_{t}}}+\kappa\boldsymbol{w}_{t}\Big]\sigma^{\prime}(\boldsymbol{\beta}_{t})&0&I-\kappa\left[\boldsymbol{\alpha}_{t}\right]&-\mu_{t}\Big[\frac{\boldsymbol{\alpha}_{t}}{\sqrt{\boldsymbol{v}_{t}}}\Big]&\frac{\mu_{t}}{2}\Big[\frac{\boldsymbol{\alpha}_{t}\boldsymbol{m}_{t}}{\boldsymbol{v}_{t}^{1.5}}\Big]&0\\ 0&0&\nabla^{2}f_{t}&\rho I&0&0\\ 0&0&2\left[\nabla f_{t}\right]\,\nabla^{2}f_{t}&0&\lambda I&0\\ \hline\cr\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{\beta}_{t}}}\\ \\ \end{array}&0&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{w}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{m}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{v}_{t}}}\\ \\ \end{array}&\hskip-7.0pt\begin{array}[]{c}\\ {\frac{\mathrm{d}\,\boldsymbol{h}_{t+1}}{\mathrm{d}\,\boldsymbol{h}_{t}}}\\ \\ \end{array}\end{array}\right]\end{split} (59) where the last equality follows from simple derivative computations in ( 57 ).