Deep Regression Ensembles

DRE.py

import numpy as np
from opt_einsum import contract

class DRE:
    """Deep Regression Ensemble

    Parameters
    ----------
    n_layers : int
        Number of hidden layers.
        
    n_features : int
        Number of random features.
        
    shrinkages : np.array
        Ridge shrinkage parameters.
        
    kernel_scales : np.array
        Scales for random features.
        
    bias : float
        Bias for random features.

    References
    ----------
    Didisheim, Antoine, Bryan Kelly, and Semyon Malamud.
    "Deep Regression Ensembles." arXiv preprint arXiv:2203.05417 (2022).
    """
    def __init__(
        self,
        n_layers : int,
        n_features : int,
        shrinkages : np.array,
        kernel_scales : np.array,
        bias : float,
    ):
        assert n_layers >= 0
        assert n_features > 0
        
        assert len(shrinkages) > 0
        assert (shrinkages > 0).all()

        assert len(kernel_scales) > 0
        assert (kernel_scales > 0).all()
        assert bias >= 0
        
        self.n_layers = n_layers
        self.n_features = n_features
        self.shrinkages = shrinkages
        self.kernel_scales = kernel_scales
        self.bias = bias
        
    def fit(self, X : np.array, y : np.array):
        P = self.n_features
        K = len(self.kernel_scales)
        L = len(self.shrinkages)
        alphas = self.shrinkages.reshape(L, 1)

        layers = []
        x = X
        
        for l in range(self.n_layers):
            kernel = np.random.normal(size=(x.shape[1], P, K)) * self.kernel_scales
            bias = np.random.uniform(low=-self.bias, high=self.bias, size=(P, K))
            z = np.maximum(contract('nd,dpk->npk', x, kernel) + bias, 0)
            
            betas = []
            for i in range(K):
                u = z[..., i]
                w, W = np.linalg.eigh(u.T @ u)
                beta = W @ ((y @ u @ W) / (w + alphas)).T
                betas.append(beta)

            beta = np.stack(betas, axis=-1)
            x = contract('npk,plk->nlk', z, beta).reshape(-1, K * L)
            
            layers.append({'kernel': kernel, 'bias': bias, 'beta': beta})
            
        w, W = np.linalg.eigh(x.T @ x)
        beta = W @ ((y @ x @ W) / (w + alphas)).T
        
        self.layers_ = layers
        self.beta_ = beta
        return self
        
    def predict(self, X : np.array):
        shape = (-1, len(self.kernel_scales) * len(self.shrinkages))

        x = X
        for l in self.layers_:
            z = np.maximum(contract('nd,dpk->npk', x, l['kernel']) + l['bias'], 0)
            x = contract('npk,plk->nlk', z, l['beta']).reshape(shape)
        
        y_hat = x @ self.beta_
        return y_hat