Hi,
I’m working on a problem which involves learning taking a vector X = [1, c, x_1, ..., x_N] (currently N ~ 100, but much more in the future) and learning a specific quadratic form, Y'PY + b, where Y = [1, c, X_bar] and X_bar is the average of the x's.
The current way I’ve been doing this is by using a bilinear layer with a bias term, e.g. something like:
class vf_bilin_3(nn.Module):
    def __init__(self):
        super().__init__()
        self.Pb = nn.Bilinear(3, 3, 1, bias=True) # Y -> Y'PY + b

    def forward(self, x):
        # forward pass
        M = x.shape[0]
        X_bar = torch.mean(x[:, 2:], dim=1, keepdims=True) 
        Y = torch.cat((x[:, 0:2].reshape([M, 2]), X_bar), dim=1) # [1, c, X_bar]
        output = self.Pb(Y, Y)
        return output

But I’m concerned that this won’t scale well when we have large numbers of observations per batch, or huge N, or on the GPU, etc.
I’d be very grateful if someone could point me in the right direction, or provide any optimization tips for this network.
– Arnav