Repository: cornellius-gp/gpytorch

Issue:
[Docs] Deep Sigma Point Processes
Hi,

The DSPPs documentaiton makes it seem like theyare coded up very similarly to the standard Deep GPs, and this does seem true for[ two layers](https://docs.gpytorch.ai/en/stable/examples/05_Deep_Gaussian_Processes/Deep_Sigma_Point_Processes.html). 

However, if I try to add a third layer (a second hidden layer) and update the input and output dims and make sure the forward call is correct, the loss function returns a vector rather than a single value (which the optimsier obviously doesn't like).

I'm guessing these losses are meant to be weighted according to the quadrature points. Is this true? Perhpas this could be more clear in the documentation (or alternatively, should this be done automatically in the backend?)


Discussion:
Hi @EvanBaker57 sorry for the slow reply. Would you mind posting a small reproducible code example?
No worries,

Below is a basic extension of the DSPP in the documentation (its necessarily a bit long). Note that this all works fine when dropped back to a two layer DSPP (because its essentially the tutorial example)

```python
import numpy as np
from matplotlib import pyplot as plt
import gpytorch
import torch
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.means import ConstantMean, LinearMean
from gpytorch.kernels import ScaleKernel, MaternKernel
from gpytorch.variational import VariationalStrategy, BatchDecoupledVariationalStrategy
from gpytorch.variational import MeanFieldVariationalDistribution
from gpytorch.models.deep_gps.dspp import DSPPLayer, DSPP
import gpytorch.settings as settings


#get some fake data
train_n = 100
x= np.linspace(0,1, train_n)
y = np.sin(2*x) + np.random.normal(size = train_n, loc = 0, scale = 0.1)
train_x = torch.from_numpy(x.reshape(-1,1)).float() 
train_y = torch.from_numpy(y).float() 


batch_size = 1000                 # Size of minibatch
num_inducing_pts = 300            # Number of inducing points in each hidden layer
num_epochs = 1000                # Number of epochs to train for
initial_lr = 0.01                 # Initial learning rate
hidden_dim = 3                    # Number of GPs (i.e., the width) in the hidden layer.
num_quadrature_sites = 8          # Number of quadrature sites (see paper for a description of this. 5-10 generally works well).

milestones = (np.linspace(0, num_epochs, 5)[1:-1]).astype(int).tolist()      # Epochs at which we will lower the learning rate by a factor of 0.1


#minibatch
from torch.utils.data import TensorDataset, DataLoader
train_dataset = TensorDataset(train_x, train_y)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)


from scipy.cluster.vq import kmeans2

# Use k-means to initialize inducing points (only helpful for the first layer)
inducing_points = (train_x[torch.randperm(min(1000 * 100, train_n))[0:num_inducing_pts], :])
inducing_points = inducing_points.clone().data.cpu().numpy()
inducing_points = torch.tensor(kmeans2(train_x.data.cpu().numpy(),
                               inducing_points, minit='matrix')[0])


class DSPPHiddenLayer(DSPPLayer):
    def __init__(self, input_dims, output_dims, num_inducing=300, inducing_points=None, mean_type='constant', Q=8):
        if inducing_points is not None and output_dims is not None and inducing_points.dim() == 2:
            # The inducing points were passed in, but the shape doesn't match the number of GPs in this layer.
            # Let's assume we wanted to use the same inducing point initialization for each GP in the layer,
            # and expand the inducing points to match this.
            inducing_points = inducing_points.unsqueeze(0).expand((output_dims,) + inducing_points.shape)
            inducing_points = inducing_points.clone() + 0.01 * torch.randn_like(inducing_points)
        if inducing_points is None:
            # No inducing points were specified, let's just initialize them randomly.
            if output_dims is None:
                # An output_dims of None implies there is only one GP in this layer
                # (e.g., the last layer for univariate regression).
                inducing_points = torch.randn(num_inducing, input_dims)
            else:
                inducing_points = torch.randn(output_dims, num_inducing, input_dims)
        else:
            # Get the number of inducing points from the ones passed in.
            num_inducing = inducing_points.size(-2)

        # Let's use mean field / diagonal covariance structure.
        variational_distribution = MeanFieldVariationalDistribution(
            num_inducing_points=num_inducing,
            batch_shape=torch.Size([output_dims]) if output_dims is not None else torch.Size([])
        )

        # Standard variational inference.
        variational_strategy = VariationalStrategy(
            self,
            inducing_points,
            variational_distribution,
            learn_inducing_locations=True
        )

        batch_shape = torch.Size([]) if output_dims is None else torch.Size([output_dims])

        super(DSPPHiddenLayer, self).__init__(variational_strategy, input_dims, output_dims, Q)

        if mean_type == 'constant':
            # We'll use a constant mean for the final output layer.
            self.mean_module = ConstantMean(batch_shape=batch_shape)
        elif mean_type == 'linear':
            # As in Salimbeni et al. 2017, we find that using a linear mean for the hidden layer improves performance.
            self.mean_module = LinearMean(input_dims, batch_shape=batch_shape)

        self.covar_module = ScaleKernel(MaternKernel(batch_shape=batch_shape, ard_num_dims=input_dims),
                                        batch_shape=batch_shape, ard_num_dims=None)

    def forward(self, x, mean_input=None, **kwargs):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
    
    
class ThreeLayerDSPP(DSPP):
    def __init__(self, train_x_shape, inducing_points, num_inducing, hidden_dim=3, Q=3):
        hidden_layer1 = DSPPHiddenLayer(
            input_dims=train_x_shape[-1],
            output_dims=hidden_dim,
            mean_type='linear',
            inducing_points=inducing_points,
            Q=Q,
        )
        hidden_layer2 = DSPPHiddenLayer(
            input_dims=hidden_layer1.output_dims,
            output_dims=hidden_dim,
            mean_type='linear',
            inducing_points=None,
            num_inducing = num_inducing,
            Q=Q,
        )
        last_layer = DSPPHiddenLayer(
            input_dims=hidden_layer2.output_dims,
            output_dims=None,
            mean_type='constant',
            inducing_points=None,
            num_inducing=num_inducing,
            Q=Q,
        )

        likelihood = GaussianLikelihood()

        super().__init__(Q)
        self.likelihood = likelihood
        self.last_layer = last_layer
        self.hidden_layer2 = hidden_layer2
        self.hidden_layer1 = hidden_layer1


    def forward(self, inputs, **kwargs):
        hidden_rep1 = self.hidden_layer1(inputs, **kwargs)
        hidden_rep2 = self.hidden_layer2(hidden_rep1, **kwargs)
        output = self.last_layer(hidden_rep2, **kwargs)
        return output

    def predict(self, loader):
        with settings.fast_computations(log_prob=False, solves=False), torch.no_grad():
            mus, variances = [], []
            for x_batch in loader:
                preds = self.likelihood(self(x_batch[0], mean_input=x_batch[0]))
                mus.append(preds.mean.cpu())
                variances.append(preds.variance.cpu())

        return torch.cat(mus, dim=-1), torch.cat(variances, dim=-1)
    
    
model = ThreeLayerDSPP(
    train_x.shape,
    inducing_points,
    num_inducing=num_inducing_pts,
    hidden_dim=hidden_dim,
    Q=num_quadrature_sites
)

model.train()

from gpytorch.mlls import DeepPredictiveLogLikelihood
adam = torch.optim.Adam([{'params': model.parameters()}], lr=initial_lr, betas=(0.9, 0.999))
sched = torch.optim.lr_scheduler.MultiStepLR(adam, milestones=milestones, gamma=0.1)


# The "beta" parameter here corresponds to \beta_{reg} from the paper, and represents a scaling factor on the KL divergence
# portion of the loss.
objective = DeepPredictiveLogLikelihood(model.likelihood, model, num_data=train_n, beta=0.05)

epochs_iter = range(num_epochs)
losses = []

for i in epochs_iter:
    minibatch_iter = train_loader
    for x_batch, y_batch in minibatch_iter:
        adam.zero_grad()
        output = model(x_batch)
        loss = -objective(output, y_batch)
        loss.backward()
        adam.step()
    sched.step()
    if i % 10 == 0:
        print('Iter %d/%d - Loss: %.3f      Noise: %.5f' % (
            i + 1, num_epochs, loss.item(), model.likelihood.noise_covar.noise))
    losses.append(loss.item())
    
plt.plot(range(0, len(losses)), losses[0:])
plt.show()
     
n_test = 1000
x_test = np.linspace(-0.2,1.2, num = n_test)
test_x = torch.from_numpy(x_test.reshape(-1,1)).float()

#do preds
test_dataset = TensorDataset(test_x)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)    
    
    
model.eval()
means, vars = model.predict(test_loader)
weights = model.quad_weights.unsqueeze(-1).exp().cpu()
mean = (weights * means).sum(0)
var = (weights * vars).sum(0) + (weights * means**2).sum(0) - (weights * means).sum(0)**2

with torch.no_grad(): 
    plt.plot(test_x[:,0], mean, zorder = 5)
    plt.fill_between(test_x[:,0], mean - 2*np.sqrt(var), mean + 2*np.sqrt(var), alpha = 0.5, zorder = 4)
    plt.scatter(train_x, train_y, color = "black", zorder = 1) #(zorder makes the points be at the top)
    plt.show()
```

Resolving patch:
diff --git a/gpytorch/models/deep_gps/dspp.py b/gpytorch/models/deep_gps/dspp.py
--- a/gpytorch/models/deep_gps/dspp.py
+++ b/gpytorch/models/deep_gps/dspp.py
@@ -43,24 +43,18 @@ def __init__(self, variational_strategy, input_dims, output_dims, num_quad_sites
         else:
             self.quad_sites = torch.nn.Parameter(torch.randn(num_quad_sites, input_dims))
 
-    def __call__(self, inputs, are_samples=False, expand_for_quadgrid=True, **kwargs):
+    def __call__(self, inputs, **kwargs):
         if isinstance(inputs, MultitaskMultivariateNormal):
-            # inputs is definitely in the second layer, and mean is n x t
+            # This is for subsequent layers. We apply quadrature here
+            # Mean, stdv are q x ... x n x t
             mus, sigmas = inputs.mean, inputs.variance.sqrt()
+            qg = self.quad_sites.view([self.num_quad_sites] + [1] * (mus.dim() - 2) + [self.input_dims])
+            sigmas = sigmas * qg
+            inputs = mus + sigmas  # q^t x n x t
+            deterministic_inputs = False
+        else:
+            deterministic_inputs = True
 
-            if expand_for_quadgrid:
-                xi_mus = mus.unsqueeze(0)  # 1 x n x t
-                xi_sigmas = sigmas.unsqueeze(0)  # 1 x n x t
-            else:
-                xi_mus = mus
-                xi_sigmas = sigmas
-
-            # unsqueeze sigmas to 1 x n x t, locations from [q] to Q^T x 1 x T.
-            # Broadcasted result will be Q^T x N x T
-            qg = self.quad_sites.view([self.num_quad_sites] + [1] * (xi_mus.dim() - 2) + [self.input_dims])
-            xi_sigmas = xi_sigmas * qg
-
-            inputs = xi_mus + xi_sigmas  # q^t x n x t
         if settings.debug.on():
             if not torch.is_tensor(inputs):
                 raise ValueError(
@@ -78,9 +72,15 @@ def __call__(self, inputs, are_samples=False, expand_for_quadgrid=True, **kwargs
         if self.output_dims is not None:
             inputs = inputs.unsqueeze(-3)
             inputs = inputs.expand(*inputs.shape[:-3], self.output_dims, *inputs.shape[-2:])
+
         # Now run samples through the GP
         output = ApproximateGP.__call__(self, inputs, **kwargs)
 
+        # If this is the first layer (deterministic inputs), expand the output
+        # This allows quadrature to be applied to future layers
+        if deterministic_inputs:
+            output = output.expand(torch.Size([self.num_quad_sites]) + output.batch_shape)
+
         if self.num_quad_sites > 0:
             if self.output_dims is not None and not isinstance(output, MultitaskMultivariateNormal):
                 mean = output.loc.transpose(-1, -2)


Test patch:
diff --git a/test/examples/test_dspp_regression.py b/test/examples/test_dspp_regression.py
new file mode 100644
--- /dev/null
+++ b/test/examples/test_dspp_regression.py
@@ -0,0 +1,136 @@
+import gpytorch
+import torch
+import unittest
+from gpytorch.likelihoods import GaussianLikelihood
+from gpytorch.means import ConstantMean, LinearMean
+from gpytorch.kernels import ScaleKernel, MaternKernel
+from gpytorch.variational import VariationalStrategy
+from gpytorch.variational import MeanFieldVariationalDistribution
+from gpytorch.models.deep_gps.dspp import DSPPLayer, DSPP
+from gpytorch.test.base_test_case import BaseTestCase
+
+
+train_n = 20
+train_x = torch.linspace(0, 1, train_n).unsqueeze(-1)
+train_y = torch.sin(2 * train_x[..., 0]) + torch.randn(train_n).mul_(0.01)
+test_n = 10
+test_x = torch.linspace(-0.2, 1.2, test_n).unsqueeze(-1)
+test_y = torch.sin(2 * test_x[..., 0])
+
+num_inducing_pts = 10            # Number of inducing points in each hidden layer
+num_epochs = 100                 # Number of epochs to train for
+initial_lr = 0.1                 # Initial learning rate
+hidden_dim = 3                   # Number of GPs (i.e., the width) in the hidden layer.
+num_quadrature_sites = 8         # Number of quadrature sites (see paper for a description of this.
+
+
+class DSPPHiddenLayer(DSPPLayer):
+    def __init__(self, input_dims, output_dims, num_inducing=300, inducing_points=None, mean_type='constant', Q=8):
+        if output_dims is None:
+            # An output_dims of None implies there is only one GP in this layer
+            # (e.g., the last layer for univariate regression).
+            inducing_points = torch.randn(num_inducing, input_dims)
+        else:
+            inducing_points = torch.randn(output_dims, num_inducing, input_dims)
+
+        # Let's use mean field / diagonal covariance structure.
+        variational_distribution = MeanFieldVariationalDistribution(
+            num_inducing_points=num_inducing,
+            batch_shape=torch.Size([output_dims]) if output_dims is not None else torch.Size([])
+        )
+
+        # Standard variational inference.
+        variational_strategy = VariationalStrategy(
+            self,
+            inducing_points,
+            variational_distribution,
+            learn_inducing_locations=True
+        )
+
+        batch_shape = torch.Size([]) if output_dims is None else torch.Size([output_dims])
+        super().__init__(variational_strategy, input_dims, output_dims, Q)
+
+        if mean_type == 'constant':
+            self.mean_module = ConstantMean(batch_shape=batch_shape)
+        elif mean_type == 'linear':
+            self.mean_module = LinearMean(input_dims, batch_shape=batch_shape)
+
+        self.covar_module = ScaleKernel(MaternKernel(batch_shape=batch_shape, ard_num_dims=input_dims),
+                                        batch_shape=batch_shape, ard_num_dims=None)
+
+    def forward(self, x, mean_input=None, **kwargs):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
+
+class ThreeLayerDSPP(DSPP):
+    def __init__(self, train_x_shape, num_inducing, hidden_dim=3, Q=3):
+        hidden_layer1 = DSPPHiddenLayer(
+            input_dims=train_x_shape[-1],
+            output_dims=hidden_dim,
+            num_inducing=num_inducing,
+            mean_type='linear',
+            Q=Q,
+        )
+        hidden_layer2 = DSPPHiddenLayer(
+            input_dims=hidden_layer1.output_dims,
+            output_dims=hidden_dim,
+            mean_type='linear',
+            num_inducing=num_inducing,
+            Q=Q,
+        )
+        last_layer = DSPPHiddenLayer(
+            input_dims=hidden_layer2.output_dims,
+            output_dims=None,
+            mean_type='constant',
+            num_inducing=num_inducing,
+            Q=Q,
+        )
+
+        likelihood = GaussianLikelihood()
+
+        super().__init__(Q)
+        self.likelihood = likelihood
+        self.last_layer = last_layer
+        self.hidden_layer2 = hidden_layer2
+        self.hidden_layer1 = hidden_layer1
+
+    def forward(self, inputs, **kwargs):
+        hidden_rep1 = self.hidden_layer1(inputs, **kwargs)
+        hidden_rep2 = self.hidden_layer2(hidden_rep1, **kwargs)
+        output = self.last_layer(hidden_rep2, expand_for_quadgrid=False, **kwargs)
+        return output
+
+
+class TestSGPRRegression(unittest.TestCase, BaseTestCase):
+    seed = 0
+
+    def test_dspp(self):
+        model = ThreeLayerDSPP(
+            train_x.shape,
+            num_inducing=num_inducing_pts,
+            hidden_dim=hidden_dim,
+            Q=num_quadrature_sites
+        )
+
+        model.train()
+
+        adam = torch.optim.Adam([{'params': model.parameters()}], lr=initial_lr, betas=(0.9, 0.999))
+        objective = gpytorch.mlls.DeepPredictiveLogLikelihood(model.likelihood, model, num_data=train_n, beta=0.05)
+
+        epochs_iter = range(num_epochs)
+        losses = []
+
+        for i in epochs_iter:
+            adam.zero_grad()
+            output = model(train_x)
+            loss = -objective(output, train_y)
+            loss.backward()
+            adam.step()
+            losses.append(loss.item())
+
+        model.eval()
+        preds = model.likelihood(model(test_x))
+        self.assertLess((preds.mean - test_y).abs().mean().item(), 0.5)
+        self.assertLess(preds.variance.mean().item(), 0.5)