\section{Introduction}
\label{intr}

One of many achievements of the Density Functional Theory \cite{pr_136_B864,pr_140_A1133} (DFT) is its ability to provide accurate total energy of an interacting many-electron system as a function of external fields such as an electrostatic field generated by nuclei in a solid. Naturally, it is important to be able to find an arrangement of nuclei in a solid (crystal structure) which corresponds to a minimal total energy of the whole system (electrons plus nuclei). Another important field of interest is a studying of the response of a solid when its nuclei are pushed slightly from their equilibrium positions. In both situations, the ability to evaluate accurate derivatives of the total energy with respect to the atomic positions (forces) represents an important tool. In the context of an equilibrium structure search (geometry optimization), availability of the forces greatly helps, providing the directions where atoms should be moved in order to reach their equilibrium positions. In the context of small deviations from the equilibrium (phonons), availability of the forces allows one to use the finite displacement method \cite{epl_32_729,prb_64_045123,prb_78_134106} to calculate phonon frequencies without employing more technically involved linear response approach \cite{rmp_73_515}.

In a family of methods based on the Augmented Plane Waves \cite{pr_51_846} (APW) basis set, such as the Linearized Augmented Plane Waves (LAPW, \cite{prb_12_3060}), the formulation of how one can evaluate atomic forces was given by Yu et al. \cite{prb_43_6411}. The work by Yu et al. demonstrated that breaking the space into non-overlapping muffin-tin (MT) spheres and the so-called interstitial region (IR), which is an important attribute of the APW-family of methods, leads to an additional contribution to the atomic forces, the Pulay term \cite{mp_17_197}. Since then, there were a few enhancements introduced, such as inclusion of additional surface terms when one uses basis sets with discontinuities across the MT boundaries (for instance in APW+lo basis set, \cite{prb_64_195134}). Also, Kl\"{u}ppelberg et al. \cite{prb_91_035105} presented a refinement of the approach by carefully taking into account the tails of the high-energy core states as well as small discontinuities in wafe functions, density, and potential at the MT spheres. Independently, Soler and Williams formulated their variant of the LAPW method \cite{prb_40_1560,prb_42_9728} with perfectly continuous basis functions, as well as the algorithm of the forces evaluation within their approach. Their construction certainly has some advantages, but the complications related to the fact that inside the MT spheres one has to deal with momentum independent functions as well as with momentum dependent plane waves, makes it inconvenient especially if one is interested in advanced methods going beyond DFT such as GW approximation.

One of the limitations of the existing formulations of the atomic force evaluation is that they are based on the non-relativistic Kohn-Sham equations. However, in materials where elements from the far end of the periodic table are present, one has to use fully relativistic approach based on the Dirac-Kohn-Sham (DKS) theory \cite{prb_7_1912,jpc_11_L943,jpc_12_2977,jpc_12_L845}. Therefore, this work has its principal goal in removing the above mentioned limitation. The derivation of the expression for forces goes closely along the lines paved in previous works by Yu et al. and by Kl\"{u}ppelberg et al., but with DKS equations as a background theory. Whereas our derivation is directly relevant to fully relativistic theory, we specifically are pointing out in the text, where the difference from the non-relativistic theory enter. Throughout the paper the atomic units are used with Rydbergs as units of energy.

\section{General derivation of the atomic force expression}
\label{gen}

The force $\mathbf{F}_{t}$ exerted on atom positioned at \textbf{t} is defined as the derivative of the free energy F of a solid: $\mathbf{F}_{t}=-\frac{d F}{d\mathbf{t}}$. Thus, it is convenient to begin with writing down an expression for the free energy which corresponds to a specific level of theory. In the context of a joint description of the relativistic and magnetic effects within the Relativistic Density Functional Theory (RDFT), the corresponding expression was developed in works by Rajagopal, Callaway, Vosko and Ramana \cite{prb_7_1912,jpc_11_L943,jpc_12_2977,jpc_12_L845}. Principal equations of this theory are briefly capitalized here for convenience. In RDFT, free energy of a solid with electronic density
$n(\mathbf{r})$ and magnetization density $\mathbf{m}(\mathbf{r})$ can be written as the following:
\begin{eqnarray}\label{etot}
F &= -T\sum_{\mathbf{k}\lambda}\ln (1+e^{-(\epsilon^{\mathbf{k}}_{\lambda}-\mu)/T})+\mu N
\nonumber \\&-\int_{\Omega_{0}} d \mathbf{r} [n(\mathbf{r})V_{eff}(\mathbf{r})+\mathbf{m}(\mathbf{r})\cdot
\mathbf{B}_{eff}(\mathbf{r})]\nonumber\\&+\int_{\Omega_{0}} d\mathbf{r} [n(\mathbf{r})V_{ext}(\mathbf{r})+\mathbf{m}(\mathbf{r})\cdot
\mathbf{B}_{ext}(\mathbf{r})] \nonumber\\&+ \int_{\Omega_{0}} d \mathbf{r}
\int_{\Omega} d \mathbf{r'} \frac{n(\mathbf{r})n(\mathbf{r'})}{|\mathbf{r}-\mathbf{r'}|}\nonumber\\&
+\int_{\Omega_{0}} d\mathbf{r} n(\mathbf{r})\epsilon_{xc}[n(\mathbf{r}),\mathbf{m}(\mathbf{r})]+E_{nn},
\end{eqnarray}
where $T$ stands for the temperature, sum runs over the Brillouin zone points $\mathbf{k}$ and band indexes $\lambda$, $\epsilon^{\mathbf{k}}_{\lambda}$ is the band energy, $\mu$ is the chemical potential, and $N$ is the total number of electrons in the unit cell. In the integrals, $\Omega_{0}$ is the volume of the primitive unit cell and $\Omega$ is the volume of the whole solid.
Effective scalar potential $V_{eff}(\mathbf{r})$ is a sum of an external scalar field $V_{ext}(\mathbf{r})$ and induced fields (Hartree (electrostatic) $V_{H}(\mathbf{r})=2 \int d\mathbf{r'}
\frac{n(\mathbf{r'})}{|\mathbf{r} - \mathbf{r'}|}$ and exchange-correlation $V_{xc}(\mathbf{r})=\frac{\delta
E_{xc} [n(\mathbf{r}),\mathbf{m}(\mathbf{r})]}{\delta n(\mathbf{r})}$):
\begin{equation} \label{veff}
V_{eff}(\mathbf{r})=V_{ext}(\mathbf{r})+V_{H}(\mathbf{r}) + V_{xc}(\mathbf{r}),
\end{equation}
whereas the effective magnetic field $\mathbf{B}_{eff}(\mathbf{r})$ represents a sum of external $\mathbf{B}_{ext}(\mathbf{r})$ and induced $\mathbf{B}_{xc}(\mathbf{r})=\frac{\delta E_{xc}
[n(\mathbf{r}),\mathbf{m}(\mathbf{r})]}{\delta \mathbf{m}(\mathbf{r})}$ magnetic fields:
\begin{equation} \label{beff}
\mathbf{B}_{eff}(\mathbf{r})=\mathbf{B}_{ext}(\mathbf{r})+ \mathbf{B}_{xc}(\mathbf{r}).
\end{equation}

$E_{xc}$ in the above formulae stands for the exchange-correlation energy which is a functional of $n(\mathbf{r})$ and $\mathbf{m}(\mathbf{r})$ : $\int_{\Omega_{0}} d\mathbf{r} n(\mathbf{r})\epsilon_{xc}[n(\mathbf{r}),\mathbf{m}(\mathbf{r})]$. $E_{nn}$ in (\ref{etot}) is the nuclear-nuclear electrostatic interaction energy.
One-electron energies $\epsilon^{\mathbf{k}}_{\lambda}$ are the eigen values of the following equations (Dirac-Kohn-Sham equations):
\begin{equation} \label{ksh}
\left(\hat{K}+V_{eff}(\mathbf{r})+ \beta
\widetilde{\boldsymbol{\sigma}} \cdot \mathbf{B}_{eff}(\mathbf{r})\right) \Psi_{\lambda}^{\mathbf{k}}(\mathbf{r})=\epsilon^{\mathbf{k}}_{\lambda}
\Psi_{\lambda}^{\mathbf{k}}(\mathbf{r}),
\end{equation}
where $\Psi_{\lambda}^{\mathbf{k}}(\mathbf{r})$ stands for the Bloch periodic band function. The kinetic energy operator $\hat{K}$ has the Dirac form (electron rest energy has been subtracted):
\begin{equation} \label{hkin}
\hat{K}=c \boldsymbol{\alpha} \cdot \mathbf{p} +(\beta-I)
\frac{c^2}{2},
\end{equation}
and $\widetilde{\boldsymbol{\sigma}}$ are the $4 \times 4$ matrices, combined from the Pauli matrices $\boldsymbol{\sigma}$:
\begin{equation}
\label{Pauli}
\widetilde{\boldsymbol{\sigma}}= \left(
\begin{array}{cc}
\boldsymbol{\sigma} & 0 \\
0 & \boldsymbol{\sigma}
\end{array} \right).
\end{equation}
$c$ in equation (\ref{hkin}) is the light velocity ($c=274.074$ in our unit system), $\mathbf{p}$ is the momentum operator ($\equiv -i \nabla$),
$\boldsymbol{\alpha}$, and $\beta$ are Dirac matrices in the standard representation, and $I$ is the unit $4 \times 4$ matrix.

Finally, with the electron energies and the band state functions available, the electronic and magnetization densities are defined as the following
\begin{equation} \label{dens}
n(\mathbf{r})=\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}
\Psi_{\lambda}^{^{\dag}\mathbf{k}}(\mathbf{r}) \Psi_{\lambda}^{\mathbf{k}}(\mathbf{r}),
\end{equation}
and
\begin{equation} \label{magn}
\mathbf{m}(\mathbf{r})=\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}
\Psi_{\lambda}^{\dag}(\mathbf{k},\mathbf{r}) \beta
\widetilde{\boldsymbol{\sigma}} \Psi_{\lambda}(\mathbf{k},\mathbf{r}),
\end{equation}
with $f^{\mathbf{k}}_{\lambda}$ being the Fermi-Dirac distribution function ($f^{\mathbf{k}}_{\lambda}=\frac{1}{1+e^{(\epsilon^{\mathbf{k}}_{\lambda}-\mu)/T}}$).

Now we are differentiating the Eq.(\ref{etot}) term by term. For the first and the second terms on the right hand side one gets:

\begin{eqnarray} \label{dif12}
-\frac{d}{d\mathbf{t}}&\left(-T\sum_{\mathbf{k}\lambda}\ln (1+e^{-(\epsilon^{\mathbf{k}}_{\lambda}-\mu)/T})+\mu N\right)
\nonumber\\=&-\sum_{\mathbf{k}\lambda}\frac{1}{1+e^{(\epsilon^{\mathbf{k}}_{\lambda}-\mu)/T}}
\left(\frac{d \epsilon^{\mathbf{k}}_{\lambda}}{d\mathbf{t}}-\frac{d \mu}{d\mathbf{t}}\right)-N\frac{d \mu}{d\mathbf{t}}
=\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\frac{d \epsilon^{\mathbf{k}}_{\lambda}}{d\mathbf{t}}.
\end{eqnarray}

The terms from the third to the sixth on the right hand side of (\ref{etot}) are represented by integrals over the unit cell. In APW-related methods, it means technically the sum of the integrals over the non-overlapping MT spheres and over the interstitial region (IR). As authors of work \cite{prb_91_035105} pointed out, the integrals should be differentiated with care, namely, the change of the integration domain when atom (and its muffin-tin sphere) moves should be taken into account. The generic differentiation formula obtained in \cite{prb_91_035105} is the following:

\begin{eqnarray} \label{int_gen}
\frac{d}{d\mathbf{t}}\int_{\Omega_{0}} d\mathbf{r}f(\mathbf{r})=\int_{\Omega_{0}}d\mathbf{r}\frac{d f(\mathbf{r})}{d\mathbf{t}}+\int_{S_{t}}d\mathbf{S}
[f^{MT}(\mathbf{r})-f^{IR}(\mathbf{r})],
\end{eqnarray}
where the surface integral is taken over the MT sphere of atom $t$. $d\mathbf{S}=\mathbf{e}dS$, and $\mathbf{e}=\frac{\mathbf{r}-\mathbf{t}}{|\mathbf{r}-\mathbf{t}|}$ denotes the normal vector on the MT sphere of atom $t$ that points into the interstitial region. $f^{MT}(\mathbf{r})$ and $f^{IR}(\mathbf{r})$ distinguish between the MT and the IR representations of the function $f$. Let us now apply the generic formula (\ref{int_gen}) to the integrals in (\ref{etot}):

\begin{eqnarray} \label{dif3}
-\frac{d}{d\mathbf{t}}&\left(-\int_{\Omega_{0}} d \mathbf{r} [n(\mathbf{r})V_{eff}(\mathbf{r})+\mathbf{m}(\mathbf{r})\cdot
\mathbf{B}_{eff}(\mathbf{r})]\right)\nonumber\\&=\int_{\Omega_{0}} d \mathbf{r} [\frac{d n(\mathbf{r})}{d\mathbf{t}}V_{eff}(\mathbf{r})+n(\mathbf{r})\frac{d V_{eff}(\mathbf{r})}{d\mathbf{t}}\nonumber\\&+\frac{d \mathbf{m}(\mathbf{r})}{d\mathbf{t}}\cdot
\mathbf{B}_{eff}(\mathbf{r})+\mathbf{m}(\mathbf{r})\cdot
\frac{d\mathbf{B}_{xc}(\mathbf{r})}{d\mathbf{t}}]\nonumber\\&+\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})V^{MT}_{eff}(\mathbf{r})-n^{IR}(\mathbf{r})V^{IR}_{eff}(\mathbf{r})]\nonumber\\&+\int_{S_{t}}d\mathbf{S}
[\mathbf{m}^{MT}(\mathbf{r})\cdot \mathbf{B}^{MT}_{eff}(\mathbf{r})-\mathbf{m}^{IR}(\mathbf{r})\cdot \mathbf{B}^{IR}_{eff}(\mathbf{r})],
\end{eqnarray}
where we have assumed that only the induced magnetic field depends of the position of atom $t$.

\begin{eqnarray} \label{dif4}
-\frac{d}{d\mathbf{t}}&\left(\int_{\Omega_{0}} d \mathbf{r} [n(\mathbf{r})V_{ext}(\mathbf{r})+\mathbf{m}(\mathbf{r})\cdot
\mathbf{B}_{ext}(\mathbf{r})]\right)\nonumber\\&=-\int_{\Omega_{0}} d \mathbf{r} [\frac{d n(\mathbf{r})}{d\mathbf{t}}V_{ext}(\mathbf{r})+n(\mathbf{r})\frac{d V_{ext}(\mathbf{r})}{d\mathbf{t}}+\frac{d \mathbf{m}(\mathbf{r})}{d\mathbf{t}}\cdot
\mathbf{B}_{ext}(\mathbf{r})]\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})V^{MT}_{ext}(\mathbf{r})-n^{IR}(\mathbf{r})V^{IR}_{ext}(\mathbf{r})]\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[\mathbf{m}^{MT}(\mathbf{r})\cdot \mathbf{B}^{MT}_{ext}(\mathbf{r})-\mathbf{m}^{IR}(\mathbf{r})\cdot \mathbf{B}^{IR}_{ext}(\mathbf{r})],
\end{eqnarray}

\begin{eqnarray} \label{dif5}
-\frac{d}{d\mathbf{t}}&\left(\int_{\Omega_{0}} d \mathbf{r}
\int_{\Omega} d \mathbf{r'} \frac{n(\mathbf{r})n(\mathbf{r'})}{|\mathbf{r}-\mathbf{r'}|}\right)=-\int_{\Omega_{0}} d \mathbf{r} \frac{d n(\mathbf{r})}{d\mathbf{t}}V_{H}(\mathbf{r})\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})V^{MT}_{H}(\mathbf{r})-n^{IR}(\mathbf{r})V^{IR}_{H}(\mathbf{r})]\nonumber\\&,
\end{eqnarray}

\begin{eqnarray} \label{dif6}
-\frac{d}{d\mathbf{t}}&\left(\int_{\Omega_{0}} d\mathbf{r} n(\mathbf{r})\epsilon_{xc}[n(\mathbf{r}),\mathbf{m}(\mathbf{r})]\right)\nonumber\\&=-\int_{\Omega_{0}} d \mathbf{r} [\frac{d n(\mathbf{r})}{d\mathbf{t}}V_{xc}(\mathbf{r})+\frac{d \mathbf{m}(\mathbf{r})}{d\mathbf{t}}\cdot
\mathbf{B}_{xc}(\mathbf{r})]\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})\epsilon^{MT}_{xc}(\mathbf{r})-n^{IR}(\mathbf{r})\epsilon^{IR}_{xc}(\mathbf{r})],
\end{eqnarray}

Collecting all derivatives together and assuming self-consistency (i.e. equations (\ref{veff}) and (\ref{beff}) are met) we obtain the following force:

\begin{eqnarray} \label{forc1}
\mathbf{F}_{t}=&-\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\frac{d \epsilon^{\mathbf{k}}_{\lambda}}{d\mathbf{t}}+\int_{\Omega_{0}} d \mathbf{r} [n(\mathbf{r})\frac{d V_{eff}(\mathbf{r})}{d\mathbf{t}}+\mathbf{m}(\mathbf{r})\cdot
\frac{d\mathbf{B}_{xc}(\mathbf{r})}{d\mathbf{t}}]\nonumber\\&+\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})V^{MT}_{xc}(\mathbf{r})-n^{IR}(\mathbf{r})V^{IR}_{xc}(\mathbf{r})]\nonumber\\&
+\int_{S_{t}}d\mathbf{S}
[\mathbf{m}^{MT}(\mathbf{r})\cdot \mathbf{B}^{MT}_{xc}(\mathbf{r})-\mathbf{m}^{IR}(\mathbf{r})\cdot \mathbf{B}^{IR}_{xc}(\mathbf{r})]\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})\epsilon^{MT}_{xc}(\mathbf{r})-n^{IR}(\mathbf{r})\epsilon^{IR}_{xc}(\mathbf{r})]\nonumber\\&+\mathbf{F}^{HF}_{t},
\end{eqnarray}
where the Hellmann-Feynman force has been introduced:

\begin{equation}\label{hell1}
\mathbf{F}^{HF}_{t}=-\int_{\Omega_{0}} d \mathbf{r} n(\mathbf{r})\frac{d V_{ext}(\mathbf{r})}{d\mathbf{t}}-\frac{d}{d\mathbf{t}} E_{nn}.
\end{equation}

Hellmann-Feynman force is proportional to the gradient of the full electrostatic potential at the center of atom $t$ (excluding the field from its nuclear), \cite{prb_43_6411}:
\begin{eqnarray}\label{hell2}
\mathbf{F}^{HF}_{t}&=2Z_{t}\frac{\partial}{\partial\mathbf{t}} \left(\int_{\Omega} d \mathbf{r} \frac{n(\mathbf{r})}{|\mathbf{t}-\mathbf{r}|}-\sum'_{\mathbf{R}\mathbf{t}'}\frac{Z_{t'}}{|\mathbf{t}-\mathbf{t}'-\mathbf{R}|}\right)\nonumber\\&=-Z_{t}\nabla V'_{el-stat}(\mathbf{r})\bigg|_{\mathbf{r}\rightarrow \mathbf{t}},
\end{eqnarray}
where $Z_{t}$ is the nuclear charge of atom $t$, the integration in the first right hand side expression is performed over the whole solid, and the sum is taken over all unit cells (indexed here by translation vector $\mathbf{R}$) and over all atoms in the unit cell (atom $t$ in the central unit cell is excluded from the sum). This consideration makes the evaluation of the Hellmann-Feynman term easy.

In order to bring the remaining terms of (\ref{forc1}) to the form convenient for evaluation one has to consider the derivative of the one-electron energies. This is done in the next section. Let us also to point out that the  derivation performed up to this point is quite generic with respect to the degree of inclusion of the relativistic effects. The only formal difference is that we use vectors of the magnetization and the magnetic field as it is usually done in the spin-polarized RDFT, instead of spin up and spin down quantities as it is done in the non-relativistic spin-polarized DFT.

\section{Specifics of differentiation of the Dirac-Kohn-Sham eigenvalues}
\label{dkse}

Differentiation of the Kohn-Sham (Dirac-Kohn-Sham) eigenvalues with respect to atomic positions is rather involved. In order to keep derivation as clear as possible we will do it in a step by step fashion. Essentially the derivation is very similar to the one done by Yu et al. \cite{prb_43_6411} and by Kl\"{u}ppelberg et al. \cite{prb_91_035105}. We repeat all the steps here to make it clear where the fully relativistic formalism enters and where the formulae are independent on the formalism (relativistic or non-relativistic). We will consider the derivatives of the valence and core states separately beginning with the valence states.

As a first step, we show explicitly that only the derivatives of the basis functions enter the expression for the forces but not the derivatives of the coefficients. It can be done generically without specification of the basis set or relativistic effects. In methods which use non-orthogonal basis sets the eigenvalues can be found as the ratio of the expectation values of the hamiltonian and overlap matrices:

\begin{equation}\label{eig}
\epsilon=\frac{\sum_{ij}A^{*}_{i}H_{ij}A_{j}}{\sum_{ij}A^{*}_{i}O_{ij}A_{j}},
\end{equation}
where sums run over the basis set indexes and $A_{i}$ are the expansion coefficients. Again, using generic differentiation which we denote as prime, we obtain:

\begin{eqnarray}\label{deig}
\epsilon'&=\sum_{ij}\left(A'^{*}_{i}H_{ij}A_{j}+A^{*}_{i}H'_{ij}A_{j}+A^{*}_{i}H_{ij}A'_{j}\right)\nonumber\\&+\epsilon \sum_{ij}\left(A'^{*}_{i}O_{ij}A_{j}+A^{*}_{i}O'_{ij}A_{j}+A^{*}_{i}O_{ij}A'_{j}\right)\nonumber\\&=\sum_{ij}\left(A'^{*}_{i}[H_{ij}-\epsilon O_{ij}]A_{j}+A^{*}_{i}[H'_{ij}-\epsilon O'_{ij}]A_{j}+A^{*}_{i}[H_{ij}-\epsilon O_{ij}]A'_{j}\right)\nonumber\\&=\sum_{ij}A^{*}_{i}[H'_{ij}-\epsilon O'_{ij}]A_{j},
\end{eqnarray}
where we have used the fact that matrix equations are solved numerically exactly (i.e. for instance $\sum_{j}[H_{ij}-\epsilon O_{ij}]A_{j}$ is zero with computer accuracy). From (\ref{deig}), it is obvious that we have to differentiate only the matrix elements but not the coefficients.

Before proceeding further, let us briefly specify the basis functions (or their combinations) which we are using. As it becomes common practice in the APW-based calculations \cite{cpc_184_2670,cpc_220_230,arx_2012_04992}, we use generic combination of an augmentation function (APW or LAPW) and local orbitals of different kind. As local orbitals, we use the so called 'lo'-orbitals which have discontinuity in its small component (in its derivative in non-relativistic formulation) at the MT sphere boundaries. It is used in combination with APW augmentation \cite{prb_64_195134} to improve variational flexibility of the basis set. Next type of the local orbital is the so called High Derivative Local Orbitals (HDLO) \cite{prb_74_045104,cpc_184_2670,cpc_220_230} which can be used in combination with LAPW or APW+lo to further enhance the accuracy of the basis set in the range of energies corresponding to the valence bands. Finally, the so called High Energy Local Orbitals (HELO's, \cite{cpc_184_2670,cpc_220_230}) can be included in a basis set to describe semicore states or high energy states in the conduction band range of energies.

Hamiltonian and overlap matrix elements are represented by the volume integrals over all MT spheres in the unit cell and over the interstitial region. For the basis functions with discontinuities at the MT surfaces (for instance if the APW+lo combination is used), matrix elements of the hamiltonian include the surface correction terms as it was specified in Ref. \cite{prb_64_195134} for the non-relativistic case and in \cite{arx_2012_04992} for the fully relativistic case. The recipe (\ref{int_gen}) is applied for the differentiation when integration domain changes. Still using generic indexes for the basis set but specifying the band index and the \textbf{k}-point (i.e. using $f^{\mathbf{k}}_{i}$ as generic basis functions and $\epsilon^{\mathbf{k}}_{\lambda}$ as eigenvalues) as well as the specific form (\ref{ksh}) of the Dirac-Kohn-Sham hamiltonian $H_{DKS}$, we obtain:

\begin{eqnarray}\label{dks}
\hspace*{-2cm}
\frac{d H^{\mathbf{k}}_{ij}}{d\mathbf{t}}-\epsilon^{\mathbf{k}}_{\lambda}\frac{d O^{\mathbf{k}}_{ij}}{d\mathbf{t}}&=
P^{\mathbf{k}\lambda}_{ij}\nonumber\\&+\int_{\Omega_{0}} d\mathbf{r}
f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})[\frac{d V_{eff}(\mathbf{r})}{d\mathbf{t}}+\beta
\widetilde{\boldsymbol{\sigma}} \cdot \frac{d\mathbf{B}_{eff}(\mathbf{r})}{d\mathbf{t}}]f^{\mathbf{k}}_{j}(\mathbf{r})\nonumber\\&\hspace*{-3cm}+\int_{S_{t}}d\mathbf{S}
\left(f^{^{\dagger}\mathbf{k}(MT)}_{i}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]f^{\mathbf{k}(MT)}_{j}(\mathbf{r})-f^{^{\dagger}\mathbf{k}(IR)}_{i}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]f^{\mathbf{k}(IR)}_{j}(\mathbf{r})\right),
\end{eqnarray}
where the terms which later will contribute to the Pulay force have been collected into the quantity $P^{\mathbf{k}\lambda t}_{ij}$:

\begin{eqnarray}\label{pulay1}
P^{\mathbf{k}\lambda t}_{ij}&=\int_{\Omega_{0}} d\mathbf{r}
\frac{d f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})}{d\mathbf{t}}[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]f^{\mathbf{k}}_{j}(\mathbf{r})\nonumber\\&+\int_{\Omega_{0}} d\mathbf{r}
f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]\frac{d f^{\mathbf{k}}_{j}(\mathbf{r})}{d\mathbf{t}}+\frac{d}{d\mathbf{t}}S^{\mathbf{k}(DISC)}_{ij}.
\end{eqnarray}
Derivatives of the terms which appear in the hamiltonian when some of the basis functions have discontinuities were denoted as $\frac{d}{d\mathbf{t}}S^{\mathbf{k}(DISC)}_{ij}$. We do not specify them here because they will be combined with other explicitly dependent on the atomic position terms in the same way as they were combined in the derivation of the matrix elements of the hamiltonian \cite{arx_2012_04992}.

At this point of the derivation we have to take into account the differences between basis functions of augmentation type (APW and LAPW) and local basis functions (lo, HDLO, and HELO). Also, taking the derivatives assumes an understanding of the quantities themselves. So, in order to avoid the repetition of a rather lengthy derivation of the basis functions and matrix elements which has been done in \cite{arx_2012_04992}, we ask the reader to have the paper \cite{arx_2012_04992} at hand for quick references (we will refer to the equations in that paper as (I-???) with '???' as the equation number). Keeping this in mind, let us proceed with formal differentiation.

For the augmentation functions defined in (I-12,39), the derivative is not zero only in the MT sphere of atom $t$:

\begin{equation} \label{dbasmt}
\frac{d}{d\mathbf{t}}\Pi^{\mathbf{k}}_{\mathbf{G}s}(\mathbf{r})= i(\mathbf{k}+\mathbf{G})\Pi^{\mathbf{k}}_{\mathbf{G}s}(\mathbf{r})-\nabla \Pi^{\mathbf{k}}_{\mathbf{G}s}(\mathbf{r}),
\end{equation}
where the first term comes from the augmentation constraints and the second from the dependence of the radial functions on atomic position. Derivative of the local functions (I-46) also has two terms stemming from a formal Bloch factor and from the same position dependence:
\begin{equation} \label{dbasloc}
\frac{d}{d\mathbf{t}}\Lambda^{\mathbf{k}}_{tnil\mu}(\mathbf{r})= i\mathbf{k}\Lambda^{\mathbf{k}}_{tnil\mu}(\mathbf{r})-\nabla \Lambda^{\mathbf{k}}_{tnil\mu}(\mathbf{r}).
\end{equation}

Let us first consider the contribution of the gradient terms (which is generic) from (\ref{dbasmt}) and (\ref{dbasloc}) into the quantity $P^{\mathbf{k}\lambda t}_{ij}$ in (\ref{pulay1}):

\begin{eqnarray}\label{pulay2}
\hspace*{-2cm}
-\int_{\Omega_{t}} d\mathbf{r}
\nabla &f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]f^{\mathbf{k}}_{j}(\mathbf{r})-\int_{\Omega_{t}} d\mathbf{r}
f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]\nabla f^{\mathbf{k}}_{j}(\mathbf{r})\nonumber\\&=-\int_{\Omega_{t}} d\mathbf{r}
\nabla \left(f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})[\hat{K}-\epsilon^{\mathbf{k}}_{\lambda}]f^{\mathbf{k}}_{j}(\mathbf{r})\right)-\int_{\Omega_{t}} d\mathbf{r}
V_{eff}(\mathbf{r})\nabla [f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})f^{\mathbf{k}}_{j}(\mathbf{r})]\nonumber\\&-\int_{\Omega_{t}} d\mathbf{r}
\nabla [f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})\beta
\widetilde{\boldsymbol{\sigma}} \cdot f^{\mathbf{k}}_{j}(\mathbf{r})]\mathbf{B}_{eff}(\mathbf{r})\nonumber\\&=-\int_{S_{t}} \mathbf{e} dS
f^{^{\dagger}\mathbf{k}(MT)}_{i}(\mathbf{r})[\hat{K}-\epsilon^{\mathbf{k}}_{\lambda}]f^{\mathbf{k}(MT)}_{j}(\mathbf{r})-\int_{\Omega_{t}} d\mathbf{r}
V_{eff}(\mathbf{r})\nabla [f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})f^{\mathbf{k}}_{j}(\mathbf{r})]\nonumber\\&-\int_{\Omega_{t}} d\mathbf{r}
\nabla [f^{^{\dagger}\mathbf{k}}_{i}(\mathbf{r})\beta
\widetilde{\boldsymbol{\sigma}} \cdot f^{\mathbf{k}}_{j}(\mathbf{r})]\mathbf{B}_{eff}(\mathbf{r}).
\end{eqnarray}

Let us now consider the contribution from the augmentation parts of the derivatives in (\ref{dbasmt}) and (\ref{dbasloc}). It is easier to take the derivative of the final matrix element, however. In this case one can automatically include the derivatives of the discontinuities because the corresponding contributions to the matrix elements have exactly the same structure of explicit dependence on the atomic positions as the volume integral contributions \cite{arx_2012_04992}. Distinguishing the cases of the matrix elements between two augmentation functions (AA) specified in (I-60,61,68), between the local and the augmentation function (BA) specified in (I-63,64,69), and between two local functions (BB, I-66,67,70) one obtains the corresponding contribution to the quantity (\ref{pulay1}):

\begin{eqnarray}\label{dhmtaa}
\hspace*{-2cm}
i(\mathbf{G}'-\mathbf{G})&\Big[F^{t}_{\mathbf{G'-G}}\sum_{il}D^{\mathbf{k}il}_{\mathbf{G}s;\mathbf{G}'s'}[\overline{h}^{til}_{\mathbf{GG}'}-\epsilon^{\mathbf{k}}_{\lambda}\overline{o}^{til}_{\mathbf{GG}'}]\nonumber\\&+\sum_{il\mu
;i'l'\mu'}\sum_{(ww')=1}^{N^{t}_{l}}y^{^{*}(w)\mathbf{k}}_{til\mu;\mathbf{G}s}
y^{(w')\mathbf{k}}_{ti'l'\mu';\mathbf{G'}s'} \int_{\Omega_{t}}
R^{^{\dagger}(w)t}_{il\mu}(\mathbf{r}) \hat{H}_{NMT}
R^{(w')t}_{i'l'\mu'}(\mathbf{r}) d \mathbf{r}\Big]
\end{eqnarray}
for the AA type, and
\begin{eqnarray}\label{dhmtba}
i\mathbf{G}'\Big[&F^{t}_{\mathbf{G'-G}}\sum_{il}D^{\mathbf{k}il}_{\mathbf{G}s;\mathbf{G}'s'}[\overline{h}^{til}_{\mathbf{GG}'}-\epsilon^{\mathbf{k}}_{\lambda}\overline{o}^{til}_{\mathbf{GG}'}]\nonumber\\&+\sum_{i'l'\mu'}\sum_{w'=1}^{N^{t}_{l}}
y^{(w')\mathbf{k}}_{ti'l'\mu';\mathbf{G'}s'} \int_{\Omega_{t}}
R^{^{\dag}(LOC)t}_{nil\mu}(\mathbf{r}) \hat{H}_{NMT}
R^{(w')t}_{i'l'\mu'}(\mathbf{r}) d \mathbf{r}\Big]
\end{eqnarray}
for the BA type. Derivatives of the matrix elements of BB type equal to zero. The above expressions (\ref{dhmtaa}) and (\ref{dhmtba}) comprise a matrix with indexes running over the whole basis set. Anticipating a convolution of this matrix with the variational coefficients (see Eq. (\ref{deig})), it is convenient to denote this convolution as $\mathbf{C}^{\mathbf{k}t}_{\lambda}$ for a future use. The equations (\ref{dhmtaa}) and (\ref{dhmtba}) are the place where most of the differences between the fully relativistic and the non-relativistic formulations are concentrated. Whereas it is not the goal of this work to give a comprehensive account of all levels of the relativistic effects, it is helpful to know where the differences are located. Particularly, if one needs to recover all non-relativistic equations, the quantities $D^{\mathbf{k}il}_{\mathbf{G}s;\mathbf{G}'s'}$, $\overline{h}^{til}_{\mathbf{GG}'}$, and $\overline{o}^{til}_{\mathbf{GG}'}$ which are defined in (I-59,60,61) for the fully relativistic case, have to be replaced with their non-relativistic analogues.

Now it is a time to perform the Brillouin zone and the band index sums in the basis set convolution of the expression (\ref{dks}) and, correspondingly, to evaluate the first term on the right hand side of (\ref{forc1}):
\begin{eqnarray}\label{dksS}
\hspace*{-2cm}
-\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\frac{d \epsilon^{\mathbf{k}}_{\lambda}}{d\mathbf{t}}&=-\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\sum_{ij}A^{^{*}\mathbf{k}}_{i\lambda}\Big[\frac{d H^{\mathbf{k}}_{ij}}{d\mathbf{t}}-\epsilon^{\mathbf{k}}_{\lambda}\frac{d O^{\mathbf{k}}_{ij}}{d\mathbf{t}}\Big]A^{\mathbf{k}}_{\lambda}\nonumber\\&=
\mathbf{F}^{Pulay}_{t,val}-\int_{\Omega_{0}} d\mathbf{r}
[n_{val}(\mathbf{r})\frac{d V_{eff}(\mathbf{r})}{d\mathbf{t}}+\mathbf{m}_{val}(\mathbf{r})\beta\widetilde{\boldsymbol{\sigma}} \cdot \frac{d\mathbf{B}_{eff}(\mathbf{r})}{d\mathbf{t}}]\nonumber\\&\hspace*{-3cm}-\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\int_{S_{t}}d\mathbf{S}
\Big[\Psi^{^{\dagger}\mathbf{k}(MT)}_{\lambda}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]\Psi^{\mathbf{k}(MT)}_{\lambda}(\mathbf{r})-\Psi^{^{\dagger}\mathbf{k}(IR)}_{\lambda}(\mathbf{r})[H_{DKS}-\epsilon^{\mathbf{k}}_{\lambda}]\Psi^{\mathbf{k}(IR)}_{\lambda}(\mathbf{r})\Big],
\end{eqnarray}
with the valence Pulay force
\begin{eqnarray}\label{pulay3}
\mathbf{F}^{Pulay}_{t,val}&=-\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\mathbf{C}^{\mathbf{k}t}_{\lambda}
+\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\int_{S_{t}} d\mathbf{S}
\Psi^{^{\dagger}\mathbf{k}(MT)}_{\lambda}(\mathbf{r})[\hat{K}-\epsilon^{\mathbf{k}}_{\lambda}]\Psi^{\mathbf{k}(MT)}_{\lambda}(\mathbf{r})\nonumber\\&+\int_{\Omega_{t}} d\mathbf{r}
V_{eff}(\mathbf{r})\nabla n_{val}(\mathbf{r})+\int_{\Omega_{t}} d\mathbf{r}
\nabla [\mathbf{m}_{val}(\mathbf{r})]\cdot\mathbf{B}_{eff}(\mathbf{r}).
\end{eqnarray}

For the core states we can formally repeat all above steps which we have done for the valence states, with a number of simplifications. The simplifications are related to the following two facts: i) each core state is an exact solution of the Dirac-Kohn-Sham equation for a spherically symmetric potential as opposite to an expansion in a basis set for the valence levels; ii) core states are strictly confined inside the corresponding MT sphere with zero values and derivatives at the boundary. As a result, all surface terms related to the augmentation or the discontinuities disappear. Equations (\ref{dksS}) and (\ref{pulay3}) for the core states, therefore, can be simplified as the following:
\begin{eqnarray}\label{dksSc}
\hspace*{-2cm}
-\sum_{c}\frac{d \epsilon_{c}}{d\mathbf{t}}&=
\mathbf{F}^{Pulay}_{t,cor}-\int_{\Omega_{0}} d\mathbf{r}
[n_{cor}(\mathbf{r})\frac{d V_{eff}(\mathbf{r})}{d\mathbf{t}}+\mathbf{m}_{cor}(\mathbf{r})\beta\widetilde{\boldsymbol{\sigma}} \cdot \frac{d\mathbf{B}_{eff}(\mathbf{r})}{d\mathbf{t}}],
\end{eqnarray}
with $c$ running over the core states of atom $t$ and with the core Pulay force
\begin{eqnarray}\label{pulay3c}
\mathbf{F}^{Pulay}_{t,cor}=\int_{\Omega_{t}} d\mathbf{r}
V_{eff}(\mathbf{r})\nabla n_{cor}(\mathbf{r})+\int_{\Omega_{t}} d\mathbf{r}
\nabla [\mathbf{m}_{cor}(\mathbf{r})]\cdot\mathbf{B}_{eff}(\mathbf{r}).
\end{eqnarray}

Finally we can include the contribution from the eigenvalue derivatives (\ref{dksS}) and (\ref{dksSc}) into a general force equation (\ref{forc1}) to finish the derivation:

\begin{eqnarray} \label{forc2}
\mathbf{F}_{t}=\mathbf{F}^{HF}_{t}+\mathbf{F}^{Pulay}_{t,cor}+\mathbf{F}^{Pulay}_{t,val}+\mathbf{F}^{Surf}_{t,kin}+\mathbf{F}^{Surf}_{t,other},
\end{eqnarray}
where we have made the following definitions:

\begin{eqnarray} \label{forc3}
\mathbf{F}^{Surf}_{t,kin}=-\sum_{\mathbf{k}\lambda}f^{\mathbf{k}}_{\lambda}\int_{S_{t}}d\mathbf{S}&
\Big[\Psi^{^{\dagger}\mathbf{k}(MT)}_{\lambda}(\mathbf{r})[\hat{K}-\epsilon^{\mathbf{k}}_{\lambda}]\Psi^{\mathbf{k}(MT)}_{\lambda}(\mathbf{r})\nonumber\\&-\Psi^{^{\dagger}\mathbf{k}(IR)}_{\lambda}(\mathbf{r})[\hat{K}-\epsilon^{\mathbf{k}}_{\lambda}]\Psi^{\mathbf{k}(IR)}_{\lambda}(\mathbf{r})\Big],
\end{eqnarray}

\begin{eqnarray} \label{forc4}
\mathbf{F}^{Surf}_{t,other}=&-\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})V^{MT}_{eff}(\mathbf{r})-n^{IR}(\mathbf{r})V^{IR}_{eff}(\mathbf{r})]\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[\mathbf{m}^{MT}(\mathbf{r})\cdot \mathbf{B}^{MT}_{eff}(\mathbf{r})-\mathbf{m}^{IR}(\mathbf{r})\cdot \mathbf{B}^{IR}_{eff}(\mathbf{r})]\nonumber\\&+\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})V^{MT}_{xc}(\mathbf{r})-n^{IR}(\mathbf{r})V^{IR}_{xc}(\mathbf{r})]\nonumber\\&
+\int_{S_{t}}d\mathbf{S}
[\mathbf{m}^{MT}(\mathbf{r})\cdot \mathbf{B}^{MT}_{xc}(\mathbf{r})-\mathbf{m}^{IR}(\mathbf{r})\cdot \mathbf{B}^{IR}_{xc}(\mathbf{r})]\nonumber\\&-\int_{S_{t}}d\mathbf{S}
[n^{MT}(\mathbf{r})\epsilon^{MT}_{xc}(\mathbf{r})-n^{IR}(\mathbf{r})\epsilon^{IR}_{xc}(\mathbf{r})],
\end{eqnarray}
\section{Performance tests}
\label{conv}

\begin{table}[t]
\caption{Structural parameters of the solids considered in this work. The parameters correspond to the equilibrium geometries with zero forces. The change in atomic positions when we evaluate the forces is specified for each case later.} \label{list_s}
\begin{center}
\begin{tabular}{@{}c c c c c c c} &Space&&&&Wyckoff&$R_{MT}$\\
Solid &group&a(\AA)&b(\AA)&c(\AA)&positions&(a$_{B}$)\\
\hline\hline
$\alpha$-U&63 &2.854 &5.869 &4.955&0;0.1025;0.25  &2.602333\\
PuCoGa$_{5}$&123 &4.2354 & &6.7939&Pu: 0;0;0&Pu, Ga(1): 2.829752\\
& & & &&Co: 0;0;1/2  &Co, Ga(4): 2.34805\\
& & & &&Ga(4): 0;1/2;0.3086  &\\
& & & &&Ga(1): 1/2;1/2;0  &\\
FePt&123 &2.7248 & &3.78&Fe: 0;0;0&Fe, Pt: 2.55\\
& & & &&Pt: 1/2;1/2;1/2  &\\
\end{tabular}
\end{center}
\end{table}

\begin{table}[b]
\caption{Principal set up parameters of the studied solids.} \label{setup_s}
\begin{center}
\begin{tabular}{@{}c c c c c c} &Core&&$L_{max}$&$L_{max}$&\\
Solid &states&Semicore&$\Psi/\rho,V$&APW+lo+HDLO & $RK_{max}$ \\
\hline\hline
$\alpha$-U&[Kr]4d,4f&5s,6s,5p,6p,5d&12/8&3&12.0  \\
PuCoGa$_{5}$&Pu: [Kr]4d,4f,5s&Pu: 6s,5p,6p,5d&Pu: 12/10&Pu: 3&9.0  \\
& Co: [Ne]&Co: 3s,3p&Co: 10/10&Co: 2&  \\
& Ga: [Ne]&Ga: 3s,3p,3d&Ga: 10/10&Ga: 2&  \\
FePt&Fe: [Ne]&Fe: 3s,3p&10/10&Fe: 2&12.0  \\
& Pt: [Kr]&Pt: 5s,5p,4d,4f&&Pt: 3&  \\
\end{tabular}
\end{center}
\end{table}

This section presents results of the calculations. In order to make presentation more compact, principal structural parameters for studied solids have been collected in Table \ref{list_s} and the most important set up parameters have been collected in Table \ref{setup_s}. The APW type of the plane waves augmentation was used for the "physically relevant" orbital momenta which roughly correspond to the shells which have electrons in a free atom. This type of augmentation was accompanied with addition of two local orbitals (lo and HDLO) in order to enhance variational freedom. For higher orbital momenta, LAPW type of augmentation was applied. The separation of the augmentation strategy into APW+lo and LAPW was suggested in Ref. \cite{prb_64_195134}. Additional use of HDLO's was advocated in \cite{cpc_184_2670,cpc_220_230} and, in the context of the fully relativistic calculations, in \cite{arx_2012_04992}. High energy Local Orbitals (HELO's) were used for the "physically relevant" orbital momenta, but their effect on the calculated values of the forces was rather small. Radii of the muffin-tin spheres were selected to be the largest allowed (no overlapping). In the cases of competing sizes the ratio was 1:1. All results presented below correspond to the fully relativistic approach (FRA). A few tests performed with simplified relativistic approach (SRA, \cite{arx_2012_04992}) have shown very little difference with FRA. All calculations have been performed for the electronic temperature $T=300K$. Exchange-correlation functional corresponded to the local density approximation (LDA) as parametrized in \cite{prb_45_13244}.

\begin{table}[t]
\caption{Calculated $\alpha$-U free energy (Ry), forces (mRy/a$_{B}$), and forces evaluated by numerical differentiation of the free energy, for the different Brillouin zone samplings. Forces were evaluated for the structure with atomic positions $\mathbf{t}=\pm 0.13\mathbf{B}+1/4\mathbf{C}$ which is slightly perturbed from the equilibrium one with atomic positions $\mathbf{t}=\pm 0.1025\mathbf{B}+1/4\mathbf{C}$. The free energy corresponding to the perturbed structure is denoted as F(0) in the table. In order to evaluate the forces numerically, two additional small distortions relative to the already perturbed structure were considered. Atomic positions for these distorted structures were $\mathbf{t}=\pm (0.13\pm\Delta)\mathbf{B}+1/4\mathbf{C}$ with $\Delta=0.0005$. Free energies are given relative to the constant -112276 Ry.} \label{U_conv}
\begin{center}
\begin{tabular}{@{}c c c c c c c} Number of & &   &  & & Numerical& Mismatch    \\
\textbf{k}-points & F($-\Delta$)&   F(0)& F($+\Delta$) & Force & force &    (\%) \\
\hline\hline
144  & -0.4468887 &-0.4463283 &-0.4457661 &  -50.6287 &  -50.6095 & -0.04 \\
384  & -0.4466470 &-0.4460889 &-0.4455300 &  -50.3776 &  -50.3570 & -0.04 \\
700  & -0.4465695 &-0.4460026 &-0.4454351 &  -51.1774 &  -51.1415 & -0.07 \\
1152 & -0.4466360 &-0.4460739 &-0.4455112 &  -50.7377 &  -50.7087 & -0.06 \\
2560 & -0.4465665 &-0.4460026 &-0.4454378 &  -50.9115 &  -50.8845 & -0.05 \\
\end{tabular}
\end{center}
\end{table}

Special remark is about core states. As authors of the Ref. \cite{prb_91_035105} stress, those core states which are not exactly confined inside their MT spheres may affect the calculated forces noticeably. Such core states were allowed in \cite{prb_91_035105} to extend beyond their MT spheres and into the interstitial region and in other MT spheres with subsequent correction of the calculated forces via the plane waves expansion of their tails. This approach allows one to minimize the size of the matrices as only valence states need to be described by the basis set. The price, however, is the increased complexity of the core states treatment. Another way to handle the "shallow" core states is to include them in the list of the semicore states. In this case the size of the matrices increases slightly, but strict confinement of the remaining ('deep") core states inside their MT spheres makes the algorithm simpler, which is especially important when one builds approaches of a higher complexity (like the GW approximation) on top of the DFT code. This approach is accepted in the FlapwMBPT code.

\begin{table}[b]
\caption{Calculated free energy (Ry), forces (mRy/a$_{B}$), and numerical forces for PuCoGa$_{5}$ and ferromagnet FePt. Forces were evaluated for the structures with Pu and Pt atoms shifted from their equilibrium positions: $\mathbf{t}_{Pu}=0.02\mathbf{C}$ and $\mathbf{t}_{Pt}=1/2\mathbf{A}+1/2\mathbf{B}+0.52\mathbf{C}$ correspondingly. Free energy corresponding to these perturbed structures is denoted as F(0) in the table. In order to evaluate the forces numerically, two additional small distortions relative to already perturbed structure were considered. Plutonium positions for these distorted structures were $\mathbf{t}_{Pu}=(0.02\pm\Delta)\mathbf{C}$ with $\Delta=0.0004$, and platinum positions were: $\mathbf{t}_{Pt}=1/2\mathbf{A}+1/2\mathbf{B}+(0.52\pm\Delta)\mathbf{C}$ with $\Delta=0.0005$. Values of forces are given for Pu and Pt atoms correspondingly. Total number of \textbf{k}-points in the Brillouin zone was 486 and 6000 for PuCoGa$_{5}$ and FePt correspondingly. Free energies are given relative to the constants -81550 Ry for PuCoGa$_{5}$ and -39414 Ry for FePt.} \label{others}
\begin{center}
\begin{tabular}{@{}c c c c c c c} & &   &  & & Numerical& Mismatch    \\
Solid & F($-\Delta$)&   F(0)& F($+\Delta$) & Force & force &    (\%) \\
\hline\hline
PuCoGa$_{5}$  &-0.9305437   &-0.9303248 &-0.9301012   &-43.0625&-43.0829  &	0.05 \\
FePt  &-0.5907798  &-0.5907130 &-0.5906452  & -18.756 & -18.843	&0.46 \\
\end{tabular}
\end{center}
\end{table}

Principal results of this work, demonstrating the accuracy of the calculated forces, are collected in Table \ref{U_conv} (for $\alpha$-uranium) and in Table \ref{others} (for PuCoGa$_{5}$ and FePt). The tables also include the free energies which were used for the numerical evaluation of the forces. For the numerical differentiation we used three point formula $F'(0)=\frac{F(\Delta)-F(-\Delta)}{2\Delta}$ with $\Delta$ specified in Tables \ref{U_conv} and \ref{others}. Let us first discuss $\alpha$-uranium. As one can see from the Table \ref{U_conv}, the deviation of the calculated forces from the numerical ones is very small (about 0.05\%), which demonstrates high accuracy of the implementation. It is interesting, that the deviation is essentially independent on the sampling of the Brillouin zone. When the number of \textbf{k}-points increases, the forces and the numerical forces change slightly, but their difference is almost constant. This fact supports the robustness of the implementation. One has to mention that the forces evaluated by numerical differentiation are not exact. Not only they depend on the step $\Delta$ in the above formula (though this dependence was rather small in all cases considered in this work), but the free energies corresponding to the shift by $+\Delta$ and $-\Delta$ are subjects to different numerical errors. For instance, MT radii can be dependent on $\Delta$ (as they were in this work). Thus, comparison of the directly and numerically evaluated forces should not be considered as a test of the directly evaluated forces against the numerical ones but, rather, as a test of the consistency of the algorithms involved in both, energies and forces.

\begin{table}[t]
\caption{Calculated components of the forces exerted on all atoms. Forces correspond to the structure distorted from the equilibrium as described in tables \ref{U_conv} and \ref{others}. Surface (kinetic) term corresponds to the contribution to the force from the discontinuity of the kinetic energy at the MT surface as it is specified in the eq. (\ref{forc3}). Surface (other) include the contributions from all other discontinuities as it is specified in the eq. (\ref{forc4}). Group of four Gallium atoms (in undisturbed structure) becomes splitted in two groups (2 atoms each) with slightly different forces which are separated by the slash in the table.} \label{compon}
\begin{center}
\begin{tabular}{@{}c c c c c c c c}
Structure &$\alpha$-U &\multicolumn{4}{c}{PuCoGa$_{5}$}    &\multicolumn{2}{c}{FePt} \\
Atom & U&Pu&Co & Ga(4) & Ga(1) &Fe &Pt \\
\hline\hline
Hellmann-Feynman  &399.284   &-996.116 &-89.242   &-381.61/24.874&67.072  &-20.333	&28.551 \\
Pulay(core)  &-440.654   &954.318 & 64.177  &242.442/-2.174&-38.228  &28.868	&-46.512 \\
Pulay(valence)  &-8.712   &-9.682 &29.136   &134.506/-0.482&-27.138  &25.793	&-14.979 \\
Surface(kinetic)  &-0.883   &8.311 &-2.408   &14.172/-12.201&0.734  &-15.713	&14.208 \\
Surface(other)  &0.054   &-0.006 &-0.0004   &0.0005/-0.0003&-0.0009  &0.086	&-0.024 \\
Total  &-50.912   &-43.063 &1.663   &9.51/10.016&2.44  &18.702	&-18.756 \\
Sum of totals (drift)  & 0  &\multicolumn{4}{c}{-0.020} &\multicolumn{2}{c}{-0.054} \\
\end{tabular}
\end{center}
\end{table}

Whereas the accuracy of the basis set used in the calculations for $\alpha$-U was specially studied in our previous work \cite{arx_2012_04992}, basis sets used in the calculations for PuCoGa$_{5}$ and FePt have not been specifically tuned to reach very high accuracy. This, most likely, explains slightly bigger mismatch between the directly and numerically evaluated forces in these two cases. Nevertheless, the mismatch is small (about 0.5\%) and acceptable in most situations. It demonstrates, that the algorithm of the force evaluation is accurate enough not only when one sort of atoms is present ($\alpha$-U) but also in materials with different atoms (PuCoGa$_{5}$) and in materials with a long range magnetic order (FePt).

Finally, table \ref{compon} presents the components of the forces for all solids studied in the work. First interesting observation is that Hellmann-Feynman and Pulay (core part) are far the biggest components (especially for actinide atoms) and they cancel each other in considerable degree. Both of them come from the inner part of the MT spheres stressing the importance of correct numerical description in that area of the unit cell. Second observation is that the kinetic surface term prevails (considerably) over all other surface terms. This fact essentially supports the approximation accepted in the Ref. \cite{prb_43_6411} where only kinetic operator discontinuity was taken into account. Careful analysis of all other discontinuities performed by authors of Ref. \cite{prb_91_035105} had shown, however, the importance of these additional terms in enhancing the accuracy of the calculated forces. Thus, the other surface contributions were kept in this work and, as one can see, they are not negligible despite their relative smallness.

\section*{Conclusions}
\label{concl}

In conclusion, a formulation of the atomic forces evaluation in the framework of the relativistic density functional theory was given. It is formulated for the APW/LAPW family of basis sets with a flexible inclusion of different kind of the local orbitals (lo, HDLO, HELO). The method has been implemented in the computer code FlapwMBPT and successfully applied to the atomic forces evaluation in $\alpha$-U, PuCoGa$_{5}$, and FePt. The formulation of the forces evaluation  in the fully relativistic framework brings in an opportunity to study, for instance, the phonon spectra in actinide materials with greater reliability than it was previously available with scalar-relativistic approaches. It can also increase the efficiency of the calculations. For example, recent successful study of the phonon spectra in $\alpha$-Plutonium \cite{sr_9_18682} used the small-displacement method \cite{cpc_180_2622} and numerical differentiation of the total energies for the force evaluation. The study could be done easier with the direct evaluation of the forces.

\section*{Acknowledgments}
\label{acknow}
This work was   supported by the U.S. Department of energy, Office of Science, Basic
Energy Sciences as a part of the Computational Materials Science Program.

\bibliographystyle{elsarticle-num}


\section{Introduction}

	Deep Neural Networks (DNNs) have shown remarkable success in many computer vision tasks. Despite the high performance in server-based DNNs powered by cutting-edge parallel computing hardware, most state-of-the-art architectures are not yet ready to be deployed on mobile devices due to the limitations on computational capacity, memory and power.

	To address this problem, many network compression and acceleration methods have been proposed. Pruning based methods \cite{han2015learning,He_2017_ICCV,Liu2017learning,Luo2017ThiNetAF} explore the sparsity in weights and filters. Quantization based methods \cite{han2015learning,Zhou2016Incremental,Courbariaux2016BinaryNet,rastegari2016xnor,Xu2018DeepNN} reduce the bit-width of network parameters. Low-rank decomposition \cite{Denton2014Exploiting,jaderberg2014speeding,Guo2018Network,Wen2017Coordinating,Alvarez2017Compression} minimizes the channel-wise and spatial redundancy by decomposing the original network into a compact one with low-rank layers. In addition, efficient architectures \cite{Sandler2018MobileNetV2,Ma2018ShuffleNet} are carefully designed to facilitate mobile deployment of deep neural networks. Different from precedent works, this paper proposes a novel approach to design low-rank networks.

	Low-rank networks can be trained directly from scratch. However, it is difficult to obtain satisfactory results for several reasons.
	(1)~\textit{Low capacity:}
	compared with the original full rank network, the capacity of a low-rank network is limited, which causes difficulties in optimizing its performances.
	(2)~\textit{Deep structure:}
	low-rank decomposition typically doubles the number of layers in a network. The additional layers make numerical optimization much more vulnerable to gradients explosion and/or vanishing.
	(3)~\textit{Heuristic rank selection:}
	the rank of decomposed network is often  chosen as a hyperparameter based on pre-trained
	networks; this may not be the optimal rank for the network trained from scratch.

	Alternatively, several previous works \cite{zhang2016accelerating,Guo2018Network,jaderberg2014speeding} attempted to decompose pre-trained models in order to get initial low-rank networks. However, the heuristically imposed low-rank could incur huge accuracy loss and network retraining is needed to recover the performance of the original network as much as possible. Some attempts were made to use sparsity regularization \cite{Wen2017Coordinating,chen2015compressing} to constrain the network into a low-rank space. Though sparsity regularization reduces the error incurred by decomposition to some extent, performance still degrades rapidly when compression rate increases.

	This paper is an extension of \cite{xu2018trained}. In this paper, we propose a new method, namely Trained Rank Pruning (TRP), for training low-rank networks. We embed the low-rank decomposition into the training process by gradually pushing the weight distribution of a well functioning network into a low-rank form, where all parameters of the original network are kept and optimized to maintain its capacity. We also propose a stochastic sub-gradient descent optimized nuclear regularization that further constrains the weights in a low-rank space to boost the TRP.
	 The proposed solution is illustrated in Fig.~\ref{fig.1}.

	Overall, our contributions are summarized below.
	\begin{enumerate}
		\setlength\itemsep{-0em}
		\item A new training method called the TRP is presented by explicitly embedding the low-rank decomposition into the network training;
		\item A nuclear regularization is optimized by stochastic sub-gradient descent to boost the performance of the TRP;
		\item Improving inference acceleration and reducing approximation accuracy loss in both channel-wise and spatial-wise decomposition methods.
	\end{enumerate}

	\begin{figure*}[h]
		\centering
		\includegraphics[width=5.5in]{framework.pdf}
		\caption{The training of TRP consists of two parts as illustrated in (a) and (b). (a) one normal iteration with forward-backward broadcast and weight update. (b) one training iteration inserted by TRP, where the low-rank approximation is first applied on filters before convolution. During backward propagation, the gradients are directly added on low-rank filters and the original weights are substituted by updated low-rank filters. (b) is applied once every $m$ iterations (\textit{i.e.} when gradient update iteration $t=zm, z=0, 1, 2, \cdots$), otherwise (a) is applied.}\label{fig.1}
	\end{figure*}
	\section{Related Works}

	A lot of works have been proposed to accelerate the inference process of deep neural networks. Briefly, these works could be categorized into three main categories: quantization, pruning, and low-rank decomposition.

	\textbf{Quantization} Weight quantization methods include training a quantized model from scratch \cite{chen2015compressing,Courbariaux2016BinaryNet,rastegari2016xnor} or converting a pre-trained model into quantized representation \cite{Zhou2016Incremental,Han2015DeepCC,Xu2018DeepNN}. The quantized weight representation includes binary value \cite{rastegari2016xnor,Courbariaux2016BinaryNet} or hash buckets \cite{chen2015compressing}.
	Note that our method is inspired by the scheme of combining quantization with training process, \textit{i.e.} we embed the low-rank decomposition into training process to explicitly guide the parameter to a low-rank form.

	\textbf{Pruning} Non-structured and structured sparsity are introduced by pruning. \cite{han2015learning} proposes to prune unimportant connections between neural units with small weights in a pre-trained CNN. \cite{Wen2016LearningSS} utilizes group Lasso strategy to learn the structure sparsity of networks. \cite{Liu2017learning} adopts a similar strategy by explicitly imposing scaling factors on each channel to measure the importance of each connection and dropping those with small weights. In \cite{He_2017_ICCV}, the pruning problem is formulated as a data recovery problem. Pre-trained filters are re-weighted by minimizing a data recovery objective function. Channels with smaller weight are pruned. \cite{Luo2017ThiNetAF} heuristically selects filters using  change of next layer's output as a criterion.

	\textbf{Low-rank decomposition} Original models are decomposed into compact ones with more lightweight layers
	\cite{jaderberg2014speeding} considers both the spatial-wise and channel-wise redundancy and proposes decomposing a filter into two cascaded asymmetric filters. \cite{zhang2016accelerating} further assumes the feature map lie in a low-rank subspace and decompose the convolution filter into $k\times k$ followed by $1\times 1$ filters via SVD. \cite{Guo2018Network} exploits the low-rank assumption of convolution filters and decompose a regular convolution into several depth-wise and point-wise convolution structures. Although these works achieved notable performance in network compression, all of them are based on the low-rank assumption. When such assumption is not completely satisfied, large prediction error may occur

	Alternatively, some other works \cite{Wen2017Coordinating,Alvarez2017Compression} implicitly utilize sparsity regularization to direct the neural network training process to learn a low-rank representation. Our work is similar to this low-rank regularization method. However, in addition to appending an implicit regularization during training, we impose an explicit sparsity constraint in our training process and prove that our approach can push the weight distribution into a low-rank form quite effectively.

	\section{Methodology}
	\subsection{Preliminaries}
	Formally, the convolution filters in a layer can be denoted by a tensor $W \in \mathbb{R}^{n \times c \times k_w \times k_h}$, where $ n $ and $ c $ are the number of filters and input channels, $ k_h $ and $ k_w $ are the height and width of the filters. An input of the convolution layer $ F_i \in  \mathbb{R}^{c \times x \times y}$ generates an output as $ F_o = W * F_i $. Channel-wise correlation \cite{zhang2016accelerating} and spatial-wise correlation \cite{jaderberg2014speeding} are explored to approximate convolution filters in a low-rank space. In this paper, we focus on these two decomposition schemes.

	However, unlike the previous works, we propose a new training scheme TRP to obtain a low-rank network without re-training after decomposition.

	\subsection{Trained Rank Pruning}

	Trained Rank Pruning (TRP) is motivated by the strategies of training quantized nets.
	One of the gradient update schemes to train quantized networks from scratch \cite{Li2017TrainingQN} is

	\begin{equation}\label{equ.1}
	w^{t+1}=Q(w^t-\alpha \triangledown f(w^t))
	\end{equation}
	where $Q(\cdot)$ is the quantization function, $w^t$ denote the parameter in the $t^{th}$ iteration. Parameters are quantized by $Q(\cdot)$ before updating the gradients.

	In contrast, we propose a simple yet effective training scheme called Trained Rank Pruning (TRP) in a periodic fashion:

	\begin{equation}\label{equ.2}
	\begin{split}
	W^{t+1}&=\left \{
	\begin{array}{c}
	W^t - \alpha \triangledown f(W^t)  \quad t \% m \neq 0 \\
	T^z - \alpha \triangledown f(T^z)  \quad t \% m = 0
	\end{array} \right. \\
	T^{z}&=\mathcal{D}(W^{t}), \quad z = t / m
	\end{split}
	\end{equation}
	where $\mathcal{D}({\cdot})$ is a  low-rank  tensor approximation operator, $\alpha$ is the learning rate, $t$ indexes the iteration and $z$ is the iteration of the operator $\mathcal{D}$ , with $m$ being the period for the low-rank approximation.

	At first glance, this TRP looks very simple. An immediate concern arises: can the iterations guarantee the rank of the parameters converge, and more importantly would not increase when they are updated in this way? A positive answer (see Theorem 2) given in our theoretical analysis will certify the legitimacy of this algorithm.

	For the network quantization, if the gradients are smaller than the quantization, the gradient information would be totally lost and become zero. However, it will not happen in TRP because the low-rank operator is applied on the weight tensor. Furthermore, we apply low-rank approximation every $m$ SGD iterations. This saves training time to a large extend. As illustrated in Fig.~\ref{fig.1}, for every $m$ iterations, we perform low-rank approximation on the original filters, while gradients are updated on the resultant low-rank form. Otherwise, the network is updated via the normal SGD.
	Our training scheme could be combined with any low-rank operators. In the proposed work, we choose the low-rank techniques proposed in \cite{jaderberg2014speeding} and \cite{zhang2016accelerating}, both of which transform the 4-dimensional filters into 2D matrix and then apply the truncated singular value decomposition (TSVD). The SVD of matrix ${W}^t$ can be written as:
	\begin{equation}\label{equ.3}
	W^t=\sum_{i=1}^{rank(W^t)}\sigma_i\cdot U_i\cdot (V_i)^T
	\end{equation}
	where $\sigma_i$ is the singular value of $W^t$ with $\sigma_1\geq \sigma_2 \geq \cdots \geq \sigma_{rank(W^t)}$, and $U_i$ and $V_i$ are the singular vectors. The parameterized TSVD($W^t;e$) is to find the smallest integer $k$ such that
	\begin{equation}\label{equ.4}
	\sum_{j=k+1}^{rank(W^t)}(\sigma_j)^2 \leq \quad e \sum_{i=1}^{rank(W^t)}(\sigma_i)^2
	\end{equation}
	where $e$ is a pre-defined hyper-parameter of the energy-pruning ratio, $e \in(0 , 1)$.

	After truncating the last $n-k$ singular values, we transform the low-rank 2D matrix  back to 4D tensor. Compared with directly training low-rank structures from scratch, the proposed TRP has following advantages.

	(1) Unlike updating the decomposed filters independently of the network training in literature \cite{zhang2016accelerating,jaderberg2014speeding}, we update the network directly on the original 4D shape of the decomposed parameters, which enable jointly network decomposition and training by preserving its discriminative capacity as much as possible

	(2) Since the gradient update is  performed based on the original network structure, there will be no exploding and vanishing gradients problems caused by additional layers.

	(3) The rank of each layer is automatically selected during the training. We will prove a theorem certifying the rank of network weights convergence and would not increase in section \ref{sec:thm}.

	\subsection{Nuclear Norm Regularization}
	Nuclear norm is widely used in matrix completion problems. Recently, it is introduced to constrain the network into low-rank space during the training process \cite{Alvarez2017Compression}.
	\begin{equation}\label{equ.5}
	\min \left\{  f\left(x; w\right)+\lambda \sum_{l=1}^L||W_l||_* \right\}
	\end{equation}
	where $f(\cdot)$ is the objective loss function, nuclear norm $||W_l||_*$ is defined as $||W_l||_*=\sum_{i=1}^{rank(W_l)}\sigma_l^i$, with $\sigma_l^i$ the singular values of $W_l$. $\lambda$ is a hyper-parameter setting the influence of the nuclear norm. In \cite{Alvarez2017Compression} the proximity operator is applied in each layer independently to solve Eq.~(\ref{equ.5}). However, the proximity operator is split from the training process and doesn't consider the influence within layers.

	In this paper, we utilize stochastic sub-gradient descent \cite{Avron2012EfficientAP} to optimize nuclear norm regularization in the training process. Let $W=U\Sigma V^T$ be the SVD of $W$ and let $U_{tru}, V_{tru}$ be $U, V$ truncated to the first $rank(W)$ columns or rows, then $U_{tru}V_{tru}^T$ is the sub-gradient of $||W||_*$ \cite{watson1992characterization}. Thus, the sub-gradient of  Eq.~(\ref{equ.5}) in a layer is
	\begin{equation}\label{equ.6}
	\triangledown f+\lambda U_{tru}V_{tru}^T
	\end{equation}

	The nuclear norm and loss function are optimized simultaneously during the training of the networks and can further be combined with the proposed TRP.

	\subsection{Theoretic Analysis}\label{sec:thm}
	In this section, we analyze the rank convergence of TRP from the perspective of matrix perturbation theory \cite{stewart1990matrix}. We prove that rank in TRP is monotonously decreasing, \textit{i.e.,} the model gradually converges to a more sparse model.

	Let A be an $m\times n$ matrix, without loss of generality, $m\geq n$.

	$\Sigma = diag\left(\sigma_1,\cdots,\sigma_n\right)$ and $\sigma_1\geq \sigma_2\geq\cdots\geq\sigma_n$. $\Sigma$ is the diagonal matrix composed by all singular values of $A$.
	Let $\widetilde{A}=A+E$ be a perturbation of $A$, and $E$ is the noise matrix.

	$\widetilde{\Sigma} = diag\left(\widetilde{\sigma}_1,\cdots,\widetilde{\sigma}_n\right)$ and $\widetilde{\sigma}_1 \geq \widetilde{\sigma}_2 \geq \cdots \geq \widetilde{\sigma}_n$. $\widetilde{\sigma}_i$ is the singular values of $\widetilde{A}$.
	The basic perturbation bounds for the singular values of a matrix are given by

	\begin{theorem}\label{theorem1}
		Mirsky's theorem \cite{mirsky1960symmetric}:
		\begin{equation}\label{equ.9}
		\sqrt{\sum_{i}|\widetilde{\sigma}_i-\sigma_i|^2}\leq||E||_F\\
		\end{equation}
	\end{theorem}

	where $||\cdot||_F$ is the Frobenius norm. Then the following corollary can be inferred from Theorem \ref{theorem1},

	\begin{corollary}\label{corollary1}
		Let $B$ be any $m\times n$ matrix of rank not greater than $k$, \textit{i.e.} the singular values of B can be denoted by $\varphi_1 \geq \cdots \geq \varphi_k \geq 0$ and $\varphi_{k+1}=\cdots=\varphi_n=0$. Then
		\begin{equation}\label{equ.10}
		||B-A||_F \geq \sqrt{\sum_{i=1}^{n}|\varphi_i-\sigma_i|^2}\geq \sqrt{\sum_{j=k+1}^{n}\sigma_{j}^2}\\
		\end{equation}
	\end{corollary}

	Below, we will analyze the training procedure of the proposed TRP. Note that $W$ below are all transformed into 2D matrix. In terms of Eq.~(\ref{equ.2}), the training process between two successive TSVD operations can be rewritten as Eq. (\ref{equ.11})
	\begin{equation}\label{equ.11}
	\begin{split}
	W^t&=T^z=TSVD(W^t;e)\\
	W^{t+m} &= T^z-\alpha \Sigma_{i=0}^{m-1}\triangledown f(W^{t+i})\\
	T^{z+1}&=TSVD(W^{t+m};e)
	\end{split}
	\end{equation}
	where $W^t$ is the weight matrix in the $t$-th iteration. $T^z$ is the weight matrix  after applying TSVD over $W^t$. $\triangledown f(W^{t+i})$ is the gradient back-propagated during the $(t+i)$-th iteration. $e\in\left(0,1\right)$ is the predefined energy threshold. Then we have following theorem.
	\begin{theorem}\label{thm:2}
		Assume that $||\alpha \triangledown f||_F$ has an upper bound G, if  $G<\frac{\sqrt{e}}{m}||W^{t+m}||_F$, then $rank(T^z)\geq rank(T^{z+1})$.
	\end{theorem}

	\begin{proof}

		We denote $\sigma^t_{j}$ and $\sigma^{t+m}_{j}$ as the singular values of $W^t$ and $W^{t+m}$ respectively. Then at the $t$-th iteration, given the energy ratio threshold $e$, the TSVD operation tries to find the singular value index $k \in [0, n-1]$ such that :

		\begin{equation}\label{equ.12}
		\begin{split}
		\sum_{j=k+1}^{n}\left(\sigma^{t}_{j}\right)^2&<e||W^t||_F^2\\
		\sum_{j=k}^{n}\left(\sigma^t_{j}\right)^2&\geq e||W^t||_F^2\\
		\end{split}
		\end{equation}

		In terms of Eq.~(\ref{equ.12}), $T^z$ is a $k$ rank matrix, i.e, the last $n-k$ singular values of $T^z$ are equal to $0$. According to Corollary \ref{corollary1}, we can derive that:
		\begin{equation}\label{equ.13}
		\begin{split}
		||W^{t+m}-T^z||_F& = ||\alpha \sum_{i=0}^{m-1}\triangledown f^{t+i}||_F\\
		&\geq \sqrt{\sum_{j=k+1}^{n}\left(\sigma^{t+m}_{j}\right)^2} \\
		\end{split}
		\end{equation}

		Given the assumption $G < \frac{\sqrt{e}}{m} ||W^{t+m}||_F $, we can get:
		\begin{equation}\label{equ.14}
		\begin{split}
		\frac{\sqrt{\sum_{j=k+1}^{n}\left(\sigma^{t+m}_{j}\right)^2}}{||W^{t+m}||_F}&\leq \frac{||\alpha \sum_{i=0}^{m-1}\triangledown f^{t+i}||_F}{||W^{t+m}||_F}\\
		&\leq \frac{\sum_{i=0}^{m-1}||\alpha \triangledown f^{t+i}||_F}{||W^{t+m}||_F}\\
		&\leq \frac{mG}{||W^{t+m}||_F}< \sqrt{e} \\
		\end{split}
		\end{equation}

		Eq.~(\ref{equ.14}) indicates that
		since the perturbations of singular values are bounded by the parameter gradients, if we properly select the TSVD energy ratio threshold $e$, we could guarantee that if $n-k$ singular values are pruned by previous TSVD iteration, then before the next TSVD, the energy for the last $n-k$ singular values is still less than the pre-defined energy threshold $e$. Thus TSVD should keep the number of pruned singular values or drop more to achieve the criterion in Eq.~(\ref{equ.12}), consequently a weight matrix with lower or same rank is obtained, \textit{i.e.} $Rank(T^{z})\geq Rank(T^{z+1})$. We further confirm our analysis about the variation of rank distribution in Section \ref{sec:exp}.
	\end{proof}

	\begin{table}[htb]
		\begin{center}
			\begin{tabular}{|l|c|c|c|}
				\hline
				Model & Top 1 ($\%$) & Speed up\\
				\hline\hline
				R-20 (baseline)&91.74&1.00$\times$\\
				\hline
				R-20 (TRP1)&90.12&1.97$\times$\\
				R-20 (TRP1+Nu)&\textbf{90.50}&\textbf{2.17$\times$}\\
				R-20 (\cite{zhang2016accelerating})&88.13&1.41$\times$\\
				\hline
				R-20 (TRP2)&90.13&2.66$\times$\\
				R-20 (TRP2+Nu)&\textbf{90.62}&\textbf{2.84$\times$}\\
				R-20 (\cite{jaderberg2014speeding})&89.49&1.66$\times$\\
				\hline
				\hline
				R-56 (baseline)&93.14&1.00$\times$\\
				\hline
				R-56 (TRP1)&\textbf{92.77}&2.31$\times$\\
				R-56 (TRP1+Nu)&91.85&\textbf{4.48$\times$}\\
				R-56 (\cite{zhang2016accelerating})&91.56&2.10$\times$\\
				\hline
				R-56 (TRP2)&\textbf{92.63}&2.43$\times$\\
				R-56 (TRP2+Nu)&91.62&\textbf{4.51$\times$}\\
				R-56 (\cite{jaderberg2014speeding})&91.59&2.10$\times$\\
				\hline
				R-56 \cite{He_2017_ICCV}&91.80&2.00$\times$\\
				R-56 \cite{li2016pruning}&91.60&2.00$\times$\\
				\hline
			\end{tabular}
		\end{center}
		\caption{Experiment results on CIFAR-10. "R-`` indicates ResNet-. }\label{tab1}\vspace*{-0.4cm}
	\end{table}

	\section{Experiments}\label{sec:exp}

	\subsection{Datasets and Baseline}
	We evaluate the performance of TRP scheme on two common datasets, CIFAR-10 \cite{AlexCifar10} and ImageNet \cite{Deng2009ImageNetAL}. The CIFAR-10 dataset consists of  colored natural images with $32\times 32$
	resolution and has totally 10 classes. The ImageNet dataset consists of 1000 classes of images for recognition task.
	For both of the datasets, we adopt ResNet \cite{He2016DeepRL} as our baseline model since it is widely used in different vision tasks
	We use ResNet-20, ResNet-56 for CIFAR-10 and ResNet-18, ResNet-50 for ImageNet. For evaluation metric, we adopt top-1 accuracy on CIFAR-10 and top-1, top-5 accuracy on ImageNet. To measure the acceleration performance, we compute the FLOPs ratio between baseline and decomposed models to obtain the final speedup rate. Wall-clock CPU and GPU time is also compared. Apart from the basic decomposition methods, we compare the performance with other state-of-the-art acceleration algorithms~\cite{He_2017_ICCV,li2016pruning,Luo2017ThiNetAF,Zhou_2019_ICCV}.

	\begin{table}[htb]
		\begin{center}
			\begin{tabular}{|l|c|c|c|}
				\hline
				Method &Top1($\%$)& Top5($\%$)& Speed up\\
				\hline\hline
				Baseline&69.10&88.94&1.00$\times$\\
				\hline
				TRP1 &\textbf{65.46}&\textbf{86.48} &1.81$\times$\\
				TRP1+Nu&65.39&86.37&\textbf{2.23$\times$}\\
				\cite{zhang2016accelerating}\footnotemark[1] &-& 83.69&1.39$\times$ \\
				\cite{zhang2016accelerating}&63.10&84.44&1.41$\times$\\
				\hline
				TRP2 &\textbf{65.51}&\textbf{86.74} &$2.60\times$\\
				TRP2+Nu&65.34&86.61&\textbf{3.18}$\times$\\
				\cite{jaderberg2014speeding}&62.80&83.72&2.00$\times$\\
				\hline
			\end{tabular}
		\end{center}
		\caption{Results of ResNet-18 on ImageNet. }\label{tab2}\vspace*{-0.2cm}
	\end{table}

	\begin{table}[htb]
		\begin{center}
			\begin{tabular}{|l|c|c|c|}
				\hline
				Method &Top1($\%$)& Top5($\%$) & Speed up\\
				\hline\hline
				Baseline&75.90&92.70&1.00$\times$\\
				\hline
				TRP1+Nu&72.69&91.41&\textbf{2.30}$\times$\\
				TRP1+Nu&\textbf{74.06}&\textbf{92.07}&1.80$\times$\\
				\cite{zhang2016accelerating}&71.80&90.2&1.50$\times$\\
				\cite{He_2017_ICCV}&-&90.80&2.00\\
				\cite{Luo2017ThiNetAF}&72.04&90.67&1.58\\
				\cite{luo2018thinet}&72.03&90.99&2.26\\
				\cite{Zhou_2019_ICCV}&71.50&90.20&2.30\\
				\hline
			\end{tabular}
		\end{center}
		\caption{Results of ResNet-50 on ImageNet.}\label{tab3}\vspace*{-0.4cm}
	\end{table}

	\subsection{Implementation Details}
	We implement our TRP scheme with NVIDIA 1080 Ti GPUs. For training on CIFAR-10, we start with base learning rate of $0.1$ to train 164 epochs and degrade the value by a factor of $10$ at the $82$-th and $122$-th epoch. For ImageNet, we directly finetune the model with TRP scheme from the pre-trained baseline with learning rate $0.0001$ for 10 epochs. We adopt SGD solver to update weight and set the weight decay value as $10^{-4}$ and momentum value as $0.9$. The accuracy improvement enabled by data dependent decomposition vanishes after fine-tuning. So we simply
	adopt the retrained data independent decomposition as our basic methods.

	\footnotetext[1]{the implementation of \cite{Guo2018Network}}

	\subsection{Results on CIFAR-10}
	\textbf{Settings.} Experiments on channel-wise decomposition (TRP1) and spatial-wise decomposition (TRP2) are both considered. The TSVD energy threshold in TRP and TRP+Nu is $0.02$ and the nuclear norm weight $\lambda$ is set as $0.0003$. We decompose both the $1\times 1$ and $3\times 3$ layers in ResNet-56.

	\textbf{Results.} As shown in Table~\ref{tab1}, for both spatial-wise and channel-wise decomposition, the proposed TRP outperforms basic methods~\cite{zhang2016accelerating,jaderberg2014speeding} on ResNet-20 and ResNet-56. Results become even better when nuclear regularization is used. For example, in the channel-wise decomposition (TRP2) of ResNet-56, results of TRP combined with nuclear regularization can even achieve $2\times$ speed up rate than \cite{zhang2016accelerating} with same accuracy drop. TRP also outperforms filter pruning~\cite{li2016pruning} and channel pruning~\cite{He_2017_ICCV}. The channel decomposed TRP trained ResNet-56 can achieve $92.77\%$ accuracy with $2.31\times$ acceleration, while \cite{He_2017_ICCV} is $91.80\%$ and \cite{li2016pruning} is $91.60\%$. With nuclear regularization, our methods can approximately double the acceleration rate of \cite{He_2017_ICCV} and \cite{li2016pruning} with higher accuracy.

	\begin{figure}[t]
		\centering
		\includegraphics[width=3in]{heatmap.pdf}\\
		\caption{Visualization of rank selection, taken from the res3-1-2 convolution layer in ResNet-20 trained on CIFAR-10.}\label{fig.5}
		\vspace*{-0.4cm}
	\end{figure}

	\subsection{Results on ImageNet}
	\textbf{Settings.} We choose ResNet-18 and ResNet-50 as our baseline models. The TSVD energy threshold $e$ is set as 0.005. $\lambda$ of nuclear norm regularization is 0.0003 for both ResNet-18 and ResNet-50. We decompose both the $3\times 3$ and $1\times 1$ Convolution layers in ResNet-50. TRP1 is the channel-wise decomposition and TRP2 is the spatial-wise decomposition.

	\begin{figure*}[htb]
		\centering
		\subfigure[Channel-wise decomposition]{

			\includegraphics[width=2.8in]{nc2.pdf}}
		\subfigure[Spatial-wise decomposition]{
			\includegraphics[width=2.8in]{vh2.pdf}}
		\caption{Ablation study on ResNet-20. Basic methods are data-independent decomposition methods (channel or spatial) with finetuning.}
		\label{fig.6}\vspace*{-0.4cm}
	\end{figure*}

	\textbf{Results.} The results on ImageNet are shown in Table~\ref{tab2} and Table~\ref{tab3}. For ResNet-18, our method outperforms the basic methods~\cite{zhang2016accelerating,jaderberg2014speeding}. For example, in the channel-wise decomposition, TRP obtains 1.81$\times$ speed up rate with 86.48\% Top5 accuracy on ImageNet which outperforms both the data-driven~\cite{zhang2016accelerating}\textsuperscript{1} and data independent~\cite{zhang2016accelerating} methods by a large margin. Nuclear regularization can increase the speed up rates with the same accuracy.

	For ResNet-50, to better validate the effectiveness of our method, we also compare TRP with pruning based methods. With $1.80\times$ speed up, our decomposed ResNet-50 can obtain $74.06\%$ Top1 and $92.07\%$ Top5 accuracy which is much higher than \cite{Luo2017ThiNetAF}. The TRP achieves $2.23\times$ acceleration which is higher than \cite{He_2017_ICCV} with the same $1.4\%$ Top5 degrade. Besides, with the same $2.30\times$ acceleration rate, our performance is better than \cite{Zhou_2019_ICCV}.

	\subsection{Rank Variation}

	To analyze the variation of rank distribution during training, we further conduct an experiment on the CIFAR-10 dataset with ResNet-20 and extract the weight from the \emph{res3-1-2} convolution layer with channel-wise decomposition as our TRP scheme. After each TSVD, we compute the normalized energy ratio $ER(i)$ for each singular value $\sigma_i$ as Eq. (\ref{eq:energy}).

	\begin{equation}\label{eq:energy}
	ER(i) = \frac{\sigma_i^2}{\sum_{j=0}^{rank(T^z)}\sigma_j^2}
	\end{equation}
	we record for totally $40$ iterations of TSVD with period $m=20$, which is equal to $800$ training iterations, and our energy theshold $e$ is pre-defined as $0.05$.
	Then we visualize the variation of $ER$ in Fig.~\ref{fig.5}. During our training, we observe that the theoretic bound value $\max_t\frac{mG}{||W^t||_F} \approx 0.092 < \sqrt{e} \approx 0.223$, which indicates that our basic assumption in theorem \ref{thm:2} always holds for the initial training stage.

	And this phenomenon is also reflected in Fig.~\ref{fig.5}, at the beginning, the energy distribution is almost uniform w.r.t each singular value, and the number of dropped singular values increases after each TSVD iteration
	and the energy distribution becomes more dense among singular values with smaller index. Finally, the rank distribution converges to a certain point where the smallest energy ratio exactly reaches our threshold $e$ and TSVD will not cut more singular values.

	\subsection{Ablation Study}

	In order to show the effectiveness of different components of our method, we compare four training schemes, basic methods \cite{zhang2016accelerating,jaderberg2014speeding}, basic methods combined with nuclear norm regularization, TRP and TRP combined with nuclear norm regularization. The results are shown in Fig.~\ref{fig.6}. We can have following observations:

	(1) \emph{Nuclear norm regularization} After combining nuclear norm regularization, basic methods improve by a large margin. Since Nuclear norm regularization constrains the filters into low rank space, the loss caused by TSVD is smaller than the basic methods.

	(2) \emph{Trained rank pruning} As depicted in Fig.~\ref{fig.6}, when the speed up rate increases, the performance of basic methods and basic methods combined with nuclear norm regularization degrades sharply. However, the proposed TRP degrades very slowly. This indicates that by reusing the capacity of the network, TRP can learn a better low-rank feature representations than basic methods. The gain of nuclear norm regularization on TRP is not as big as basic methods because TRP has already induced the parameters into low-rank space by embedding TSVD in training process.
	\begin{table}[htb]
		\begin{center}
			\begin{tabular}{|l|c|c|}
				\hline
				Model & GPU time (ms)&CPU time (ms)\\
				\hline\hline
				Baseline&0.45&118.02\\
				\hline
				TRP1+Nu (channel)&0.33&64.75\\
				\hline
				TRP2+Nu (spatial)&0.31&49.88\\
				\hline
			\end{tabular}
		\end{center}
		\caption{Actual inference time per image on ResNet-18.}\label{tab4}
	\end{table}
	\vspace*{-0.2cm}
	\subsection{Runtime Speed up of Decomposed Networks}
	We further evaluate the actual runtime speed up of the compressed Network as shown in Table \ref{tab4}. Our experiment is conducted on a platform with one Nvidia 1080Ti GPU and Xeon E5-2630 CPU. The models we used are the original ResNet-18 and decomposed models by TRP1+Nu and TRP2+Nu.
	From the results, we observe that on CPU our TRP scheme achieves more salient acceleration performance. Overall the spatial decomposition combined with our TRP+Nu scheme has better performance. Because cuDNN is not friendly for $1\times 3$ and $3\times1$ kernels, the actual speed up of spatial-wise decomposition is not as obvious as the reduction of FLOPs.

	\section{Conclusion}

	In this paper, we proposed a new scheme Trained Rank Pruning (TRP) for training low-rank networks. It leverages capacity and structure of the original network by embedding the low-rank approximation in the training process. Furthermore, we propose stochastic sub-gradient descent optimized nuclear norm regularization to boost the TRP. The proposed TRP can be incorporated with any low-rank decomposition method. On CIFAR-10 and ImageNet datasets, we have shown that our methods can outperform basic methods and other pruning based methods both in channel-wise decmposition and spatial-wise decomposition.
	\clearpage
	\bibliographystyle{named}