Hi @erroneus,
Finally got some time to dig through your DIM code and the paper. Very impressive and very thorough.
The code is obviously highly modular, but it makes it a little difficult to track how the DIM network gets constructed and the loss calculated.  I made a small, hard-coded network to implement the local/global DIM loss calculation for a specific tensor shape. This was written from your posted DIM code and the various figures in the paper (and appendix). Can you take a look and let me know if I’m understanding things correctly?
(I’m fairly new to all this, so apologies for any obviously dumb mistakes)
Input A & B: 256 x 256 x 256 tensor representing a 3D image volume (“Atlas” and “Moving” volumes)
Volumes are greyscale (single channel).
When I use actual image data, the network trains, and seems to generate numbers that accurately correlate with hard-calculated mutual information. Which is fantastic, but only if I’m actually doing this right. Appreciate any insight/critiques/etc. Thanks!
A = torch.randint(0, 256, (2, 1, 256, 256, 256)).float().requires_grad_(True)
B = torch.randint(0, 256, (2, 1, 256, 256, 256)).float().requires_grad_(True)

def get_positive_expectation(p_samples, average=True):
    #Measure = JSD (Simplified from DIM Code for clarity)
    log_2 = math.log(2.)
    Ep = log_2 - F.softplus(-p_samples)  # Note JSD will be shifted
    if average:
        return Ep.mean()
    else:
        return Ep

def get_negative_expectation(q_samples, average=True):
    #Measure = JSD (Simplified from DIM Code for clarity)
    log_2 = math.log(2.)
    Eq = F.softplus(-q_samples) + q_samples - log_2  # Note JSD will be shifted
    if average:
        return Eq.mean()
    else:
        return Eq

def loss_calc(lmap, gmap):
    #The fenchel_dual_loss from the DIM code
    #Reshape tensors dims to (N, Channels, chunks)
    lmap = lmap.reshape(2,128,-1)
    gmap = gmap.squeeze()
    
    N, units, n_locals = lmap.size()
    n_multis = gmap.size(2)

    # First we make the input tensors the right shape.
    l = lmap.view(N, units, n_locals)
    l = lmap.permute(0, 2, 1)
    l = lmap.reshape(-1, units)

    m = gmap.view(N, units, n_multis)
    m = gmap.permute(0, 2, 1)
    m = gmap.reshape(-1, units)
    
    u = torch.mm(m, l.t())
    u = u.reshape(N, n_multis, N, n_locals).permute(0, 2, 3, 1)
    
    mask = torch.eye(N).to(l.device)
    n_mask = 1 - mask
    
    E_pos = get_positive_expectation(u, average=False).mean(2).mean(2)
    E_neg = get_negative_expectation(u, average=False).mean(2).mean(2)
    
    E_pos = (E_pos * mask).sum() / mask.sum()
    E_neg = (E_neg * n_mask).sum() / n_mask.sum()
    loss = E_neg - E_pos
    
    return loss

class Mixed_Dim(torch.nn.Module):
    def __init__(self):
        super(Mixed_Dim, self).__init__()
        #Local Feature Map: Local feature map of size [N, 128, 8, 8, 8]
        self.lmapnet = Sequential(Conv3d(1, 8, 2, 2), ReLU(), BatchNorm3d(8),
                                  Conv3d(8, 16, 2, 2), ReLU(), BatchNorm3d(16),
                                  Conv3d(16, 32, 2, 2), ReLU(), BatchNorm3d(32),
                                  Conv3d(32, 64, 2, 2), ReLU(), BatchNorm3d(64),
                                  Conv3d(64, 128, 2, 2), ReLU(), BatchNorm3d(128) 
                                 )
        
        #Global Feature Map: Global feature map of size [N, 128, 1, 1, 1]
        self.gmapnet = Sequential(Conv3d(128, 128, 2, 2), ReLU(), BatchNorm3d(128),
                                  Conv3d(128, 128, 2, 2), ReLU(), BatchNorm3d(128),
                                  Conv3d(128, 128, 2, 2), ReLU()
                                 )
        #Per paper, global map is activated:
        self.gfc1 = Sequential(Linear(1, 512), ReLU(), BatchNorm3d(128),
                               Linear(512, 512))
        self.gfc2 = Sequential(Linear(1, 512), ReLU(), BatchNorm3d(128))
        
        #Per paper, local map is activated:
        self.lfc1 = Sequential(Conv3d(128, 128, 1, 1), ReLU(), BatchNorm3d(128),
                               Conv3d(128, 128, 1, 1))
        self.lfc2 = Sequential(Conv3d(128, 128, 1, 1), ReLU(), BatchNorm3d(128))
        self.Laynorm = LayerNorm([128, 8, 8, 8])
        
    def forward(self, moving, atlas):
        #Local feature maps of moving and atlas
        lmap_moving = self.lmapnet(moving)
        lmap_atlas = self.lmapnet(atlas)
        
        #Global feature map of atlas
        gmap_atlas = self.gmapnet(lmap_atlas)
        
        #Encode Global feature map
        gout1, gout2 = self.gfc1(gmap_atlas), self.gfc2(gmap_atlas)
        gmap_atlas_enc = self.gfc1(gmap_atlas) + self.gfc2(gmap_atlas)
        
        #Encode Local feature maps
        lmap_atlas_enc = self.Laynorm(self.lfc1(lmap_atlas) + self.lfc2(lmap_atlas))
        lmap_moving_enc = self.Laynorm(self.lfc1(lmap_moving) + self.lfc2(lmap_moving))
        
        return lmap_atlas_enc, lmap_moving_enc, gmap_atlas_enc
        
model_dim = Mixed_Dim().cuda().train()
optim = torch.optim.Adam(model_dim.parameters(), lr=0.01)
A = A.cuda()
B = B.cuda()
for epoch in range(100):
    loc_atlas, loc_moving, glob_atlas = model_dim(B, A)
    loss1 = loss_calc(loc_moving, glob_atlas)
    loss2 = loss_calc(loc_atlas, glob_atlas)
    loss = (loss2 + loss1)
    model_dim.zero_grad()
    print(loss.item())
    loss.backward()
    optim.step()