Mixed models from scratch: REML and random slopes

$$$$$$$$$$In my previous post I stopped right before working out the REML objective function so I’ll start with that, and then I’ll check my work by fitting a random slopes model using my equations and comparing with the output of lme4. All the notation in this post is defined in part 1, but as a quick reminder my model is
$$y=X\beta +Z\gamma +\epsilon$$
with
$$\gamma \sim \mathcal{N}(\mathbf{0},{\sigma }^2{\mathrm{\Omega }}^{-1})\perp \epsilon \sim \mathcal{N}(\mathbf{0},{\sigma }^{2}I)$$
and I’m parameterizing $\mathrm{\Omega }$ with a vector $\theta$.

When I use maximum likelihood for $({\hat{\sigma }}^2,\hat{\theta })$, like I did in the previous post, it can make some things easier, like likelihood ratio tests, but ${\hat{\sigma }}^2$ is known to be biased and underestimate the true variability due to having too large of a denominator. Intuitively this comes from how the sample mean is based on the same data so the data tend to be more clustered around the mean than a denominator of $n$ accounts for.

Restricted maximum likelihood (REML) can be used to help with this. An improper uniform prior is put on $\beta$ and it is integrated out so then the likelihood that is optimized is only a function of ${\sigma }^2$ and $\theta$. The marginal distribution of $y$ after integrating out $\gamma$ is
$$y\mid \beta \sim \mathcal{N}(X\beta ,{\sigma }^2(I-S_{\theta }{)}^{-1})$$
where I’m now including the dependence on $\beta$. This means that the density for REML, which I’ll denote by $p_r$, is
$$\begin{array}{rl}{p}_r(y\mid \theta ,{\sigma }^2)& ={\int }_{{\mathbb{R}}^p}\mathcal{N}(y\mid X\beta ,{\sigma }^2(I-S_{\theta }{)}^{-1})\text{d}\beta \\ & =\frac{|I-S_{\theta }{|}^{1/2}}{(2\pi {\sigma }^2{)}^{n/2}}{\int }_{{\mathbb{R}}^p}\exp (-\frac{1}{2{\sigma }^2}(y-X\beta {)}^T(I-S_{\theta })(y-X\beta ))\text{d}\beta .\end{array}$$
I want to turn this into a quadratic form in $\beta$ so I’ll expand inside the exponential and complete the square. I’ll let $R_{\theta }=I-S_{\theta }$ and $A_{\theta }=(X^T{R}_{\theta }X{)}^{-1}{X}^T{R}_{\theta }$ to make the notation easier ($A_{\theta }$ gives the weighted pseudoinverse of $X$ and I’m using “$R$” for “residual” since this matrix gives the residuals of a generalized ridge regression on $Z$). This gives me
$$(y-X\beta {)}^T{R}_{\theta }(y-X\beta )=y^T{R}_{\theta }y–2{\beta }^T{X}^T{R}_{\theta }y+{\beta }^T{X}^T{R}_{\theta }X\beta .$$
For $v\ne \mathbf{0}$ consider $v^T{X}^T{R}_{\theta }Xv=u^T{R}_{\theta }u$ with $u=Xv$. Since $X$ is full column rank $v\ne \mathbf{0}⟹u\ne \mathbf{0}$, and then $u^T{R}_{\theta }u>0$ since $R_{\theta }$ is PD. This means that $X^T{R}_{\theta }X$ is also PD and therefore invertible. This means I can complete the square to get
$$(y-X\beta {)}^T{R}_{\theta }(y-X\beta )=(\beta –A_{\theta }y{)}^T{X}^T{R}_{\theta }X(\beta –A_{\theta }y)+y^T(R_{\theta }–R_{\theta }X{A}_{\theta })y.$$
Plugging this in,
$$\begin{array}{rl}{p}_r(y\mid \theta ,{\sigma }^2)& =\frac{|R_{\theta }{|}^{1/2}}{(2\pi {\sigma }^2{)}^{n/2}}\exp (-\frac{1}{2{\sigma }^2}{y}^T(R_{\theta }–R_{\theta }X{A}_{\theta })y)\times \\ & {\int }_{{\mathbb{R}}^p}\exp (-\frac{1}{2{\sigma }^2}(\beta –A_{\theta }y{)}^T{X}^T{R}_{\theta }X(\beta –A_{\theta }y))\text{d}\beta .\end{array}$$
The integral is now that of the kernel of a multivariate Gaussian so all together
$$p_r(y\mid \theta ,{\sigma }^2)=\frac{|R_{\theta }{|}^{1/2}}{(2\pi {\sigma }^2{)}^{n/2}}\exp (-\frac{1}{2{\sigma }^2}{y}^T(R_{\theta }–R_{\theta }X{A}_{\theta })y)\cdot \frac{(2\pi {\sigma }^2{)}^{p/2}}{|X^T{R}_{\theta }X{|}^{1/2}}.$$
For interpreting this, I can note that
$$H_R:=X(X^T{R}_{\theta }X{)}^{-1}{X}^T{R}_{\theta }$$
is the projection matrix for a GLS regression of $y$ on $X$ with weights given by $R_{\theta }^{-1}=(I-S_{\theta }{)}^{-1}$ which is exactly the marginal variance of $y$ (after integrating out $\gamma$ but before integrating out $\beta$). Thus the quadratic form in the exponential is
$$y^T{R}_{\theta }(I-H_R)y=⟨(I-S_{\theta })y,(I–H_R)y⟩$$
so this gives the agreement between the residuals of a ridge regression of $y$ on $Z$ and a GLS regression of $y$ on $X$. The density will be small when this is large, which happens when both residuals are large and the angle between them is small. If both residuals are small in norm then $y$ will be more likely, which corresponds to how a more likely $y$ here is in both the column spaces of $X$ and $Z$ (with the appropriate shrinkage and weighting). Comparing this with the marginal likelihood after just integrating our $\gamma$ I can see that now both $X$ and $Y$ appear as column spaces which reflects that both have been integrated out.

Taking logs, I have
$$\begin{array}{rl}{\ell }_r(\theta ,{\sigma }^2\mid y)& =-\frac{1}{2{\sigma }^2}{y}^T(R_{\theta }–R_{\theta }X(X^T{R}_{\theta }X{)}^{-1}{X}^T{R}_{\theta })y–\frac{n-p}{2}\log {\sigma }^2\\ & +\frac{1}{2}\log det{R}_{\theta }–\frac{1}{2}\log det\left(X^T{R}_{\theta }X\right)\end{array}$$
as my REML log likelihood. To get ${\hat{\sigma }}^2$ I have
$$\begin{array}{rl}\frac{\partial {\ell }_r}{\partial {\sigma }^2}& =\frac{1}{2{\sigma }^4}{y}^T(R_{\theta }–R_{\theta }X(X^T{R}_{\theta }X{)}^{-1}{X}^T{R}_{\theta })y–\frac{n-p}{2{\sigma }^2}\stackrel{\text{set}}{=}0\\ & ⟹\frac{1}{n-p}{y}^T(R_{\theta }–R_{\theta }X(X^T{R}_{\theta }X{)}^{-1}{X}^T{R}_{\theta })y={\hat{\sigma }}^2.\end{array}$$
When I don’t use REML I have a coefficient of $\frac{1}{n}$ here so this $\frac{1}{n-p}$ is how I’m avoiding underestimating the variability.

I can now profile out ${\sigma }^2$ from ${\ell }_r$ to get ${\ell }_{rp}$ which is given by
$${\ell }_{rp}(\theta \mid y)=–\frac{n-p}{2}\log {\hat{\sigma }}_{\theta }^2+\frac{1}{2}\log det{R}_{\theta }–\frac{1}{2}\log det\left(X^T{R}_{\theta }X\right)$$
(up to a constant). Optimizing this gives me $\hat{\theta }$ and then I can get $\hat{\beta }$ and ${\hat{\sigma }}^2$ and the rest proceeds as before.

Application: random slopes

I’ll now check these equations (from my previous post and this one) by fitting a random slopes model on the sleepstudy dataset and comparing the results with those from lme4. I won’t be worrying about efficiency or numerical stability so I’ll be explicitly inverting matrices and all that.

For a random slopes model I have $m$ groups and I want to fit a model with a group-level intercept and slope term. I’ll have $n_i$ units in group $i$ for $n=\sum _{i=1}^m{n}_i$ observations in all. The slope data for group $i$ will be collected into ${\mathbf{x}}_i\in {\mathbb{R}}^{n_i}$. I’ll use $x_{ij}=({\mathbf{x}}_i{)}_j$ to simplify my notation. In scalar form my model is
$$y_{ij}={\beta }_0+{\alpha }_i+({\beta }_1+{\eta }_i)x_{ij}+{\epsilon }_{ij}$$
with ${\alpha }_i$ as the random intercept, ${\eta }_i$ as the random slope, and $j=1,\dots ,n_i$ indexing units within a group. Written in matrix notation I have
$$y=X\beta +Z\gamma +\epsilon$$
with
$$X=\left(\begin{array}{cc}{\mathbf{1}}_{n_1}& {\mathbf{x}}_1\\ ⋮& ⋮\\ {\mathbf{1}}_{n_m}& {\mathbf{x}}_m\end{array}\right)\in {\mathbb{R}}^{n\times 2}$$
giving the features for the fixed-effects part, and
$$Z=\left(\begin{array}{cc}\underset{i=1}{\overset{m}{⨁}}{\mathbf{1}}_{n_i}& \underset{i=1}{\overset{m}{⨁}}{\mathbf{x}}_i\end{array}\right)\in {\mathbb{R}}^{n\times (2m)}.$$
My random effects are $\gamma =(\genfrac{}{}{0ex}{}{\alpha }{\eta })$ so the multiplication $Z\gamma$ picks out the correct elements of $\alpha$ and $\eta$, and multiplies that element of $\eta$ by the right $x_{ij}$ value.

For the distributions of my random effects and error I’ll have
$$\gamma =(\genfrac{}{}{0ex}{}{\alpha }{\eta })\sim \mathcal{N}((\genfrac{}{}{0ex}{}{\mathbf{0}}{\mathbf{0}}),\left(\begin{array}{cc}{\sigma }_{\alpha }^{2}I& \rho {\sigma }_{\alpha }{\sigma }_{\eta }I\\ \rho {\sigma }_{\alpha }{\sigma }_{\eta }I& {\sigma }_{\eta }^{2}I\end{array}\right))\perp \epsilon \sim \mathcal{N}(\mathbf{0},{\sigma }^{2}I)$$
so ${\sigma }_{\alpha }^2$ and ${\sigma }_{\eta }^2$ are the marginal variances of $\alpha$ and $\eta$ respectively, and $\rho$ is the correlation between ${\alpha }_i$ and ${\gamma }_i$. $\rho$ is useful because it allows the predictor $x$ to be translated without changing the model (Bates et al. mention this on page 7 of the lme4 manual; I’ll show later why this is the case).

In the previous post I wanted to have $\mathrm{Var}(\gamma )={\sigma }^2{\mathrm{\Omega }}^{-1}$ which means in this case I have
$${\mathrm{\Omega }}^{-1}=\left(\begin{array}{cc}{\sigma }_{\alpha }^2/{\sigma }^{2}I& \rho {\sigma }_{\alpha }{\sigma }_{\eta }/{\sigma }^{2}I\\ \rho {\sigma }_{\alpha }{\sigma }_{\eta }/{\sigma }^{2}I& {\sigma }_{\eta }^2/{\sigma }^{2}I\end{array}\right).$$
This is a $2\times 2$ block matrix so I could work out the inverse to explicitly get $\mathrm{\Omega }$ but I’ll just numerically invert it in my example at the end. I don’t want to have ${\sigma }^2$s in $\mathrm{\Omega }$ so I’ll parameterize this by
$$\theta =({\sigma }_{\alpha }/\sigma ,{\sigma }_{\eta }/\sigma ,\rho )\in [0,\mathrm{\infty }{)}^2\times [-1,1].$$
The advantage to taking $\rho$ to be the correlation versus the covariance is that its bounds don’t depend on the variances in question.

The variance of $y$ is therefore
$$\mathrm{Var}(y)={\sigma }^2(Z{\mathrm{\Omega }}^{-1}{Z}^T+I)$$
which is exactly what I had in my previous post. I now can use my previous results to fit this model, but before I do I want to examine the actual structure of $\mathrm{Var}(y)$ to see what my random slopes assumption does to the covariance matrix.

$\mathrm{Var}(y)$ is block diagonal because I’m still assuming independence between groups, so I’ll look at just a single block which I’ll denote by $\mathrm{\Psi }$. I can work out
$$\mathrm{\Psi }={\sigma }_{\alpha }^2\mathbf{1}{\mathbf{1}}^T+{\sigma }_{\gamma }^2\mathbf{x}{\mathbf{x}}^T+\rho {\sigma }_{\alpha }{\sigma }_{\gamma }(\mathbf{1}{\mathbf{x}}^T+x{\mathbf{1}}^T)+{\sigma }^{2}I.$$
The ${\sigma }_{\alpha }^2\mathbf{1}{\mathbf{1}}^T+{\sigma }^{2}I$ part is easy to interpret, and is what comes from the random intercepts plus error: there is a baseline level of correlation between all observations within this group but the independent individual-level noise (i.e. $\epsilon$) prevents this from being constrained to a subspace (which is what would happen if my covariance matrix was just the singular ${\sigma }_{\alpha }^2\mathbf{1}{\mathbf{1}}^T$).

The more complicated part is ${\sigma }_{\gamma }^2\mathbf{x}{\mathbf{x}}^T+\rho {\sigma }_{\alpha }{\sigma }_{\gamma }(\mathbf{1}{\mathbf{x}}^T+x{\mathbf{1}}^T)$. I’ll interpret it by looking at a particular element of $\mathrm{\Psi }$.

$$\begin{array}{rl}{\mathrm{\Psi }}_{ij}& ={\sigma }_{\alpha }^2+{\sigma }^2{\delta }_{ij}+{\sigma }_{\gamma }^2{x}_i{x}_j+\rho {\sigma }_{\alpha }{\sigma }_{\gamma }(x_i+x_j)\\ & ={\sigma }_{\alpha }^2+{\sigma }^2{\delta }_{ij}+({\sigma }_{\gamma }{x}_i)({\sigma }_{\gamma }{x}_j)+\rho {\sigma }_{\alpha }({\sigma }_{\gamma }{x}_i+{\sigma }_{\gamma }{x}_j)\\ & ={\sigma }_{\alpha }^2(1–{\rho }^2)+{\sigma }^2{\delta }_{ij}+({\sigma }_{\gamma }{x}_i+\rho {\sigma }_{\alpha })({\sigma }_{\gamma }{x}_j+\rho {\sigma }_{\alpha })\end{array}$$
where ${\delta }_{ij}$ is the Kronecker delta, and this form suggests an interpretation via kernel functions. This is also where $\rho \ne 0$ can be seen to give me translation invariance in the $x$s since adding a constant to each $x_i$ can be canceled out via the $\rho {\sigma }_{\alpha }$ term. If $\rho =0$ there’d be no way to cancel that change out.

Define a function $k:\mathbb{R}\times \mathbb{R}\to \mathbb{R}$ by
$$k(x_i,x_j)=({\sigma }_{\gamma }{x}_i+\rho {\sigma }_{\alpha })({\sigma }_{\gamma }{x}_j+\rho {\sigma }_{\alpha })$$

Result 1: $k$ is positive semidefinite, which makes it a valid kernel function.

Pf: Let $\mathbf{x}\in {\mathbb{R}}^M$ and consider the matrix $K$ with $K_{ij}=k({\mathbf{x}}_i,{\mathbf{x}}_j)$. For simplicity I’ll say $k(x,x^{\prime })=(ax+b)(a{x}^{\prime }+b)$. I know
$$K=a^2\mathbf{x}{\mathbf{x}}^T+ab(\mathbf{x}{\mathbf{1}}^T+\mathbf{1}{\mathbf{x}}^T)+b^2\mathbf{1}{\mathbf{1}}^T$$
(I basically had this before I even defined $k$). Letting $v\in {\mathbb{R}}^M$ I have
$$\begin{gathered} v^{T}Kv=a^2(v^T\mathbf{x}{)}^2+2ab(v^T\mathbf{x})(v^T\mathbf{1})+b^2(v^T\mathbf{1}{)}^2 \\ =(a{v}^T\mathbf{x}+b{v}^T\mathbf{1}{)}^2\ge 0 \end{gathered}$$
for any $v\in {\mathbb{R}}^M$, which makes $K$ PSD. $K$ is not strictly PD in general though since if $v\perp \text{span}\{\mathbf{x},\mathbf{1}\}$, which is at most a 2 dimensional space, then I’ll have $v^{T}Kv=0$ even if $v\ne \mathbf{0}$. This means $\text{rank}(K)\le 2$ so unless $M=2$ I can’t have $K$ be PD.

$\mathrm{<img\; draggable="false"\; role="img"\; class="emoji"\; alt="◻"\; src="https://s.w.org/images/core/emoji/15.0.3/svg/25fb.svg"\; loading="lazy">}$

I could also have established that $k$ is PSD by noting that $k(x,x^{\prime })=⟨\phi (x),\phi (x^{\prime })⟩$ with the feature map $\phi (x)={\sigma }_{\gamma }x+\rho {\sigma }_{\alpha }$, so $K$ is a Gram matrix.

Looking at it this way shows that the covariance is large and positive when $x_i$ and $x_j$ are large with the same signs, and is large and negative when $x_i$ and $x_j$ are large with different signs. This reflects how when I choose to use random slopes I’m expecting that some groups would grow faster than the usual rate, or perhaps slower than the usual rate, but that either way there’d be an increasing covariance between $x$s as they get larger.

It’s worth commenting on how this is not an isotropic kernel. With an isotropic kernel like the squared exponential kernel $(\mathbf{x},{\mathbf{x}}^{\prime })\mapsto \exp (-\gamma \Vert \mathbf{x}-{\mathbf{x}}^{\prime }{\Vert }^2)$ it only depends on the distance, so the covariance matrix is diagonally dominant and decays in a “banded” fashion away from the diagonal. With my kernel here I will still get a decently prominent diagonal but not always to the same extent since it’s only PSD rather than PD, and the actual magnitudes of the input values still matter so I don’t get a banded decay away from the diagonal.

Here are 4 example covariance matrices with this kernel, just to get a sense of what they tend to look like. Each is made with a linear sequence of length $50$ for $\mathbf{x}$. Red values are lower and white values are higher.

 

 

This shows how when all of $\mathbf{x}$ is positive then the covariance is at its largest for the final values, while when $\mathbf{x}$ has both positive and negative values then the covariance is at its largest at both extremes.

I’ll now actually fit this model. I’ll use both ML and REML to compare the two. I won’t worry about the BLUPs because the parameter estimates are the real sticking point.

library(lme4)
mod_ml <- lmer(Reaction ~ Days + (1 + Days | Subject), data = sleepstudy, REML=FALSE)
mod_reml <- lmer(Reaction ~ Days + (1 + Days | Subject), data = sleepstudy, REML=TRUE)

## setting up my data ~~~~~~~~~~~~~~~~~~~~
X <- cbind(1, sleepstudy$Days)
Z <- model.matrix(~ 0 + Subject + Subject:Days, data = sleepstudy)
y <- sleepstudy$Reaction
logdet <- function(A) as.numeric(determinant(A)$mod)

## defining functions for R, beta, sigma2, and theta ~~~~~~~~~~
Rfunc <- function(theta, Z) {
   # reminder: theta = (s_alpha/sigma, s_eta/sigma, rho)
   m <- ncol(Z)/2
   Omega_inv <- rbind(
      cbind(theta[1]^2 * diag(m), prod(theta) * diag(m)),
      cbind(prod(theta) * diag(m), theta[2]^2 * diag(m))
   )
   diag(nrow(Z)) - Z %*% solve(t(Z) %*% Z + solve(Omega_inv)) %*% t(Z)
}

bfunc <- function(R, X, y) { # beta hat
   solve(t(X) %*% R %*% X) %*% t(X) %*% R %*% y
}

s2func <- function(R, X, y, REML) { # sigma^2 hat
   bhat <- bfunc(R, X, y)
   if(!REML) {
      t(y - X %*% bhat) %*% R %*% (y - X %*% bhat) / length(y)
   } else {
      t(y) %*% (R - R %*% X %*% solve(t(X) %*% R %*% X) %*% t(X) %*% R) %*% y /
         (length(y) - ncol(X))
   }
}

# the actual objective function for theta, \ell_p or \ell_{rp}
negloglik <- function(theta, X, y, Z, REML) {
   R <- Rfunc(theta, Z)
   if(!REML) {
      out <- .5 * logdet(R) - (length(y)/2) * log(s2func(R, X, y, REML))
   } else {
      out <- .5 * logdet(R) -.5 * logdet(t(X) %*% R %*% X) -
         (length(y)-ncol(X))/2 * log(s2func(R, X, y, REML))
   }
   -as.numeric(out)
}

## optimizing the log likelihoods ~~~~~~~~~~~~~~~~~~~
opt_ml <- optim(c(.1, .1, .1), negloglik, X=X, y=y, Z=Z, REML=FALSE,
      lower = c(0.001,0.001,-1), upper=c(Inf,Inf,1), method="L-BFGS-B")
opt_reml <- optim(c(.1, .1, .1), negloglik, X=X, y=y, Z=Z, REML=TRUE,
      lower = c(0.001,0.001,-1), upper=c(Inf,Inf,1), method="L-BFGS-B")

## collecting the results and comparing to lme4 ~~~~~~~~~~~~~~~~
results <- function(theta_hat, REML, X, y, Z) {
   Rhat <- Rfunc(theta_hat, Z)
   s2 <- s2func(Rhat, X, y, REML)
   out <- list(
      beta_hat = bfunc(Rhat, X, y), s2 = s2, s2_alpha = theta_hat[1]^2 * s2,
      s2_eta = theta_hat[2]^2 * s2, rho = theta_hat[3]
   )
   lapply(out, as.numeric) # the helper funcs return matrices
}

results(opt_ml$par, FALSE, X, y, Z)
summary(mod_ml)

results(opt_reml$par, TRUE, X, y, Z)
summary(mod_reml)

Everything agrees so I’m good to go.