Inconsistent usage of the linear layer in the SABlock

[KumoLiu]

There are three linear in the original generative repo which now is combined within one linear layer. We need add an argument to make it consistent with before.

https://github.com/Project-MONAI/GenerativeModels/blob/main/generative/networks/nets/diffusion_model_unet.py#L379-L381

MONAI/monai/networks/blocks/selfattention.py

Line 89 in 56ee32e

self.qkv = nn.Linear(self.hidden_input_size, self.inner_dim * 3, bias=qkv_bias)

[ericspod]

Perhaps we change the current implementation to match this instead? I don't know if it's significant but making this an option will be problematic with Torchscript which doesn't like class variables being optional.

[KumoLiu]

I did a quick test, looks the outcome shows significantly different. In this PR, I utilized identity to get around the torchscript conversion issue. Could you please help review it, thanks!

import torch.nn as nn
import torch
from monai.utils import set_determinism

set_determinism(0)

hidden_input_size = 360
inner_dim = 360
qkv_bias = True
qkv = nn.Linear(hidden_input_size, inner_dim * 3, bias=qkv_bias)


q = nn.Linear(hidden_input_size, inner_dim, bias=qkv_bias)
k = nn.Linear(hidden_input_size, inner_dim, bias=qkv_bias)
v = nn.Linear(hidden_input_size, inner_dim, bias=qkv_bias)

input = torch.rand(2, 512, 360)
qkv_out = qkv(input)
q_out = q(input)
k_out = k(input)
v_out = v(input)

print((qkv_out - torch.cat([q_out, k_out, v_out], dim=-1)).abs().sum().item()) --> 410444.28125