How to calculate the gradient and Hessian for composite function?

Question

I have a function that can be written as follows: $$f(\vec{x},\vec{y}) = f_x(\vec{x}) + f_y(\vec{y}) + k\bigg(f_x(\vec{x})-f_y(\vec{y})\bigg)^2$$

I have to take the first derivative (gradient) and second derivative (hessian) of this function with respect to the combined $\vec{x}$ and $\vec{y}$. By combined, I mean a vector which the concatenation of both $\vec{x}$ and $\vec{y}$ like this: $$\vec{z} = (x_1,x_2,x_3,...,y_1,y_2,y_3,...)$$

So far I have been able to derive the analytic expression for the gradient $\nabla f$, which is simple because $f_x$ depends only on $x$ terms, and $f_y$ depends only on y terms. So, I can simply expand the square expression, collect the gradient terms, and concatenate. However, I am struggling to derive an expression for the Hessian matrix ($\nabla^2 f$) because of the cross-terms.

Note: The gradient and Hessian of each of the $f_x(\vec{x})$ and $f_y(\vec{y})$ with respect to their own vectors are possible to calculate (available). So, I have to write the gradient and Hessian of the composite function in those terms, as I need those for a program I am coding. Both $f_x$ and $f_y$ are functions that produce scalar numbers as ouptut. $k$ is a scalar constant.

I am a chemist, so I am not very familiar with linear algebra. If the Hessian (or gradient) can be calculated quickly with a matrix operation then it would be easier to code for. Any help is appreciated.

Answer 1

Let $h(x,y) = f_x(\vec{x})+f_y(\vec{y})+k\Big(f_x(\vec{x})-f_y(\vec{y})\Big)^2$. Then

$$\nabla_{\vec{x},\vec{y}}h(\vec{x},\vec{y})=\begin{bmatrix}\nabla_\vec{x}\\\nabla_\vec{y}\end{bmatrix}h(\vec{x},\vec{y})\\ = \begin{bmatrix}\nabla_\vec{x}f_x(\vec{x}) + k\nabla_\vec{x}f_x(\vec{x})^2-2kf_y(\vec{y})\nabla_\vec{x}f_x(\vec{x}) \\ \nabla_\vec{y}f_y(\vec{y}) + k\nabla_\vec{y}f_y(\vec{y})^2-2kf_x(\vec{x})\nabla_\vec{y}f_y(\vec{y})\end{bmatrix}$$ and $$H_{\vec{x},\vec{y}}h(\vec{x},\vec{y})=\begin{bmatrix}\partial^2_{\vec{x},\vec{x}} & \partial^2_{\vec{x},\vec{y}} \\ \partial^2_{\vec{y},\vec{x}} & \partial^2_{\vec{y},\vec{y}} \end{bmatrix}h(\vec{x},\vec{y})$$ $$=\begin{bmatrix} H_\vec{x}f_x(\vec{x}) + kH_\vec{x}f_x(\vec{x})^2-2kf_y(\vec{y})H_\vec{x}f_x(\vec{x}) & -2k\nabla f_x^T \nabla f_y \\ -2k\nabla f_y^T \nabla f_x & H_\vec{y}f_y(\vec{y}) + kH_\vec{y}f_y(\vec{y})^2-2kf_x(\vec{x})H_\vec{y}f_y(\vec{y}) \end{bmatrix}$$

I've written these in block matrix notation. If you can compute the regular gradients and hessians of $f_i$ and $f_i^2$, then these formulae should get you to the gradient and hessian of $h$

Answer 2

For the cross-term of the Hessian, you can write \begin{eqnarray} \mathbf{g}_x \equiv \frac{\partial f}{\partial \mathbf{x}} &=& -2 k f_y(\mathbf{y}) \frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}} \\ d\mathbf{g}_x &=& -2 k f_y(\mathbf{y}) \frac{\partial^2 f_x(\mathbf{x})}{\partial \mathbf{x}^2} d\mathbf{x} -2 k \frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}} \left( \frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}} \right)^T d\mathbf{y} \end{eqnarray} and similarly \begin{eqnarray} \mathbf{g}_y \equiv \frac{\partial f}{\partial \mathbf{y}} &=& -2 k f_x(\mathbf{x}) \frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}} \\ d\mathbf{g}_y &=& -2 k f_x(\mathbf{x}) \frac{\partial^2 f_y(\mathbf{y})}{\partial \mathbf{y}^2} d\mathbf{y} -2 k \frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}} \left( \frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}} \right)^T d\mathbf{x} \end{eqnarray} The Hessian is easily found by regrouping terms \begin{equation} -2k \begin{bmatrix} f_y(\mathbf{y}) \frac{\partial^2 f_x(\mathbf{x})}{\partial \mathbf{x}^2} & \frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}} \left( \frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}} \right)^T \\ \frac{\partial f_y(\mathbf{y})}{\partial \mathbf{y}} \left( \frac{\partial f_x(\mathbf{x})}{\partial \mathbf{x}} \right)^T & f_x(\mathbf{x}) \frac{\partial^2 f_y(\mathbf{y})}{\partial \mathbf{y}^2} \end{bmatrix} \end{equation}

Answer 3

$ \def\BR#1{\Big(#1\Big)} \def\LR#1{\left(#1\right)} \def\op#1{\operatorname{#1}} \def\trace#1{\op{Tr}\LR{#1}} \def\qiq{\quad\implies\quad} \def\qif{\quad\iff\quad} \def\a{\alpha}\def\b{\beta}\def\T{{\large\tt0}} \def\l{\lambda}\def\t{\lambda}\def\p{\partial} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\hess#1#2{\frac{\p^2 #1}{\p #2^2}} \def\m#1{\left[\begin{array}{r}#1\end{array}\right]} \def\mr#1{\left[\begin{array}{l}#1\end{array}\right]} \def\mm#1{\left[\begin{array}{c|c}#1\end{array}\right]} \def\c#1{\color{red}{#1}} \def\CLR#1{\c{\LR{#1}}} \def\fracLR#1#2{\LR{\frac{#1}{#2}}} \def\o{{\tt1}} \def\bbR#1{{\mathbb R}^{#1}} $The inverse of concatenating matrix and vector variables is to partition them into blocks.

Towards that end, define block-matrix analogs of the cartesian basis vectors $\LR{e_k\in\bbR 2}$ $$\eqalign{ &E_1 = \mr{I_n\\ \T^T},\quad &E_2 = \mr{\T\\I_m} &\qiq &z = E_1x + E_2y = \m{x\\y} \\ &E_1^TE_1 = I_n,&E_1^TE_2 = \T &\qiq &x=E_1^Tz \\ &E_2^TE_1 = \T^T,&E_2^TE_2 = I_m &\qiq &y=E_2^Tz \\ &E_1E_1^T + E_2E_2^T &= I_{(m+n)} \\ &\T \in\bbR{n\times m} &x\in\bbR{n\times\o} \qquad &y\in\bbR{m\times\o}\qquad &z\in\bbR{(m+n)\times\o} \\ }$$ and give the known function/gradient/Hessian variables single-letter names $$\eqalign{ \a &= \a(x),\qquad &a = a(x) = \grad{\a}{x},\qquad &A = A(x) = \grad{a}{x} = \hess{\a}{x} \\ \b &= \b(y),\qquad &b = b(y) = \grad{\b}{y},\qquad &B = B(y) = \grad{b}{y} = \hess{\b}{y} \\ }$$ Rewrite the function and calculate its gradient $$\eqalign{ f &= \a + \b + k\LR{\a-\b}^2 \\ df &= d\a + d\b + \c{2k\LR{\a-\b}}\LR{d\a-d\b} \\ &= d\a + d\b + \c{\t}\LR{d\a-d\b} \\ &= \LR{\o+\t}\c{d\a} + \LR{\o-\t}\c{d\b} \\ &= \LR{\o+\t}\c{a:dx} + \LR{\o-\t}\c{b:dy} \\ &= \LR{\o+\t}a:E_1^Tdz + \LR{\o-\t}b:E_2^Tdz \\ &= \LR{\o+\t}E_1a:dz + \LR{\o-\t}E_2b:dz \\ &= \BR{\LR{\o+\t}E_1a + \LR{\o-\t}E_2b\,}:dz \\ \grad{f}{z} &= \BR{\LR{\o+\t}E_1a + \LR{\o-\t}E_2b\,} \;=\; g\;\Big({\rm gradient}\Big) \\ g &= \mm{\LR{\o+\t}a \\ \LR{\o-\t}b} \\ }$$ Before tackling the Hessian, here's the gradient of the $\l$ variable introduced above $$\eqalign{ c &= \LR{E_1a - E_2b} = \m{a\\-b} \\ \t &= 2k\LR{\a-\b} \\ d\t &= 2k\LR{d\a-d\b} \\ &= 2k\LR{E_1a-E_2b}:dz \\ &= 2kc^Tdz \\ }$$ The Hessian is simply the gradient of the gradient, therefore $$\eqalign{ dg &= \LR{\o+\t}E_1\,\c{da} + \LR{\o-\t}E_2\,\c{db} + \LR{E_1a-E_2b}\c{d\t} \\ &= \LR{\o+\t}E_1\,\c{A\,dx} + \LR{\o-\t}E_2\,\c{B\,dy} + \LR{E_1a-E_2b}\c{2kc^Tdz} \\ &= \LR{\o+\t}E_1AE_1^Tdz + \LR{\o-\t}E_2BE_2^Tdz + 2kcc^Tdz \\ &= \BR{\LR{\o+\t}E_1AE_1^T + \LR{\o-\t}E_2BE_2^T + 2kcc^T}\,dz \\ \grad{g}{z} &= \BR{\LR{\o+\t}E_1AE_1^T + \LR{\o-\t}E_2BE_2^T + 2kcc^T} \;=\; H\;\Big({\rm Hessian}\Big) \\ H &= \mm{ 2kaa^T+\LR{\o+\t}A & -2kab^T \\\hline -2kba^T & 2kbb^T+\LR{\o-\t}B \\ } \;=\; \hess{\!f}{z} \\ \\ }$$

---

In the above $(:)$ denotes the Frobenius product, which is extremely useful in Matrix Calculus $$\eqalign{ A:B &= \sum_{i=1}^m\sum_{j=1}^n A_{ij}B_{ij} \;=\; \trace{A^TB} \\ A:A &= \|A\|^2_{\c F}\qquad \Big({\rm \c{Frobenius}\:norm}\Big) \\ }$$ This is also called the double-dot or double contraction product.
When applied to vectors $(n={\tt1})$ it reduces to the standard dot product.

The properties of the underlying trace function allow the terms in a Frobenius product to be rearranged in many interesting ways, e.g. $$\eqalign{ A:B &= B:A \\ A:B &= A^T:B^T \\ C:\LR{AB} &= \LR{CB^T}:A &= \LR{A^TC}:B \\ }$$