<jupyter_start><jupyter_text>MathJax.Hub.Config({    TeX: { equationNumbers: { autoNumber: "AMS" } }}); IntroductionThough the Bayesian paradigm is theoretically appealing, it has proved difficult to apply in deep learning. Consider a neural network $f_\theta: X \to Y$. For illustrative purposes, suppose the model is given by$$p(y|x,\theta) \propto \exp\{-|| y - f_\theta(x) ||^2/2\}$$Let $\varphi(\theta)$ be the prior distribution and $p(\mathcal D_n |\theta) = \Pi_{i=1}^n p(y_i | x_i, \theta)$ be the likelihood of the data $\mathcal D_n = \{(x_i,y_i)\}_{i=1}^n$. Then the posterior distribution of $\theta$ is given by$$p(\theta | \mathcal D_n) \propto p(\mathcal D_n | \theta) \varphi(\theta).$$The normalizing constant in the posterior distribution is an intractable integral. Many methods for sampling from the posterior, e.g., MCMC, variational inference, will encounter extra difficulty when $\theta$ is a high dimensional neural network weight. Leaving aside the fact that the posterior is hard to sample, let's see why we should care about the posterior distribution in the first place. In the Bayesian paradigm, we predict using a distribution$$p(y|x, \mathcal D_n) = \int p(y|x,\theta) p(\theta|\mathcal D_n) \,d\theta$$Thus instead of providing a point estimate $\hat y$ for input $\hat x$, we can form an entire distribution estimate which lends itself naturally to uncertainty quantification. There is yet another advantage to Bayesian prediction, perhaps less well documented in the deep learning community. Neural networks are singular models. In terms of generalization, Bayesian prediction is better than MAP or MLE for singular models. This phenomenon is what we seek to elucidate in this work.  Generalization errorSuppose $\hat q_n(y|x)$ is some estimate of the true unknown conditional density $q(y|x)$. The generalization error of the prediction $\hat q(y|x)$ is defined as$$G(n) = KL (q(y|x) || \hat q(y|x) ) = \int \int q(y|x) \log \frac{q(y|x)}{\hat q_n(y|x)} \,dy q(x) \,dx$$We've written this in terms of the sample size $n$ to remind ourselves that $\hat q_n$ is formed using $\mathcal D_n$. We will consider three predictors of $q(y|x)$:+ Bayesian predictive distribution $\hat q(y|x) = p(y|x,\mathcal D_n)$+ MAP $\hat q(y|x) = p(y|x,\theta_n^{MAP}$)+ MLE $\hat q(y|x) = p(y|x,\theta_n^{MLE}$)To average out the randomness in the dataset $\mathcal D_n$ used to form the predictors, we will ultimately look at the average generalization error \begin{equation}E_n G(n) \label{eq:avgGn}\end{equation}where $E_n$ denotes expectation over the dataset $\mathcal D_n$. In simulations, we calculate the average generalization error using a held-out-test set $T_{n'} = \{(x_i',y_i')\}_{i=1}^{n'}$ as\begin{equation}\frac{1}{n'} \sum_{i=1}^{n'} \log q(y_i'|x_i') - E_n \frac{1}{n'} \sum_{i=1}^{n'} \log \hat q_n(y_i'|x_i')\label{eq:computed_avgGn}\end{equation}Assume the held-out test set is large enough so that the difference between \eqref{eq:avgGn} and \eqref{eq:computed_avgGn} is negligible. We will refer to them interchangeably as the average generalization error.  Bayes versus MAP versus MLE in singular modelsA singular model is ...Neural networks are singular models because ...Since $p(y|x,\theta)$ is singular, we look to Watanabe's singular learnng theory to understand the generalization error of the various predictors considered. Assume the true distribution is realizable by the model, i.e., $q(y|x) = p(y|x,\theta_0)$ for some $\theta_0$ (though we will investigate violations of this assumption in the experiments).We have the following asymptotic expansions of the generalization error for singular models. + For Bayes, $E_n G(n) = \lambda/n + o(1/n)$ where $\lambda \in \mathbb Q^+$ is a positive rational number known as the learning coefficient. The learning coefficient is completely determined by the the truth-model-prior triplet $( q(y|x), p(y|x,\theta), \varphi(\theta) )$. Most of the time $\lambda << dim(\theta)/2$+ For MAP and MLE, $E_n G(n) = C/n + o(1/n)$ (different $C$'s for MAP and MLE). Basically $C$ is the maximum of some Gaussian process, which can easily be greater than $dim(\theta)/2$. Watanabe's Main Theorem 6.4 gives the precise formulation.  Last layer Bayesian Though the Bayes predictive distribution is provably superior to MAP and MLE, it relies on the intractable posterior distribution. As a workaround, we form the predictive distribution only in the last layer(s) of a deep neural network. Since a neural network $f_\theta$ is naturally hierarchical, we break it up into $f_\theta = h_w \circ g_v$ where $\theta = (v,w)$ with $dim(w)$ being small enough to manageably perform MCMC. Let $\theta_{map} = (v_{map},w_{map})$. Let $\tilde x_i = g_{v_{map}}(x_i)$. Define a new transformed dataset $\tilde{\mathcal D_n} = \{(\tilde x_i, y_i) \}_{i=1}^n$. We perform MCMC to sample the posterior over $w$:$$p(w | \tilde{\mathcal D_n}) \propto p(\tilde{\mathcal D_n} | w) \varphi(w) = \Pi_{i=1}^n \exp\{-|| y_i - h_w \circ g_{v_{map}}(x_i) ||^2/2\} \varphi(w)$$Define the last layer Bayesian predictive distribution to be $$p_{LLB}(y|x, \mathcal D_n) = \int p(y|x,(v_{map},w)) p(w|\tilde{\mathcal D_n}) \,dw$$We have to be careful when we speak of the generalization error of the LLB predictive distribution. For a proper comparison to $E_n G_{map}(n)$ or $E_n G_{mle}(n)$, we have to look at \begin{equation}E_n G_{LLB}(n) =  KL (q(y|x) || p_{LLB}(y|x, \mathcal D_n) )\label{G_LLB}\end{equation}where $E_n$ averages out the randomness of both $D_n$ and  $v_{map}$. An alternative is to condition on $v_{map}$, using $g_{v_{map}}$ as a feature extractor in a preprocessing step. Then, assuming realizability $q(y|x) = p(y|x,(v_0,w_0))$ we may examine the generalizaton error\begin{equation}E_{\mathcal D_n | v_{map}} KL( p(y|x,(v_{map},w_0)) || p_{LLB}(y|x, \mathcal D_n) )\label{G_LLB_vfixed}\end{equation}The average generalization error in \eqref{G_LLB_vfixed} is distinctly different from the one for LLB in \eqref{G_LLB}. The nice thing about \eqref{G_LLB_vfixed} is that we know its asymptotic expansion is $\lambda/n$ where $\lambda$ corresponds to the triplet $( p(y|x,(v_{map},w_0)), p(y| x, (v_{map},w)), \varphi(w))$. For certain functions $h_w$ where the true $\lambda$ is known, we can verify this in the experiments. (Not implemented yet).So we have various (average) generalization errors we are interested in + $E_n G_{post}(n)$ corresponding to the full posterior over $\theta$+ $E_n G_{map}(n)$+ $E_n G_{mle}(n)$+ $E_n G_{LLB}(n)$+ Generalization error in \eqref{G_LLB_vfixed}. (don't have a name for it yet)Theory guarantees the following relationships:+ $E_n G_{post}(n)$ has smaller learning coefficient than either $E_n G_{map}(n)$ or $E_n G_{mle}(n)$+ \eqref{G_LLB_vfixed} = known $\lambda$/n + o(1/n) for certain $h_w$. The question is, what is the behavior of $E_n G_{LLB}(n)$? Two hunches+ Though LLB is different from full posterior predictive distribution, it's still better than map or mle, i.e., $E_n G_{LLB}(n) \le E_n G_{map}(n)$.+ $E_n G_{LLB}(n)$ is different from \eqref{G_LLB_vfixed} but perhaps asymptotially they share the same leading term $\lambda/n$?We'll set out to investigate these conjectures in the following experiment.  ExperimentsConsider a neural network where $h_w$ is a reduced-rank regression network and $g_v$ is a block of linear composed with ReLU layers.<jupyter_code>class Model(nn.Module):

    def __init__(self, input_dim, output_dim, ffrelu_hidden, rr_hidden):
        super(Model, self).__init__()

        self.feature_map = nn.Sequential(
            nn.Linear(input_dim, ffrelu_hidden),
            nn.ReLU(),
            nn.Linear(ffrelu_hidden, ffrelu_hidden),
            nn.ReLU(),
            nn.Linear(ffrelu_hidden, input_dim),
        )

        self.rr = nn.Sequential(
            nn.Linear(input_dim, rr_hidden, bias=False),
            nn.Linear(rr_hidden, output_dim, bias=False)
        )

    def forward(self, x):
        x = self.feature_map(x)
        return self.rr(x)<jupyter_output><empty_output><jupyter_text>We use the following values for the parameters+ input_dim = output_dim = rr_hidden = 3+ ffrelu_hidden = 5We generate the training data as follows: + $X \sim N(0,\sigma_x^2 I_3)$+ realizability    + realizable $Y \sim N(f_{\theta_0}(X),\sigma_y^2 I_3)$    + not realizable $Y \sim N(h_{w_0}(X),\sigma_y^2 I_3)$    For the testing data, we consider a possible change in the support of $X$:+ $X \sim \alpha N(0,\sigma_x^2 I_3)$ for some scale $\alpha \in \mathbb R^+$.<jupyter_code>def get_data(args):

    train_size = int(args.n)
    valid_size = int(args.n * 0.5)
    test_size = int(10000)

    a = Normal(0.0, 1.0)
    a_params = 0.2 * a.sample((args.input_dim, args.rr_hidden))
    b = Normal(0.0, 1.0)
    b_params = 0.2 * b.sample((args.rr_hidden, args.output_dim))

    X_rv = MultivariateNormal(torch.zeros(args.input_dim), torch.eye(args.input_dim))
    y_rv = MultivariateNormal(torch.zeros(args.output_dim), torch.eye(args.output_dim))
    true_model = Model(args.input_dim, args.output_dim, args.ffrelu_hidden, args.rr_hidden)
    true_model.eval()

    with torch.no_grad():
        # training +valid data
        X = X_rv.sample(torch.Size([train_size+valid_size]))
        if args.realizable:
            true_mean = true_model(X)
        else:
            true_mean = torch.matmul(torch.matmul(X, a_params), b_params)
        y = true_mean + y_rv.sample(torch.Size([train_size+valid_size]))
        dataset_train, dataset_valid = torch.utils.data.random_split(TensorDataset(X, y),[train_size,valid_size])

        # testing data
        X = args.X_test_std * X_rv.sample(torch.Size([test_size]))
        if args.realizable:
            true_mean = true_model(X)
        else:
            true_mean = torch.matmul(torch.matmul(X, a_params), b_params)
        y = true_mean + y_rv.sample(torch.Size([test_size]))
        dataset_test = TensorDataset(X, y)
        oracle_mse = (torch.norm(y - true_mean, dim=1)**2).mean()
        entropy = -torch.log((2 * np.pi) ** (-args.output_dim / 2) * torch.exp(
            -(1 / 2) * torch.norm(y - true_mean, dim=1) ** 2)).mean()

        train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=args.batchsize, shuffle=True)
        valid_loader = torch.utils.data.DataLoader(dataset_valid, batch_size=args.batchsize, shuffle=True)
        test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=args.batchsize, shuffle=True)

    return train_loader, valid_loader, test_loader, oracle_mse, entropy<jupyter_output><empty_output><jupyter_text>Note that the way MAP training is usually conducted in deep learning may involve early stopping. Watanabe's theory does not account for early stopping. We will try MAP training with and without early stopping.<jupyter_code>def map_train(args, X_train, Y_train, X_valid, Y_valid, X_test, Y_test, oracle_mse):

    model = Model(args.input_dim, args.output_dim, args.ffrelu_hidden, args.rr_hidden)
    opt = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, weight_decay=5e-4)
    early_stopping = EarlyStopping(patience=10, verbose=False, taskid=args.taskid)

    # TODO: is it necessary to implement mini batch SGD?
    for it in range(5000):
        model.train()
        y_pred = model(X_train).squeeze()
        l = (torch.norm(y_pred - Y_train, dim=1)**2).mean()
        l.backward()
        opt.step()
        opt.zero_grad()

        model.eval()
        with torch.no_grad():
            valid_loss = (torch.norm(model(X_valid).squeeze() - Y_valid, dim=1)**2).mean()
            
        if it % 100 == 0:
            model.eval()
            ytest_pred = model(X_test).squeeze()
            test_loss = (torch.norm(ytest_pred - Y_test, dim=1)**2).mean()
            print('MSE: train {:.3f}, validation {:.3f}, test {:.3f}, oracle on test set {:.3f}'.format(l.item(), valid_loss, test_loss.item(), oracle_mse))

        if args.early_stopping:
            early_stopping(valid_loss, model)
            if early_stopping.early_stop:
                print("Early stopping")
                break

    return model<jupyter_output><empty_output><jupyter_text>The last function we will need is for drawing samples from the last layers. Although the code supports implicit variational inference (and even explicit variational inference), only the MCMC (NUTS) will be considered in the experiments.<jupyter_code>def lastlayer_approxinf(model, args, X_train, Y_train, X_valid, Y_valid, X_test, Y_test):
    
    transformed_X_train = model.feature_map(X_train)
    transformed_X_valid = model.feature_map(X_valid)
    transformed_X_test = model.feature_map(X_test)

    transformed_train_loader = torch.utils.data.DataLoader(
        TensorDataset(Tensor(transformed_X_train), torch.as_tensor(Y_train, dtype=torch.long)),
        batch_size=args.batchsize, shuffle=True)
    transformed_valid_loader = torch.utils.data.DataLoader(
        TensorDataset(Tensor(transformed_X_valid), torch.as_tensor(Y_valid, dtype=torch.long)), 
        batch_size=args.batchsize, shuffle=True)

    if args.posterior_method == 'ivi':

        # parameters for train_implicitVI
        mc = 1
        beta_index = 0
        args.betas = [1.0]
        saveimgpath = None
        args.dataset = 'reducedrank_synthetic'
        args.H = args.rr_hidden

        # TODO: strip train_implicitVI to simplest possible inputs
        G = train_implicitVI(transformed_train_loader, transformed_valid_loader, args, mc, beta_index, saveimgpath)
        eps = torch.randn(args.R, args.epsilon_dim)
        sampled_weights = G(eps)
        list_of_param_dicts = weights_to_dict(args, sampled_weights)

        pred_prob = 0
        output_dim = transformed_X_test.shape[1]
        for param_dict in list_of_param_dicts:
            mean = torch.matmul(torch.matmul(transformed_X_test, param_dict['a']), param_dict['b'])
            pred_prob += (2 * np.pi) ** (-output_dim / 2) * torch.exp(-(1 / 2) * torch.norm(Y_test - mean, dim=1) ** 2)

    elif args.posterior_method == 'mcmc':

        wholex = transformed_train_loader.dataset[:][0]
        wholey = transformed_train_loader.dataset[:][1]
        beta = 1.0

        kernel = NUTS(conditioned_pyro_rr, adapt_step_size=True)
        mcmc = MCMC(kernel, num_samples=args.R, warmup_steps=args.num_warmup, disable_progbar=True)
        mcmc.run(pyro_rr, wholex, wholey, args.rr_hidden, beta)
        sampled_weights = mcmc.get_samples()

        pred_prob = 0
        output_dim = wholey.shape[1]
        for r in range(0, args.R):
            mean = torch.matmul(torch.matmul(transformed_X_test, sampled_weights['a'][r,:,:]), sampled_weights['b'][r,:,:])
            pred_prob += (2 * np.pi) ** (-output_dim / 2) * torch.exp(-(1 / 2) * torch.norm(Y_test - mean, dim=1) ** 2)

    return -torch.log(pred_prob / args.R).mean()<jupyter_output><empty_output><jupyter_text>We set the arguments below. We are interested in examining the following three factors+ realizability or not+ support of $X_{test}$+ early stopping or notFor the realizable plus non-early-stopping setting, we expect the learning coefficient for the last layer Bayes predictive distribution to match the theoretically known values of $\lambda$ for reduced rank regression. Watanabe's theory, as it stands, seems unable to say anything meaningful in the case $q$ is unrealizable by the model or in the case that MAP is trained using early stopping.The computational overhead comes from the following variables:+ num_warmup: burn-in for MCMC+ MCs: number of training datasets for assessing $E_n$+ num_n: number of sample sizes considered for drawing learning curve $1/n$ versus $E_n G(n)$+ R: number of MCMC samples from the posteriorWe've set these numbers to be very small for the purpose of running quickly in the notebook.<jupyter_code>class Args:

    taskid = 1

    input_dim = 3
    output_dim = 3
    X_test_std = 1.0 # play around with
    realizable = 1 # play around with

    ffrelu_hidden = 5
    rr_hidden = 3
    early_stopping = 0 # play around with

    posterior_method = 'mcmc'
    num_warmup = 10


    batchsize = 50
    #   epsilon_mc = 100
    #   epochs = 100
    #   pretrainDepochs = 100
    #   trainDepochs = 20
    #   n_hidden_D = 128
    #   num_hidden_layers_D = 1
    #   n_hidden_G = 128
    #   num_hidden_layers_G = 1
    #   lr_primal = 0.01
    #   lr_dual = 0.001

    MCs = 5
    R = 10
    num_n = 10
    
    no_cuda = True
    log_interval = 500

args=Args()

args.cuda = not args.no_cuda and torch.cuda.is_available()

# change to boolean
if args.early_stopping == 0:
    args.early_stopping = False
else:
    args.early_stopping = True
if args.realizable == 0:
    args.realizable = False
else:
    args.realizable = True

args.epsilon_dim = args.rr_hidden*(args.input_dim + args.output_dim)
# TODO: w_dim and total_param_count depend on model and shouldn't be hardcoded as follows
args.w_dim = args.rr_hidden*(args.input_dim + args.output_dim)
total_param_count = (args.input_dim + args.rr_hidden + args.input_dim) * args.rr_hidden + args.w_dim<jupyter_output><empty_output><jupyter_text>Below we plot the learning curve, i.e., $n$ versus $E_n G(n)$. (Actually we'll plot $1/n$ versus $E_n G(n)$ to easily assess whether the slope matches the theoretical $\lambda$.)<jupyter_code>avg_lastlayerbayes_gen_err = np.array([])
std_lastlayerbayes_gen_err = np.array([])
avg_map_gen_err = np.array([])
std_map_gen_err = np.array([])
avg_entropy = np.array([])
std_entropy = np.array([])

n_range = np.round(1/np.linspace(1/200,1/10000,args.num_n))

for n in n_range:

    map_gen_err = np.empty(args.MCs)
    lastlayerbayes_gen_err = np.empty(args.MCs)
    entropy_array = np.empty(args.MCs)

    args.n = n

    for mc in range(0, args.MCs):

        train_loader, valid_loader, test_loader, oracle_mse, entropy = get_data(args)
        entropy_array[mc] = entropy

        X_train = train_loader.dataset[:][0]
        Y_train = train_loader.dataset[:][1]
        X_valid = valid_loader.dataset[:][0]
        Y_valid = valid_loader.dataset[:][1]
        X_test = test_loader.dataset[:][0]
        Y_test = test_loader.dataset[:][1]

        model = map_train(args, X_train, Y_train, X_valid, Y_valid, X_test, Y_test, oracle_mse)

        model.eval()
        map_gen_err[mc] = -torch.log((2*np.pi)**(-args.output_dim /2) * torch.exp(-(1/2) * torch.norm(Y_test-model(X_test), dim=1)**2)).mean() - entropy

        Bmap = list(model.parameters())[-1]
        Amap = list(model.parameters())[-2]
        params = (args.input_dim, args.output_dim, np.linalg.matrix_rank(torch.matmul(Bmap, Amap).detach().numpy()), args.rr_hidden)
        trueRLCT = theoretical_RLCT('rr', params)
        print('true RLCT {}'.format(trueRLCT))

        lastlayerbayes_gen_err[mc] = lastlayer_approxinf(model, args, X_train, Y_train, X_valid, Y_valid, X_test, Y_test) - entropy

        print('n = {}, mc {}, gen error: map {}, bayes last layer {}'
              .format(n, mc, map_gen_err[mc], lastlayerbayes_gen_err[mc]))


    print('average gen error: MAP {}, bayes {}'
          .format(map_gen_err.mean(), lastlayerbayes_gen_err.mean()))

    avg_lastlayerbayes_gen_err = np.append(avg_lastlayerbayes_gen_err, lastlayerbayes_gen_err.mean())
    std_lastlayerbayes_gen_err = np.append(std_lastlayerbayes_gen_err, lastlayerbayes_gen_err.std())
    avg_map_gen_err = np.append(avg_map_gen_err, map_gen_err.mean())
    std_map_gen_err = np.append(std_map_gen_err, map_gen_err.std())
    avg_entropy = np.append(avg_entropy, entropy_array.mean())
    std_entropy = np.append(std_entropy, entropy_array.std())

print('avg LLB gen err {}, std {}'.format(avg_lastlayerbayes_gen_err, std_lastlayerbayes_gen_err))
print('avg MAP gen err {}, std {}'.format(avg_map_gen_err, std_map_gen_err))


if args.realizable:
    ols_map = OLS(avg_map_gen_err, 1 / n_range).fit()
    map_slope = ols_map.params[0]

    ols_llb = OLS(avg_lastlayerbayes_gen_err, 1 / n_range).fit()
    llb_intercept = 0.0
    llb_slope = ols_llb.params[0]
else:
    ols_map = OLS(avg_map_gen_err, add_constant(1 / n_range)).fit()
    map_slope = ols_map.params[1]

    ols_llb = OLS(avg_lastlayerbayes_gen_err, add_constant(1 / n_range)).fit()
    llb_intercept = ols_llb.params[0]
    llb_slope = ols_llb.params[1]

print('estimated RLCT {}'.format(llb_slope))
#
fig, ax = plt.subplots()
ax.errorbar(1/n_range, avg_lastlayerbayes_gen_err, yerr=std_lastlayerbayes_gen_err, fmt='-o', c='r', label='En G(n) for last layer Bayes predictive')
ax.errorbar(1/n_range, avg_map_gen_err, yerr=std_map_gen_err, fmt='-o', c='g', label='En G(n) for MAP')
plt.plot(1 / n_range, llb_intercept + llb_slope / n_range, 'r--', label='ols fit for last-layer-Bayes')
plt.xlabel('1/n')
plt.title('map slope {:.2f}, parameter count {}, LLB slope {:.2f}, true RLCT {}'.format(map_slope, total_param_count, llb_slope, trueRLCT))
plt.legend()
plt.savefig('taskid{}.png'.format(args.taskid))
plt.show()<jupyter_output><empty_output><jupyter_text>The graph produced plots $1/n$ versus average generalizaton error $E_n G(n)$ for MAP predictor and LLB predictive distribution. Vertical bars indicate 1 std deviation over the different Monte Carlo training-testing sets. The LLB learning coefficient (LLB slope) indicated in the title of the graph should match the "true RLCT." In addition, the MAP learning coefficient (map slope) should be bigger than the LLB slope according to theory.  Below we plot the learning curve, i.e., $n$ versus $E_n G(n)$. (Actually we'll plot $1/n$ versus $E_n G(n)$ to easily assess whether the slope matches the theoretical $\lambda$.)<jupyter_code>avg_lastlayerbayes_gen_err = np.array([])
std_lastlayerbayes_gen_err = np.array([])
avg_map_gen_err = np.array([])
std_map_gen_err = np.array([])
avg_entropy = np.array([])
std_entropy = np.array([])

n_range = np.round(1/np.linspace(1/200,1/10000,args.num_n))

for n in n_range:

    map_gen_err = np.empty(args.MCs)
    lastlayerbayes_gen_err = np.empty(args.MCs)
    entropy_array = np.empty(args.MCs)

    args.n = n

    for mc in range(0, args.MCs):

        train_loader, valid_loader, test_loader, oracle_mse, entropy = get_data(args)
        entropy_array[mc] = entropy

        X_train = train_loader.dataset[:][0]
        Y_train = train_loader.dataset[:][1]
        X_valid = valid_loader.dataset[:][0]
        Y_valid = valid_loader.dataset[:][1]
        X_test = test_loader.dataset[:][0]
        Y_test = test_loader.dataset[:][1]

        model = map_train(args, X_train, Y_train, X_valid, Y_valid, X_test, Y_test, oracle_mse)

        model.eval()
        map_gen_err[mc] = -torch.log((2*np.pi)**(-args.output_dim /2) * torch.exp(-(1/2) * torch.norm(Y_test-model(X_test), dim=1)**2)).mean() - entropy

        Bmap = list(model.parameters())[-1]
        Amap = list(model.parameters())[-2]
        params = (args.input_dim, args.output_dim, np.linalg.matrix_rank(torch.matmul(Bmap, Amap).detach().numpy()), args.rr_hidden)
        trueRLCT = theoretical_RLCT('rr', params)
        print('true RLCT {}'.format(trueRLCT))

        lastlayerbayes_gen_err[mc] = lastlayer_approxinf(model, args, X_train, Y_train, X_valid, Y_valid, X_test, Y_test) - entropy

        print('n = {}, mc {}, gen error: map {}, bayes last layer {}'
              .format(n, mc, map_gen_err[mc], lastlayerbayes_gen_err[mc]))


    print('average gen error: MAP {}, bayes {}'
          .format(map_gen_err.mean(), lastlayerbayes_gen_err.mean()))

    avg_lastlayerbayes_gen_err = np.append(avg_lastlayerbayes_gen_err, lastlayerbayes_gen_err.mean())
    std_lastlayerbayes_gen_err = np.append(std_lastlayerbayes_gen_err, lastlayerbayes_gen_err.std())
    avg_map_gen_err = np.append(avg_map_gen_err, map_gen_err.mean())
    std_map_gen_err = np.append(std_map_gen_err, map_gen_err.std())
    avg_entropy = np.append(avg_entropy, entropy_array.mean())
    std_entropy = np.append(std_entropy, entropy_array.std())

print('avg LLB gen err {}, std {}'.format(avg_lastlayerbayes_gen_err, std_lastlayerbayes_gen_err))
print('avg MAP gen err {}, std {}'.format(avg_map_gen_err, std_map_gen_err))


if args.realizable:
    ols_map = OLS(avg_map_gen_err, 1 / n_range).fit()
    map_slope = ols_map.params[0]

    ols_llb = OLS(avg_lastlayerbayes_gen_err, 1 / n_range).fit()
    llb_intercept = 0.0
    llb_slope = ols_llb.params[0]
else:
    ols_map = OLS(avg_map_gen_err, add_constant(1 / n_range)).fit()
    map_slope = ols_map.params[1]

    ols_llb = OLS(avg_lastlayerbayes_gen_err, add_constant(1 / n_range)).fit()
    llb_intercept = ols_llb.params[0]
    llb_slope = ols_llb.params[1]

print('estimated RLCT {}'.format(llb_slope))
#
fig, ax = plt.subplots()
ax.errorbar(1/n_range, avg_lastlayerbayes_gen_err, yerr=std_lastlayerbayes_gen_err, fmt='-o', c='r', label='En G(n) for last layer Bayes predictive')
ax.errorbar(1/n_range, avg_map_gen_err, yerr=std_map_gen_err, fmt='-o', c='g', label='En G(n) for MAP')
plt.plot(1 / n_range, llb_intercept + llb_slope / n_range, 'r--', label='ols fit for last-layer-Bayes')
plt.xlabel('1/n')
plt.title('map slope {:.2f}, parameter count {}, LLB slope {:.2f}, true RLCT {}'.format(map_slope, total_param_count, llb_slope, trueRLCT))
plt.legend()
plt.savefig('taskid{}.png'.format(args.taskid))
plt.show()<jupyter_output>MSE: train 3.095, validation 3.522, test 3.048, oracle on test set 3.008
MSE: train 3.063, validation 3.535, test 3.034, oracle on test set 3.008
MSE: train 3.044, validation 3.544, test 3.029, oracle on test set 3.008
MSE: train 3.030, validation 3.553, test 3.028, oracle on test set 3.008
MSE: train 3.021, validation 3.562, test 3.031, oracle on test set 3.008
MSE: train 3.015, validation 3.569, test 3.035, oracle on test set 3.008
MSE: train 3.010, validation 3.574, test 3.039, oracle on test set 3.008
MSE: train 3.007, validation 3.578, test 3.042, oracle on test set 3.008
MSE: train 3.005, validation 3.579, test 3.044, oracle on test set 3.008
MSE: train 3.003, validation 3.580, test 3.046, oracle on test set 3.008
MSE: train 3.001, validation 3.579, test 3.047, oracle on test set 3.008
MSE: train 2.999, validation 3.580, test 3.049, oracle on test set 3.008
MSE: train 2.997, validation 3.580, test 3.050, oracle on test set 3.008
MSE: train 2.995, validation 3.579, test 3.050, ora[...]
{'fraction_non_alphanumeric': 0.12818756184821498, 'fraction_numerical': 0.01819289030981198, 'mean_word_length': 3.972558667676003, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 28, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '488988', 'license': 'MIT', 'n_tokens_mistral': 9566, 'n_tokens_neox': 8980, 'n_words': 2466, 'path': 'notebooks/lastlayerbayesian.ipynb', 'repo_name': 'suswei/RLCT', 'size': '26274'}