In [0]:
#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
|
|
|
|
In this notebook we introduce Generalized Linear Models via a worked example. We solve this example in two different ways using two algorithms for efficiently fitting GLMs in TensorFlow Probability: Fisher scoring for dense data, and coordinatewise proximal gradient descent for sparse data. We compare the fitted coefficients to the true coefficients and, in the case of coordinatewise proximal gradient descent, to the output of R's similar glmnet algorithm. Finally, we provide further mathematical details and derivations of several key properties of GLMs.
A generalized linear model (GLM) is a linear model ($\eta = x^\top \beta$) wrapped in a transformation (link function) and equipped with a response distribution from an exponential family. The choice of link function and response distribution is very flexible, which lends great expressivity to GLMs. The full details, including a sequential presentation of all the definitions and results building up to GLMs in unambiguous notation, are found in "Derivation of GLM Facts" below. We summarize:
In a GLM, a predictive distribution for the response variable $Y$ is associated with a vector of observed predictors $x$. The distribution has the form:
\begin{align*} p(y \, |\, x) &= m(y, \phi) \exp\left(\frac{\theta\, T(y) - A(\theta)}{\phi}\right) \\ \theta &:= h(\eta) \\ \eta &:= x^\top \beta \end{align*}Here $\beta$ are the parameters ("weights"), $\phi$ a hyperparameter representing dispersion ("variance"), and $m$, $h$, $T$, $A$ are characterized by the user-specified model family.
The mean of $Y$ depends on $x$ by composition of linear response $\eta$ and (inverse) link function, i.e.:
$$ \mu := g^{-1}(\eta) $$where $g$ is the so-called link function. In TFP the choice of link function and model family are jointly specifed by a tfp.glm.ExponentialFamily subclass. Examples include:
tfp.glm.Normal, aka "linear regression"tfp.glm.Bernoulli, aka "logistic regression"tfp.glm.Poisson, aka "Poisson regression"tfp.glm.BernoulliNormalCDF, aka "probit regression".TFP prefers to name model families according to the distribution over Y rather than the link function since tfp.Distributions are already first-class citizens. If the tfp.glm.ExponentialFamily subclass name contains a second word, this indicates a non-canonical link function.
GLMs have several remarkable properties which permit efficient implementation of the maximum likelihood estimator. Chief among these properties are simple formulas for the gradient of the log-likelihood $\ell$, and for the Fisher information matrix, which is the expected value of the Hessian of the negative log-likelihood under a re-sampling of the response under the same predictors. I.e.:
\begin{align*} \nabla_\beta\, \ell(\beta\, ;\, \mathbf{x}, \mathbf{y}) &= \mathbf{x}^\top \,\text{diag}\left(\frac{ {\textbf{Mean}_T}'(\mathbf{x} \beta) }{ {\textbf{Var}_T}(\mathbf{x} \beta) }\right) \left(\mathbf{T}(\mathbf{y}) - {\textbf{Mean}_T}(\mathbf{x} \beta)\right) \\ \mathbb{E}_{Y_i \sim \text{GLM} | x_i} \left[ \nabla_\beta^2\, \ell(\beta\, ;\, \mathbf{x}, \mathbf{Y}) \right] &= -\mathbf{x}^\top \,\text{diag}\left( \frac{ \phi\, {\textbf{Mean}_T}'(\mathbf{x} \beta)^2 }{ {\textbf{Var}_T}(\mathbf{x} \beta) }\right)\, \mathbf{x} \end{align*}where $\mathbf{x}$ is the matrix whose $i$th row is the predictor vector for the $i$th data sample, and $\mathbf{y}$ is vector whose $i$th coordinate is the observed response for the $i$th data sample. Here (loosely speaking), ${\text{Mean}_T}(\eta) := \mathbb{E}[T(Y)\,|\,\eta]$ and ${\text{Var}_T}(\eta) := \text{Var}[T(Y)\,|\,\eta]$, and boldface denotes vectorization of these functions. Full details of what distributions these expectations and variances are over can be found in "Derivation of GLM Facts" below.
In [0]:
import numpy as np
import pandas as pd
import scipy
import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp
tfd = tfp.distributions
In [0]:
def make_dataset(n, d, link, scale=1., dtype=np.float32):
model_coefficients = tfd.Uniform(
low=-1., high=np.array(1, dtype)).sample(d, seed=42)
radius = np.sqrt(2.)
model_coefficients *= radius / tf.linalg.norm(model_coefficients)
mask = tf.random.shuffle(tf.range(d)) < int(0.5 * d)
model_coefficients = tf.where(
mask, model_coefficients, np.array(0., dtype))
model_matrix = tfd.Normal(
loc=0., scale=np.array(1, dtype)).sample([n, d], seed=43)
scale = tf.convert_to_tensor(scale, dtype)
linear_response = tf.linalg.matvec(model_matrix, model_coefficients)
if link == 'linear':
response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
elif link == 'probit':
response = tf.cast(
tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
dtype)
elif link == 'logit':
response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
else:
raise ValueError('unrecognized true link: {}'.format(link))
return model_matrix, response, model_coefficients, mask
In [0]:
x, y, model_coefficients_true, _ = [t.numpy() for t in make_dataset(
n=int(1e5), d=100, link='probit')]
DATA_DIR = '/tmp/glm_example'
tf.io.gfile.makedirs(DATA_DIR)
with tf.io.gfile.GFile('{}/x.csv'.format(DATA_DIR), 'w') as f:
np.savetxt(f, x, delimiter=',')
with tf.io.gfile.GFile('{}/y.csv'.format(DATA_DIR), 'w') as f:
np.savetxt(f, y.astype(np.int32) + 1, delimiter=',', fmt='%d')
with tf.io.gfile.GFile(
'{}/model_coefficients_true.csv'.format(DATA_DIR), 'w') as f:
np.savetxt(f, model_coefficients_true, delimiter=',')
The function tfp.glm.fit implements Fisher scoring, which takes as some of its arguments:
model_matrix = $\mathbf{x}$response = $\mathbf{y}$model = callable which, given argument $\boldsymbol{\eta}$, returns the triple $\left(
{\textbf{Mean}_T}(\boldsymbol{\eta}),
{\textbf{Var}_T}(\boldsymbol{\eta}),
{\textbf{Mean}_T}'(\boldsymbol{\eta})
\right)$.We recommend that model be an instance of the tfp.glm.ExponentialFamily class. There are several pre-made implementations available, so for most common GLMs no custom code is necessary.
In [0]:
@tf.function(autograph=False)
def fit_model():
model_coefficients, linear_response, is_converged, num_iter = tfp.glm.fit(
model_matrix=x, response=y, model=tfp.glm.BernoulliNormalCDF())
log_likelihood = tfp.glm.BernoulliNormalCDF().log_prob(y, linear_response)
return (model_coefficients, linear_response, is_converged, num_iter,
log_likelihood)
[model_coefficients, linear_response, is_converged, num_iter,
log_likelihood] = [t.numpy() for t in fit_model()]
print(('is_converged: {}\n'
' num_iter: {}\n'
' accuracy: {}\n'
' deviance: {}\n'
'||w0-w1||_2 / (1+||w0||_2): {}'
).format(
is_converged,
num_iter,
np.mean((linear_response > 0.) == y),
2. * np.mean(log_likelihood),
np.linalg.norm(model_coefficients_true - model_coefficients, ord=2) /
(1. + np.linalg.norm(model_coefficients_true, ord=2))
))
Fisher scoring is a modification of Newton's method to find the maximum-likelihood estimate
$$ \hat\beta := \underset{\beta}{\text{arg max}}\ \ \ell(\beta\ ;\ \mathbf{x}, \mathbf{y}). $$Vanilla Newton's method, searching for zeros of the gradient of the log-likelihood, would follow the update rule
$$ \beta^{(t+1)}_{\text{Newton}} := \beta^{(t)} - \alpha \left( \nabla^2_\beta\, \ell(\beta\ ;\ \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}}^{-1} \left( \nabla_\beta\, \ell(\beta\ ;\ \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} $$where $\alpha \in (0, 1]$ is a learning rate used to control the step size.
In Fisher scoring, we replace the Hessian with the negative Fisher information matrix:
\begin{align*} \beta^{(t+1)} &:= \beta^{(t)} - \alpha\, \mathbb{E}_{ Y_i \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(x_i^\top \beta^{(t)}), \phi) } \left[ \left( \nabla^2_\beta\, \ell(\beta\ ;\ \mathbf{x}, \mathbf{Y}) \right)_{\beta = \beta^{(t)}} \right]^{-1} \left( \nabla_\beta\, \ell(\beta\ ;\ \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} \\[3mm] \end{align*}[Note that here $\mathbf{Y} = (Y_i)_{i=1}^{n}$ is random, whereas $\mathbf{y}$ is still the vector of observed responses.]
By the formulas in "Fitting GLM Parameters To Data" below, this simplifies to
\begin{align*} \beta^{(t+1)} &= \beta^{(t)} + \alpha \left( \mathbf{x}^\top \text{diag}\left( \frac{ \phi\, {\textbf{Mean}_T}'(\mathbf{x} \beta^{(t)})^2 }{ {\textbf{Var}_T}(\mathbf{x} \beta^{(t)}) }\right)\, \mathbf{x} \right)^{-1} \left( \mathbf{x}^\top \text{diag}\left(\frac{ {\textbf{Mean}_T}'(\mathbf{x} \beta^{(t)}) }{ {\textbf{Var}_T}(\mathbf{x} \beta^{(t)}) }\right) \left(\mathbf{T}(\mathbf{y}) - {\textbf{Mean}_T}(\mathbf{x} \beta^{(t)})\right) \right). \end{align*}tfp.glm.fit_sparse implements a GLM fitter more suited to sparse data sets, based on the algorithm in Yuan, Ho and Lin 2012. Its features include:
We first present an example usage of the code. Details of the algorithm are further elaborated in "Algorithm Details for tfp.glm.fit_sparse" below.
In [0]:
model = tfp.glm.Bernoulli()
model_coefficients_start = tf.zeros(x.shape[-1], np.float32)
@tf.function(autograph=False)
def fit_model():
return tfp.glm.fit_sparse(
model_matrix=tf.convert_to_tensor(x),
response=tf.convert_to_tensor(y),
model=model,
model_coefficients_start=model_coefficients_start,
l1_regularizer=800.,
l2_regularizer=None,
maximum_iterations=10,
maximum_full_sweeps_per_iteration=10,
tolerance=1e-6,
learning_rate=None)
model_coefficients, is_converged, num_iter = [t.numpy() for t in fit_model()]
coefs_comparison = pd.DataFrame({
'Learned': model_coefficients,
'True': model_coefficients_true,
})
print(('is_converged: {}\n'
' num_iter: {}\n\n'
'Coefficients:').format(
is_converged,
num_iter))
coefs_comparison
Out[0]:
Note that the learned coefficients have the same sparsity pattern as the true coefficients.
In [0]:
# Save the learned coefficients to a file.
with tf.io.gfile.GFile('{}/model_coefficients_prox.csv'.format(DATA_DIR), 'w') as f:
np.savetxt(f, model_coefficients, delimiter=',')
In [0]:
suppressMessages({
library('glmnet')
})
In [0]:
data_dir <- '/tmp/glm_example'
x <- as.matrix(read.csv(paste(data_dir, '/x.csv', sep=''),
header=FALSE))
y <- as.matrix(read.csv(paste(data_dir, '/y.csv', sep=''),
header=FALSE, colClasses='integer'))
In [0]:
fit <- glmnet(
x = x,
y = y,
family = "binomial", # Logistic regression
alpha = 1, # corresponds to l1_weight = 1, l2_weight = 0
standardize = FALSE,
intercept = FALSE,
thresh = 1e-30,
type.logistic = "Newton"
)
In [0]:
write.csv(as.matrix(coef(fit, 0.008)),
paste(data_dir, '/model_coefficients_glmnet.csv', sep=''),
row.names=FALSE)
In [0]:
DATA_DIR = '/tmp/glm_example'
with tf.io.gfile.GFile('{}/model_coefficients_glmnet.csv'.format(DATA_DIR),
'r') as f:
model_coefficients_glmnet = np.loadtxt(f,
skiprows=2 # Skip column name and intercept
)
with tf.io.gfile.GFile('{}/model_coefficients_prox.csv'.format(DATA_DIR),
'r') as f:
model_coefficients_prox = np.loadtxt(f)
with tf.io.gfile.GFile(
'{}/model_coefficients_true.csv'.format(DATA_DIR), 'r') as f:
model_coefficients_true = np.loadtxt(f)
In [0]:
coefs_comparison = pd.DataFrame({
'TFP': model_coefficients_prox,
'R': model_coefficients_glmnet,
'True': model_coefficients_true,
})
coefs_comparison
Out[0]:
tfp.glm.fit_sparseWe present the algorithm as a sequence of three modifications to Newton's method. In each one, the update rule for $\beta$ is based on a vector $s$ and a matrix $H$ which approximate the gradient and Hessian of the log-likelihood. In step $t$, we choose a coordinate $j^{(t)}$ to change, and we update $\beta$ according to the update rule:
\begin{align*} u^{(t)} &:= \frac{ \left( s^{(t)} \right)_{j^{(t)}} }{ \left( H^{(t)} \right)_{j^{(t)},\, j^{(t)}} } \\[3mm] \beta^{(t+1)} &:= \beta^{(t)} - \alpha\, u^{(t)} \,\text{onehot}(j^{(t)}) \end{align*}This update is a Newton-like step with learning rate $\alpha$. Except for the final piece (L1 regularization), the modifications below differ only in how they update $s$ and $H$.
In coordinatewise Newton's method, we set $s$ and $H$ to the true gradient and Hessian of the log-likelihood:
\begin{align*} s^{(t)}_{\text{vanilla}} &:= \left( \nabla_\beta\, \ell(\beta \,;\, \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} \\ H^{(t)}_{\text{vanilla}} &:= \left( \nabla^2_\beta\, \ell(\beta \,;\, \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} \end{align*}The gradient and Hessian of the log-likelihood are often expensive to compute, so it is often worthwhile to approximate them. We can do so as follows:
To further reduce the cost of the vanilla update steps, we can set $H$ to the negative Fisher information matrix (efficiently computable using the formulas in "Fitting GLM Parameters to Data" below) rather than the exact Hessian:
\begin{align*} H_{\text{Fisher}}^{(t+1)} &:= \mathbb{E}_{Y_i \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(x_i^\top \beta^{(t+1)}), \phi)} \left[ \left( \nabla_\beta^2\, \ell(\beta\, ;\, \mathbf{x}, \mathbf{Y}) \right)_{\beta = \beta^{(t+1)}} \right] \\ &= -\mathbf{x}^\top \,\text{diag}\left( \frac{ \phi\, {\textbf{Mean}_T}'(\mathbf{x} \beta^{(t+1)})^2 }{ {\textbf{Var}_T}(\mathbf{x} \beta^{(t+1)}) }\right)\, \mathbf{x} \\ s_{\text{Fisher}}^{(t+1)} &:= s_{\text{vanilla}}^{(t+1)} \\ &= \left( \mathbf{x}^\top \,\text{diag}\left(\frac{ {\textbf{Mean}_T}'(\mathbf{x} \beta^{(t+1)}) }{ {\textbf{Var}_T}(\mathbf{x} \beta^{(t+1)}) }\right) \left(\mathbf{T}(\mathbf{y}) - {\textbf{Mean}_T}(\mathbf{x} \beta^{(t+1)})\right) \right) \end{align*}To incorporate L1 regularization, we replace the update rule
$$ \beta^{(t+1)} := \beta^{(t)} - \alpha\, u^{(t)} \,\text{onehot}(j^{(t)}) $$with the more general update rule
\begin{align*} \gamma^{(t)} &:= -\frac{\alpha\, r_{\text{L1}}}{\left(H^{(t)}\right)_{j^{(t)},\, j^{(t)}}} \\[2mm] \left(\beta_{\text{reg}}^{(t+1)}\right)_j &:= \begin{cases} \beta^{(t+1)}_j &\text{if } j \neq j^{(t)} \\ \text{SoftThreshold} \left( \beta^{(t)}_j - \alpha\, u^{(t)} ,\ \gamma^{(t)} \right) &\text{if } j = j^{(t)} \end{cases} \end{align*}where $r_{\text{L1}} > 0$ is a supplied constant (the L1 regularization coefficient) and $\text{SoftThreshold}$ is the soft thresholding operator, defined by
$$ \text{SoftThreshold}(\beta, \gamma) := \begin{cases} \beta + \gamma &\text{if } \beta < -\gamma \\ 0 &\text{if } -\gamma \leq \beta \leq \gamma \\ \beta - \gamma &\text{if } \beta > \gamma. \end{cases} $$This update rule has the following two inspirational properties, which we explain below:
In the limiting case $r_{\text{L1}} \to 0$ (i.e., no L1 regularization), this update rule is identical to the original update rule.
This update rule can be interpreted as applying a proximity operator whose fixed point is the solution to the L1-regularized minimization problem
To see (1), note that if $r_{\text{L1}} = 0$ then $\gamma^{(t)} = 0$, hence
\begin{align*} \left(\beta_{\text{reg}}^{(t+1)}\right)_{j^{(t)}} &= \text{SoftThreshold} \left( \beta^{(t)}_{j^{(t)}} - \alpha\, u^{(t)} ,\ 0 \right) \\ &= \beta^{(t)}_{j^{(t)}} - \alpha\, u^{(t)}. \end{align*}Hence
\begin{align*} \beta_{\text{reg}}^{(t+1)} &= \beta^{(t)} - \alpha\, u^{(t)} \,\text{onehot}(j^{(t)}) \\ &= \beta^{(t+1)}. \end{align*}To see (2), first note (see Wikipedia) that for any $\gamma > 0$, the update rule
$$ \left(\beta_{\text{exact-prox}, \gamma}^{(t+1)}\right)_{j^{(t)}} := \text{prox}_{\gamma \lVert \cdot \rVert_1} \left( \beta^{(t)}_{j^{(t)}} + \frac{\gamma}{r_{\text{L1}}} \left( \left( \nabla_\beta\, \ell(\beta \,;\, \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} \right)_{j^{(t)}} \right) $$satisfies (2), where $\text{prox}$ is the proximity operator (see Yu, where this operator is denoted $\mathsf{P}$). The right-hand side of the above equation is computed here:
$$ \left(\beta_{\text{exact-prox}, \gamma}^{(t+1)}\right)_{j^{(t)}} = \text{SoftThreshold} \left( \beta^{(t)}_{j^{(t)}} + \frac{\gamma}{r_{\text{L1}}} \left( \left( \nabla_\beta\, \ell(\beta \,;\, \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} \right)_{j^{(t)}} ,\ \gamma \right). $$In particular, setting $\gamma = \gamma^{(t)} = -\frac{\alpha\, r_{\text{L1}}}{\left(H^{(t)}\right)_{j^{(t)}, j^{(t)}}}$ (note that $\gamma^{(t)} > 0$ as long as the negative log-likelihood is convex), we obtain the update rule
$$ \left(\beta_{\text{exact-prox}, \gamma^{(t)}}^{(t+1)}\right)_{j^{(t)}} = \text{SoftThreshold} \left( \beta^{(t)}_{j^{(t)}} - \alpha \frac{ \left( \left( \nabla_\beta\, \ell(\beta \,;\, \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}} \right)_{j^{(t)}} }{ \left(H^{(t)}\right)_{j^{(t)}, j^{(t)}} } ,\ \gamma^{(t)} \right). $$We then replace the exact gradient $\left( \nabla_\beta\, \ell(\beta \,;\, \mathbf{x}, \mathbf{y}) \right)_{\beta = \beta^{(t)}}$ with its approximation $s^{(t)}$, obtaining
\begin{align*} \left(\beta_{\text{exact-prox}, \gamma^{(t)}}^{(t+1)}\right)_{j^{(t)}} &\approx \text{SoftThreshold} \left( \beta^{(t)}_{j^{(t)}} - \alpha \frac{ \left(s^{(t)}\right)_{j^{(t)}} }{ \left(H^{(t)}\right)_{j^{(t)}, j^{(t)}} } ,\ \gamma^{(t)} \right) \\ &= \text{SoftThreshold} \left( \beta^{(t)}_{j^{(t)}} - \alpha\, u^{(t)} ,\ \gamma^{(t)} \right). \end{align*}Hence
$$ \beta_{\text{exact-prox}, \gamma^{(t)}}^{(t+1)} \approx \beta_{\text{reg}}^{(t+1)}. $$Consider a family of probability distributions parameterized by parameter vector $\theta$, having probability densities $\left\{p(\cdot | \theta)\right\}_{\theta \in \mathcal{T}}$. The score of an outcome $y$ at parameter vector $\theta_0$ is defined to be the gradient of the log likelihood of $y$ (evaluated at $\theta_0$), that is,
$$ \text{score}(y, \theta_0) := \left[\nabla_\theta\, \log p(y | \theta)\right]_{\theta=\theta_0}. $$We have
\begin{align*} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[\text{score}(Y, \theta_0)\right] &:=\mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[\left(\nabla_\theta \log p(Y|\theta)\right)_{\theta=\theta_0}\right] \\ &\stackrel{\text{(1)}}{=} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[\frac{\left(\nabla_\theta p(Y|\theta)\right)_{\theta=\theta_0}}{p(Y|\theta=\theta_0)}\right] \\ &\stackrel{\text{(2)}}{=} \int_{\mathcal{Y}} \left[\frac{\left(\nabla_\theta p(y|\theta)\right)_{\theta=\theta_0}}{p(y|\theta=\theta_0)}\right] p(y | \theta=\theta_0)\, dy \\ &= \int_{\mathcal{Y}} \left(\nabla_\theta p(y|\theta)\right)_{\theta=\theta_0}\, dy \\ &\stackrel{\text{(3)}}{=} \left[\nabla_\theta \left(\int_{\mathcal{Y}} p(y|\theta)\, dy\right) \right]_{\theta=\theta_0} \\ &\stackrel{\text{(4)}}{=} \left[\nabla_\theta\, 1 \right]_{\theta=\theta_0} \\ &= 0, \end{align*}where we have used: (1) chain rule for differentiation, (2) definition of expectation, (3) passing differentiation under the integral sign (using the regularity conditions), (4) the integral of a probability density is 1.
Under mild regularity conditions (permitting us to pass differentiation under the integral),
$$ \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \text{score}(Y, \theta_0) \text{score}(Y, \theta_0)^\top \right] = -\mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \left(\nabla_\theta^2 \log p(Y | \theta)\right)_{\theta=\theta_0} \right] $$where $\nabla_\theta^2 F$ denotes the Hessian matrix, whose $(i, j)$ entry is $\frac{\partial^2 F}{\partial \theta_i \partial \theta_j}$.
The left-hand side of this equation is called the Fisher information of the family $\left\{p(\cdot | \theta)\right\}_{\theta \in \mathcal{T}}$ at parameter vector $\theta_0$.
We have
\begin{align*} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \left(\nabla_\theta^2 \log p(Y | \theta)\right)_{\theta=\theta_0} \right] &\stackrel{\text{(1)}}{=} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \left(\nabla_\theta^\top \frac{ \nabla_\theta p(Y | \theta) }{ p(Y|\theta) }\right)_{\theta=\theta_0} \right] \\ &\stackrel{\text{(2)}}{=} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \frac{ \left(\nabla^2_\theta p(Y | \theta)\right)_{\theta=\theta_0} }{ p(Y|\theta=\theta_0) } - \left(\frac{ \left(\nabla_\theta\, p(Y|\theta)\right)_{\theta=\theta_0} }{ p(Y|\theta=\theta_0) }\right) \left(\frac{ \left(\nabla_\theta\, p(Y|\theta)\right)_{\theta=\theta_0} }{ p(Y|\theta=\theta_0) }\right)^\top \right] \\ &\stackrel{\text{(3)}}{=} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \frac{ \left(\nabla^2_\theta p(Y | \theta)\right)_{\theta=\theta_0} }{ p(Y|\theta=\theta_0) } - \text{score}(Y, \theta_0) \,\text{score}(Y, \theta_0)^\top \right], \end{align*}where we have used (1) chain rule for differentiation, (2) quotient rule for differentiation, (3) chain rule again, in reverse.
To complete the proof, it suffices to show that
$$ \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \frac{ \left(\nabla^2_\theta p(Y | \theta)\right)_{\theta=\theta_0} }{ p(Y|\theta=\theta_0) } \right] \stackrel{\text{?}}{=} 0. $$To do that, we pass differentiation under the integral sign twice:
\begin{align*} \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)}\left[ \frac{ \left(\nabla^2_\theta p(Y | \theta)\right)_{\theta=\theta_0} }{ p(Y|\theta=\theta_0) } \right] &= \int_{\mathcal{Y}} \left[ \frac{ \left(\nabla^2_\theta p(y | \theta)\right)_{\theta=\theta_0} }{ p(y|\theta=\theta_0) } \right] \, p(y | \theta=\theta_0)\, dy \\ &= \int_{\mathcal{Y}} \left(\nabla^2_\theta p(y | \theta)\right)_{\theta=\theta_0} \, dy \\ &= \left[ \nabla_\theta^2 \left( \int_{\mathcal{Y}} p(y | \theta) \, dy \right) \right]_{\theta=\theta_0} \\ &= \left[ \nabla_\theta^2 \, 1 \right]_{\theta=\theta_0} \\ &= 0. \end{align*}If $a$, $b$ and $c$ are scalar-valued functions, $c$ twice differentiable, such that the family of distributions $\left\{p(\cdot | \theta)\right\}_{\theta \in \mathcal{T}}$ defined by
$$ p(y|\theta) = a(y) \exp\left(b(y)\, \theta - c(\theta)\right) $$satisfies the mild regularity conditions that permit passing differentiation with respect to $\theta$ under an integral with respect to $y$, then
$$ \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ b(Y) \right] = c'(\theta_0) $$and
$$ \text{Var}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ b(Y) \right] = c''(\theta_0). $$(Here $'$ denotes differentiation, so $c'$ and $c''$ are the first and second derivatives of $c$. )
For this family of distributions, we have $\text{score}(y, \theta_0) = b(y) - c'(\theta_0)$. The first equation then follows from the fact that $\mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ \text{score}(y, \theta_0) \right] = 0$. Next, we have
\begin{align*} \text{Var}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ b(Y) \right] &= \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ \left(b(Y) - c'(\theta_0)\right)^2 \right] \\ &= \text{the one entry of } \mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ \text{score}(y, \theta_0) \text{score}(y, \theta_0)^\top \right] \\ &= \text{the one entry of } -\mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ \left(\nabla_\theta^2 \log p(\cdot | \theta)\right)_{\theta=\theta_0} \right] \\ &= -\mathbb{E}_{Y \sim p(\cdot | \theta=\theta_0)} \left[ -c''(\theta_0) \right] \\ &= c''(\theta_0). \end{align*}A (scalar) overdispersed exponential family is a family of distributions whose densities take the form
$$ p_{\text{OEF}(m, T)}(y\, |\, \theta, \phi) = m(y, \phi) \exp\left(\frac{\theta\, T(y) - A(\theta)}{\phi}\right), $$where $m$ and $T$ are known scalar-valued functions, and $\theta$ and $\phi$ are scalar parameters.
[Note that $A$ is overdetermined: for any $\phi_0$, the function $A$ is completely determined by the constraint that $\int p_{\text{OEF}(m, T)}(y\ |\ \theta, \phi=\phi_0)\, dy = 1$ for all $\theta$. The $A$'s produced by different values of $\phi_0$ must all be the same, which places a constraint on the functions $m$ and $T$.]
Under the same conditions as "Lemma about the derivative of the log partition function," we have
$$ \mathbb{E}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta, \phi)} \left[ T(Y) \right] = A'(\theta) $$and
$$ \text{Var}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta, \phi)} \left[ T(Y) \right] = \phi A''(\theta). $$By "Lemma about the derivative of the log partition function," we have
$$ \mathbb{E}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta, \phi)} \left[ \frac{T(Y)}{\phi} \right] = \frac{A'(\theta)}{\phi} $$and
$$ \text{Var}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta, \phi)} \left[ \frac{T(Y)}{\phi} \right] = \frac{A''(\theta)}{\phi}. $$The result then follows from the fact that expectation is linear ($\mathbb{E}[aX] = a\mathbb{E}[X]$) and variance is degree-2 homogeneous ($\text{Var}[aX] = a^2 \,\text{Var}[X]$).
In a generalized linear model, a predictive distribution for the response variable $Y$ is associated with a vector of observed predictors $x$. The distribution is a member of an overdispersed exponential family, and the parameter $\theta$ is replaced by $h(\eta)$ where $h$ is a known function, $\eta := x^\top \beta$ is the so-called linear response, and $\beta$ is a vector of parameters (regression coefficients) to be learned. In general the dispersion parameter $\phi$ could be learned too, but in our setup we will treat $\phi$ as known. So our setup is
$$ Y \sim p_{\text{OEF}(m, T)}(\cdot\, |\, \theta = h(\eta), \phi) $$where the model structure is characterized by the distribution $p_{\text{OEF}(m, T)}$ and the function $h$ which converts linear response to parameters.
Traditionally, the mapping from linear response $\eta$ to mean $\mu := \mathbb{E}_{Y \sim p_{\text{OEF}(m, T)}(\cdot\, |\, \theta = h(\eta), \phi)}\left[ Y\right]$ is denoted
$$ \mu = g^{-1}(\eta). $$This mapping is required to be one-to-one, and its inverse, $g$, is called the link function for this GLM. Typically, one describes a GLM by naming its link function and its family of distributions -- e.g., a "GLM with Bernoulli distribution and logit link function" (also known as a logistic regression model). In order to fully characterize the GLM, the function $h$ must also be specified. If $h$ is the identity, then $g$ is said to be the canonical link function.
Define
$$ {\text{Mean}_T}(\eta) := \mathbb{E}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(\eta), \phi)} \left[ T(Y) \right] $$and
$$ {\text{Var}_T}(\eta) := \text{Var}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(\eta), \phi)} \left[ T(Y) \right]. $$Then we have
$$ h'(\eta) = \frac{\phi\, {\text{Mean}_T}'(\eta)}{{\text{Var}_T}(\eta)}. $$By "Mean and variance of the sufficient statistic," we have
$$ {\text{Mean}_T}(\eta) = A'(h(\eta)). $$Differentiating with the chain rule, we obtain $$ {\text{Mean}_T}'(\eta) = A''(h(\eta))\, h'(\eta), $$
and by "Mean and variance of the sufficient statistic,"
$$ \cdots = \frac{1}{\phi} {\text{Var}_T}(\eta)\ h'(\eta). $$The conclusion follows.
The properties derived above lend themselves very well to fitting GLM parameters $\beta$ to a data set. Quasi-Newton methods such as Fisher scoring rely on the gradient of the log likelihood and the Fisher information, which we now show can be computed especially efficiently for a GLM.
Suppose we have observed predictor vectors $x_i$ and associated scalar responses $y_i$. In matrix form, we'll say we have observed predictors $\mathbf{x}$ and response $\mathbf{y}$, where $\mathbf{x}$ is the matrix whose $i$th row is $x_i^\top$ and $\mathbf{y}$ is the vector whose $i$th element is $y_i$. The log likelihood of parameters $\beta$ is then
$$ \ell(\beta\, ;\, \mathbf{x}, \mathbf{y}) = \sum_{i=1}^{N} \log p_{\text{OEF}(m, T)}(y_i\, |\, \theta = h(x_i^\top \beta), \phi). $$To simplify the notation, let's first consider the case of a single data point, $N=1$; then we will extend to the general case by additivity.
We have
\begin{align*} \ell(\beta\, ;\, x, y) &= \log p_{\text{OEF}(m, T)}(y\, |\, \theta = h(x^\top \beta), \phi) \\ &= \log m(y, \phi) + \frac{\theta\, T(y) - A(\theta)}{\phi}, \quad\text{where}\ \theta = h(x^\top \beta). \end{align*}Hence by the chain rule,
$$ \nabla_\beta \ell(\beta\, ; \, x, y) = \frac{T(y) - A'(\theta)}{\phi}\, h'(x^\top \beta)\, x. $$Separately, by "Mean and variance of the sufficient statistic," we have $A'(\theta) = {\text{Mean}_T}(x^\top \beta)$. Hence, by "Claim: Expressing $h'$ in terms of the sufficient statistic," we have
$$ \cdots = \left(T(y) - {\text{Mean}_T}(x^\top \beta)\right) \frac{{\text{Mean}_T}'(x^\top \beta)}{{\text{Var}_T}(x^\top \beta)} \,x. $$Differentiating a second time, by the product rule we obtain
\begin{align*} \nabla_\beta^2 \ell(\beta\, ;\, x, y) &= \left[ -A''(h(x^\top \beta))\, h'(x^\top \beta) \right] h'(x^\top \beta)\, x x^\top + \left[ T(y) - A'(h(x^\top \beta)) \right] h''(x^\top \beta)\, xx^\top ] \\ &= \left( -{\text{Mean}_T}'(x^\top \beta)\, h'(x^\top \beta) + \left[T(y) - A'(h(x^\top \beta))\right] \right)\, x x^\top. \end{align*}By "Mean and variance of the sufficient statistic," we have
$$ \mathbb{E}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(x^\top \beta), \phi)} \left[ T(y) - A'(h(x^\top \beta)) \right] = 0. $$Hence
\begin{align*} \mathbb{E}_{Y \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(x^\top \beta), \phi)} \left[ \nabla_\beta^2 \ell(\beta\, ;\, x, y) \right] &= -{\text{Mean}_T}'(x^\top \beta)\, h'(x^\top \beta) x x^\top \\ &= -\frac{\phi\, {\text{Mean}_T}'(x^\top \beta)^2}{{\text{Var}_T}(x^\top \beta)}\, x x^\top. \end{align*}We now extend the $N=1$ case to the general case. Let $\boldsymbol{\eta} := \mathbf{x} \beta$ denote the vector whose $i$th coordinate is the linear response from the $i$th data sample. Let $\mathbf{T}$ (resp. ${\textbf{Mean}_T}$, resp. ${\textbf{Var}_T}$) denote the broadcasted (vectorized) function which applies the scalar-valued function $T$ (resp. ${\text{Mean}_T}$, resp. ${\text{Var}_T}$) to each coordinate. Then we have
\begin{align*} \nabla_\beta \ell(\beta\, ;\, \mathbf{x}, \mathbf{y}) &= \sum_{i=1}^{N} \nabla_\beta \ell(\beta\, ;\, x_i, y_i) \\ &= \sum_{i=1}^{N} \left(T(y) - {\text{Mean}_T}(x_i^\top \beta)\right) \frac{{\text{Mean}_T}'(x_i^\top \beta)}{{\text{Var}_T}(x_i^\top \beta)} \, x_i \\ &= \mathbf{x}^\top \,\text{diag}\left(\frac{ {\textbf{Mean}_T}'(\mathbf{x} \beta) }{ {\textbf{Var}_T}(\mathbf{x} \beta) }\right) \left(\mathbf{T}(\mathbf{y}) - {\textbf{Mean}_T}(\mathbf{x} \beta)\right) \\ \end{align*}and
\begin{align*} \mathbb{E}_{Y_i \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(x_i^\top \beta), \phi)} \left[ \nabla_\beta^2 \ell(\beta\, ;\, \mathbf{x}, \mathbf{Y}) \right] &= \sum_{i=1}^{N} \mathbb{E}_{Y_i \sim p_{\text{OEF}(m, T)}(\cdot | \theta = h(x_i^\top \beta), \phi)} \left[ \nabla_\beta^2 \ell(\beta\, ;\, x_i, Y_i) \right] \\ &= \sum_{i=1}^{N} -\frac{\phi\, {\text{Mean}_T}'(x_i^\top \beta)^2}{{\text{Var}_T}(x_i^\top \beta)}\, x_i x_i^\top \\ &= -\mathbf{x}^\top \,\text{diag}\left( \frac{ \phi\, {\textbf{Mean}_T}'(\mathbf{x} \beta)^2 }{ {\textbf{Var}_T}(\mathbf{x} \beta) }\right)\, \mathbf{x}, \end{align*}where the fractions denote element-wise division.
We now verify the above formula for gradient of the log likelihood numerically using tf.gradients, and verify the formula for Fisher information with a Monte Carlo estimate using tf.hessians:
In [0]:
def VerifyGradientAndFIM():
model = tfp.glm.BernoulliNormalCDF()
model_matrix = np.array([[1., 5, -2],
[8, -1, 8]])
def _naive_grad_and_hessian_loss_fn(x, response):
# Computes gradient and Hessian of negative log likelihood using autodiff.
predicted_linear_response = tf.linalg.matvec(model_matrix, x)
log_probs = model.log_prob(response, predicted_linear_response)
grad_loss = tf.gradients(-log_probs, [x])[0]
hessian_loss = tf.hessians(-log_probs, [x])[0]
return [grad_loss, hessian_loss]
def _grad_neg_log_likelihood_and_fim_fn(x, response):
# Computes gradient of negative log likelihood and Fisher information matrix
# using the formulas above.
predicted_linear_response = tf.linalg.matvec(model_matrix, x)
mean, variance, grad_mean = model(predicted_linear_response)
v = (response - mean) * grad_mean / variance
grad_log_likelihood = tf.linalg.matvec(model_matrix, v, adjoint_a=True)
w = grad_mean**2 / variance
fisher_info = tf.linalg.matmul(
model_matrix,
w[..., tf.newaxis] * model_matrix,
adjoint_a=True)
return [-grad_log_likelihood, fisher_info]
@tf.function(autograph=False)
def compute_grad_hessian_estimates():
# Monte Carlo estimate of E[Hessian(-LogLikelihood)], where the expectation is
# as written in "Claim (Fisher information)" above.
num_trials = 20
trial_outputs = []
np.random.seed(10)
model_coefficients_ = np.random.random(size=(model_matrix.shape[1],))
model_coefficients = tf.convert_to_tensor(model_coefficients_)
for _ in range(num_trials):
# Sample from the distribution of `model`
response = np.random.binomial(
1,
scipy.stats.norm().cdf(np.matmul(model_matrix, model_coefficients_))
).astype(np.float64)
trial_outputs.append(
list(_naive_grad_and_hessian_loss_fn(model_coefficients, response)) +
list(
_grad_neg_log_likelihood_and_fim_fn(model_coefficients, response))
)
naive_grads = tf.stack(
list(naive_grad for [naive_grad, _, _, _] in trial_outputs), axis=0)
fancy_grads = tf.stack(
list(fancy_grad for [_, _, fancy_grad, _] in trial_outputs), axis=0)
average_hess = tf.reduce_mean(tf.stack(
list(hess for [_, hess, _, _] in trial_outputs), axis=0), axis=0)
[_, _, _, fisher_info] = trial_outputs[0]
return naive_grads, fancy_grads, average_hess, fisher_info
naive_grads, fancy_grads, average_hess, fisher_info = [
t.numpy() for t in compute_grad_hessian_estimates()]
print("Coordinatewise relative error between naively computed gradients and"
" formula-based gradients (should be zero):\n{}\n".format(
(naive_grads - fancy_grads) / naive_grads))
print("Coordinatewise relative error between average of naively computed"
" Hessian and formula-based FIM (should approach zero as num_trials"
" -> infinity):\n{}\n".format(
(average_hess - fisher_info) / average_hess))
VerifyGradientAndFIM()
[1]: Guo-Xun Yuan, Chia-Hua Ho and Chih-Jen Lin. An Improved GLMNET for L1-regularized Logistic Regression. Journal of Machine Learning Research, 13, 2012. http://www.jmlr.org/papers/volume13/yuan12a/yuan12a.pdf
[2]: skd. Derivation of Soft Thresholding Operator. 2018. https://math.stackexchange.com/q/511106
[3]: Wikipedia Contributors. Proximal gradient methods for learning. Wikipedia, The Free Encyclopedia, 2018. https://en.wikipedia.org/wiki/Proximal_gradient_methods_for_learning
[4]: Yao-Liang Yu. The Proximity Operator. https://www.cs.cmu.edu/~suvrit/teach/yaoliang_proximity.pdf