How to implement Laplace Posteriori Approximation on BERT in PyTorch?
I'm trying to implement the Laplace Posteriori Approximation on the last layer for the classification results obtained by BERT model. I get an error regarding input size, and after I fix it by extracting just embeddings and class labels from BERT to feed them into Laplace, I get another bunch of errors regarding input dimensions that I don't know how to debug.
As this is something I didn't find on the internet, and includes relatively new libraries, I will post here just the first error I got, code that might help in debugging and useful links.
I will update post if needed.
Of course, if someone knows how to implement Laplace Posteriori Approximation with BERT in some other library like Scikit or Trax, it would be helpful. Also, some other Transformer classification model with some other confidence approximation will be useful for me. Any help is appreciated!
Code:
# Import
import pandas as pd
import torch
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
from torch import nn
from transformers import BertTokenizer
from transformers import BertModel
from transformers import BertForSequenceClassification
from sklearn.model_selection import train_test_split
import time
import os


#Toy Data
data_a_b_c = [&quot;&quot;&quot;category a. This is category a. In category a we talk about animals. 
This category includes lions, fish, tigers, birds, elephants, mouses, dogs, cats, and all other animals.&quot;&quot;&quot;] * 60 \
+  [&quot;&quot;&quot;category b. This is category b. In category b we talk about people. This category members are
Abraham Maslow, John Lennon, Drazen Petrovic, Nikola Tesla, Slavoljub Penkala, Nenad Bakic and Larry Page.&quot;&quot;&quot;] * 60 \
+ [&quot;&quot;&quot;category c. This is category c. Category c is dedicated to car brands like Lamborgini, Rimac-Buggati, BMW, Mercedes,
Honda, Opel, Wolkswagen, and etc.&quot;&quot;&quot;] * 60

label_0_1_2 =  [0] * 60 + [1] * 60 + [2] * 60
d = {'text': data_a_b_c, 'labels': label_0_1_2}
df = pd.DataFrame(data=d)

print(df.head(3))
print(df.tail(3))
print(df.info())

# Parameters
tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
batch_size = 2
learning_rate = 3e-4
epochs = 3
device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
labels = pd.Series(df.labels.values).to_dict()
num_classes = 3

print(f'Tokenizer: {tokenizer}, Batch size:{batch_size}, Learning rate:{learning_rate}, Epochs:{epochs}')
print('Device: ', device)
print('Number of possible classes: ', num_classes)

# Model Architecture
class TransformerModel(nn.Module):
    def __init__(self, num_classes, dropout=0.5):
        super(TransformerModel, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-multilingual-cased')
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(768, num_classes)
        self.relu = nn.ReLU()

    def forward(self, input_id, mask):
        _, pooled_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False)
        dropout_output = self.dropout(pooled_output)
        linear_output = self.linear(dropout_output)
        final_layer = self.relu(linear_output)

        return final_layer

# Prepare Data Function
def prepare_data(data, labels):
    texts = tokenizer(data, padding='max_length', max_length=512, truncation=True, return_tensors=&quot;pt&quot;)
    input_ids = texts['input_ids']
    attention_mask = texts['attention_mask']
    train_dataset = TensorDataset(input_ids, attention_mask, torch.LongTensor(labels))
    dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    return dataloader

#Run Training Function
def run_training(train_dataloader, val_dataloader, epochs=epochs, lr=learning_rate):

    def train(dataloader):
        model.train()
        total_acc, total_count = 0, 0
        log_interval = 128
        start_time = time.time()

        for idx, (input_id, mask, label) in enumerate(train_dataloader):
           # print(idx)
            mask = mask.to(device)
            input_id = input_id.to(device)
            label = label.type(torch.LongTensor).to(device)

            output = model(input_id, mask)

            optimizer.zero_grad()
            loss = criterion(output, label)
            loss.backward()
           # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
            optimizer.step()

            total_acc += (output.argmax(1) == label).sum().item()
            total_count += label.size(0)

            if idx % log_interval == 0 and idx &gt; 0:
                elapsed = time.time() - start_time
                print('| epoch {:3d} | {:5d}/{:5d} batches '
                      '| accuracy {:8.3f}'.format(epoch, idx, len(dataloader),
                                                  total_acc / total_count))
                total_acc, total_count = 0, 0
                start_time = time.time()

    def evaluate(dataloader):
        model.eval()
        total_acc, total_count = 0, 0
        with torch.no_grad():
              for idx, (input_id, mask, label) in enumerate(dataloader):
                mask = mask.to(device)
                input_id = input_id.to(device)
                label = label.to(device)
                output = model(input_id, mask)
                total_acc += (output.argmax(1) == label).sum().item()
                total_count += label.size(0)
        return total_acc / total_count

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.1)

    cuda = torch.cuda.is_available()
    device = torch.device(&quot;cuda&quot; if cuda else &quot;cpu&quot;)
    device = 'cuda'
    model.to(device)

    total_accu = None
    for epoch in range(1, epochs + 1):
        epoch_start_time = time.time()
        train(train_dataloader)
        accu_val = evaluate(val_dataloader)
        if total_accu is not None and total_accu &gt; accu_val:
            scheduler.step()
        else:
            total_accu = accu_val
        print('-' * 59)
        print('| end of epoch {:3d} | time: {:5.2f}s | '
              'valid accuracy {:8.3f} '.format(epoch,
                                               time.time() - epoch_start_time,
                                               accu_val))
        print('-' * 59)

# Data Split And Preparation
X_train, X_test, y_train, y_test = train_test_split(df.text.values.tolist(), df.labels.values.tolist(), test_size=0.2, random_state=2)
train_dataloader = prepare_data(X_train, y_train)
val_dataloader = prepare_data(X_test, y_test)

# Run The Model
model = TransformerModel(num_classes) 
run_training(train_dataloader, val_dataloader)
print('finished')

# Save And Load The Model (if needed)
PATH = &quot;.../Torch_BERT_model&quot;
torch.save(model, os.path.join(PATH, &quot;Toy_Data_BERT.pth&quot;))
model = torch.load(os.path.join(PATH, &quot;Toy_Data_BERT.pth&quot;))
print(model)

# Laplace
from laplace import Laplace
la = Laplace(model, 'classification', subset_of_weights='last_layer', hessian_structure='full')
la.fit(train_dataloader)

Error I get:

--------------------------------------------------------------------------- ValueError                                Traceback (most recent call
last) ~\AppData\Local\Temp\ipykernel_7144\3779742208.py in &lt;cell line:
2&gt;()
1 la = Laplace(model, 'classification', subset_of_weights='last_layer', hessian_structure='full')
----&gt; 2 la.fit(train_dataloader)
~\anaconda3\lib\site-packages\laplace\lllaplace.py in fit(self,
train_loader, override)
98
99         if self.model.last_layer is None:
--&gt; 100             X, _ = next(iter(train_loader))
101             with torch.no_grad():
102                 try:
ValueError: too many values to unpack (expected 2)

Useful link for Laplace implementation with examples:
https://aleximmer.github.io/Laplace/#full-example-optimization-of-the-marginal-likelihood-and-prediction
Code that might help in debugging:
for x in train_dataloader:
        print(&quot;The length of batch is:&quot;, len(x))
        print()
        print(&quot;The batch looks like:&quot;, x)
        print()
        
        print(&quot;The length of the first element in the batch is:&quot;) #embedding
        print(len(x[0]))
        print(&quot;The length of the second element in the batch is:&quot;) #1 if place is filled with word, 0 if it's empty?
        print(len(x[1]))
        print(&quot;The length of the third element in the batch is:&quot;) #category
        print(len(x[2]))
        print()
        
        print(&quot;The lengths of the first tensor and second tensor in the first element in the batch is:&quot;)
        print(len(x[0][0]), len(x[0][1])) # = max_length (512)
        print(&quot;The lengths of the first tensor and second tensor in the second element in the batch is:&quot;)
        print(len(x[1][0]), len(x[1][1])) # = max_length (512)
        print()
        print()


The laplace library expects that the dataloader returns two parameters (X,y) and that the model requires exactly one argument to make its prediction (code). But your model forward pass requires two arguments, namely input_id and mask, and your dataloader returns three arguments input_id, mask, and labels.
There are several ways to work around this limitation (e.g. return a dict with input_ids and attention_mask). The way that requires the least understanding of the internals of the laplace library is to generate the attention mask at runtime in the forward pass (not great for the performance):
class TransformerModel(nn.Module):
    def __init__(self, num_classes, pad_id, dropout=0.5):
        super(TransformerModel, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-multilingual-cased')
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(768, num_classes)
        self.relu = nn.ReLU()

        self.pad_id = pad_id
    def forward(self, input_id):
        mask = (input_ids!=self.pad_id).type(input_ids.dtype)

        _, pooled_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False)
        dropout_output = self.dropout(pooled_output)
        linear_output = self.linear(dropout_output)
        final_layer = self.relu(linear_output)

        return final_layer

model = TransformerModel(num_classes, tokenizer.pad_token_id)

