RuntimeError: mat1 and mat2 shapes cannot be multiplied (16880x128 and 1x256), unable to fix this for an attention based layer

prayuj_sachdev: <p>I have an Attention-based feature fusion layer for two input feature embeddings. The model is defined here:</p>
<p>class AFF(nn.Module):<br>
def <strong>init</strong>(self, in_features, out_features):<br>
super(AFF, self).<strong>init</strong>()<br>
self.in_features = in_features<br>
self.out_features = out_features</p>
<pre><code># Attention mechanism
self.attention = nn.Linear(in_features * 2, 1)
self.attention.weight = nn.Parameter(torch.Tensor(self.in_features * 2, 1))
nn.init.xavier_uniform_(self.attention.weight) 
self.softmax = nn.Softmax(dim=1)
    
# Fusion layer
self.fc = nn.Linear(in_features * 2, out_features)
self.relu = nn.ReLU()
</code></pre>
<p>def forward(self, x1, x2):<br>
print(“Shape of x1:”, x1.shape)<br>
print(“Shape of x2:”, x2.shape)</p>
<pre><code>if x1.size(0) != x2.size(0):
  if x1.size(0) &lt; x2.size(0):
    padding = torch.nn.functional.pad(x1, (0,0,0, x2.size(0) - x1.size(0)), mode='constant', value=0)
    x_concat = torch.cat((padding, x2), dim=1)
  else:
    padding = torch.nn.functional.pad(x2, (0,0,0, x1.size(0) - x2.size(0)), mode='constant', value=0)
    x_concat = torch.cat((x1, padding), dim=1)
else:
  x_concat = torch.cat((x1, x2), dim=1)

print("Shape of x_concat:", x_concat.shape)

# Calculate attention weights
attn_weights = self.softmax(self.attention(x_concat))

# Apply attention to feature embeddings
x1_weighted = torch.mul(x1, attn_weights[:, :1])
x2_weighted = torch.mul(x2, attn_weights[:, 1:])

# Concatenate weighted feature embeddings
x_weighted_concat = torch.cat((x1_weighted, x2_weighted), dim=1)

# Fusion layer
fused_features = self.fc(x_weighted_concat)
fused_features = self.relu(fused_features)

return fused_features
</code></pre>
<p>How can I fix this error when I run it in my training loop?<br>
The shapes of x1, x2 and x_concat are:<br>
Shape of x1: torch.Size([16880, 64])<br>
Shape of x2: torch.Size([2956, 64])<br>
Shape of x_concat: torch.Size([16880, 128])</p>

ptrblck: <p>Your custom weight initialization is failing since the weight matrix inside the <code>nn.Linear</code> module has the shape <code>[out_features, in_features]</code>.<br>
Use:</p>
<pre data-code-wrap="python"><code class="lang-python">self.attention.weight = nn.Parameter(torch.Tensor(1, self.in_features * 2))
nn.init.xavier_uniform_(self.attention.weight) 
</code></pre>
<p>and this issue should be gone.<br>
Afterwards you will see a new failure:</p>
<pre data-code-wrap="python"><code class="lang-python">x2_weighted = torch.mul(x2, attn_weights[:, 1:])

RuntimeError: The size of tensor a (64) must match the size of tensor b (0) at non-singleton dimension 1
</code></pre>
<p>since you are again indexing the weight matrix in a wrong way.<br>
Fixing the indexing dimensions via:</p>
<pre data-code-wrap="python"><code class="lang-python">        # Apply attention to feature embeddings
        x1_weighted = torch.mul(x1, attn_weights[:1, :])
        x2_weighted = torch.mul(x2, attn_weights[1:, :])
</code></pre>
<p>won’t resolve this error since the shape will still not match:</p>
<pre data-code-wrap="python"><code class="lang-python">RuntimeError: The size of tensor a (2956) must match the size of tensor b (16879) at non-singleton dimension 0
</code></pre>
<p>so you would need to check how exactly these weights should be multiplied with <code>x2</code> since they are not compatible.</p>

prayuj_sachdev: <p>The changed weight initialization worked, leading to a successful matrix multiplication. When later multiplying the weights with the inputs, I padded the smaller matrix to make it so that the size matches the weight matrix.</p>