Using Laplace in knowledge graph embedding models

[Demirrr]

Dear all,

Thank you for this open-source project. I did enjoy your paper and presentation. I wanted to use laplace API to calibrate knowledge graph embedding (KGE) models in the link prediction task. Yet, I could not do it. In below, I aimed to briefly described the situation

class DistMult(BaseKGE):
    def __init__(self, args):
        super().__init__(args.learning_rate)
        self.name = 'DistMult'
        self.loss = torch.nn.BCEWithLogitsLoss()
        self.embedding_dim = args.embedding_dim
        self.emb_ent_real = nn.Embedding(args.num_entities, args.embedding_dim)
        self.emb_rel_real = nn.Embedding(args.num_relations, args.embedding_dim) 
        ...
    def forward_k_vs_all(self, x):
        e1_idx, rel_idx = x[:, 0], x[:, 1]
        # (1)
        # (1.1) Real embeddings of head entities
        emb_head_real = self.input_dp_ent_real(self.bn_ent_real(self.emb_ent_real(e1_idx)))
        # (1.2) Real embeddings of relations
        emb_rel_real = self.input_dp_rel_real(self.bn_rel_real(self.emb_rel_real(rel_idx)))
        # (2) Compute logits
        return torch.mm(self.hidden_dropout(self.bn_hidden_real(emb_head_real * emb_rel_real)),
                        self.emb_ent_real.weight.transpose(1, 0))

where

class BaseKGE(pl.LightningModule): ...

forward_k_vs_all(self, x) returns a 2D torch.tensor. After the model is trained , the below part is executed

from laplace import Laplace
  la = Laplace(model, 'classification',
               subset_of_weights='all',
               hessian_structure='diag')
  la.fit(dataset.train_dataloader())

and NotImplementedError(NotImplementedError: Extension saving to diag_ggn_exact does not have an extension for Module <class 'core.models.real.DistMult'> is raised.

Similarly, using

la = Laplace(model, 'classification',subset_of_weights='subnetwork',hessian_structure='full')

raises TypeError: init() missing 1 required positional argument: 'subnetwork_indices'

I was wondering whether you have any suggestion pertaining the above descired problem.

Cheers!

[georgezefko]

For your question regarding the subnetwork you are missing an argument "subnetwork_indices". Before Laplace you need to apply the following steps first subnetwork_mask = ModuleNameSubnetMask(model, module_names=module_name (eg. 'fc1')) subnetwork_mask.select() subnetwork_indices = subnetwork_mask.indices

if you run it in cuda you might need to adjust the last part to

subnetwork_indices=subnetwork_mask.indices.type(torch.LongTensor)

[Demirrr]

@georgezefko Very appreciated. Problem is solved!

[edaxberger]

Thanks a lot @Demirrr for raising this issue and to @georgezefko for commenting and providing a solution to this (and thanks to both for your interest in and usage of our library, we really appreciate it) -- I'm glad that it works now, and hope that you'll get reasonable results; please let us know if there's anything else we can help with!

Just FYI, in addition to ModuleNameSubnetMask there exist other ways to select/specify subnetworks -- there are a few examples in the README of this repository (need to scroll a bit down to see the code examples). A list of all currently supported options can be found as subclasses of SubnetMask within the following file: https://github.com/AlexImmer/Laplace/blob/main/laplace/utils/subnetmask.py

@georgezefko The issue you mention when using CUDA will soon be resolved once PR #87 is merged, addressing issue #86 raised earlier -- thanks again for your interest and help!