[Python] What seems to be wrong with my ML code, implemented from a paper's descriptions?

I came across a (short and curt) technical report that claims to be SOTA on keyword spotting, but it didn't share its code and had a very short explanation of its network. My job requires me to implement this model, which I did below, but upon training it using the Google Speech Commands dataset (as used in the paper), the results are slightly better than random guessing (3.8~4% in a 35-class dataset).

This leads me to believe that I have done something wrong when implementing the code, but I cannot find what it is, even after hours of searching. Please, show me the error of my ways.

https://prnt.sc/75qtYX8sz-fX (Company policy forbids direct upload of images)
# Waveformer
# https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=10546828&tag=1

import torch
import torch.nn as nn


class CLCA(nn.Module):
def __init__(self, emb_dim, head_dim, mlp_dim, eps=1e-3, activation=None):
super(CLCA, self).__init__()
self.emb_dim = emb_dim
self.mlp_dim = mlp_dim
self.target_dim = mlp_dim // 2
self.head_dim = head_dim
self.num_heads = emb_dim // head_dim
self.eps = eps

# Convolutional Layers
self.x1layer = nn.Sequential(
nn.Conv1d(self.emb_dim, self.emb_dim, kernel_size=5, stride=4, padding=1, groups=emb_dim), # Depthwise
nn.Conv1d(self.emb_dim, self.head_dim * self.mlp_dim, kernel_size=1) # Pointwise
)
self.x2layer1 = nn.Conv1d(self.emb_dim, self.emb_dim, kernel_size=5, stride=2, padding=2, groups=emb_dim)
self.x2layer2 = nn.Conv1d(self.emb_dim, self.head_dim * self.mlp_dim, kernel_size=1)
self.x3layer = nn.Conv1d(self.emb_dim, self.target_dim, kernel_size=1) # output channels should be E'

self.outlayer = nn.Conv1d(self.head_dim * self.mlp_dim, self.target_dim, kernel_size=1)

# Activation Layer
self.activation = activation if activation is not None else torch.nn.ELU()

def linear_attention(self, Q, K, V):
"""
Linear attention from "Transformers are RNNs"(https://arxiv.org/pdf/2006.16236)
aram Q: Query, shape [B, H, L/2, D]
aram K: Key, shape [B, H, L/4, D]
aram V: Value, shape [B, H, L/4, D]
:return: linearized attention
"""
Q = self.activation(Q) + 1
K = self.activation(K) + 1

# # https://github.com/idiap/fast-transformers/blob/master/fast_transformers/attention/linear_attention.py

# My implementation
KV = torch.einsum("bhld,bhlm->bhmd", K, V) # result should be BHDD (bhxd,bhxd->bhdd)
Z = 1. / (torch.einsum("bhld,bhd->bhl", Q, K.sum(dim=2)) + self.eps) # sum over length
V = torch.einsum("bhld,bhmd,bhl->bhlm", Q, KV, Z)


return V

def forward(self, x):
# Embedding dimension = channel
B, E, L = x.size() # Batch, embedding dim, len(audiolen)

x1 = self.x1layer(x) # (B, H*D, L/4)
x1 = torch.reshape(x1, (B, self.head_dim, self.mlp_dim, -1))
x1 = x1.transpose(2, 3) # (B, H, L/4, D)

x2 = self.x2layer1(x) # (B, C, L/2)
x3 = self.x3layer(x2) # (B, C', L/2)

x2 = self.x2layer2(x2)
x2 = torch.reshape(x2, (B, self.head_dim, self.mlp_dim, -1))
x2 = x2.transpose(2, 3) # (B, H, L/2, D)

out = self.linear_attention(x2, x1, x1) # (B, H, L/2, D)
out = out.transpose(2, 3) # (B, H, D, L/2)
out = out.reshape([B, -1, L//2])
out = self.outlayer(out) # (B, C', L/2)
out = out + x3

return out.transpose(1, 2)


class PWFF(nn.Module):
def __init__(self, mlp_dim):
super(PWFF, self).__init__()
self.target_dim = mlp_dim // 2
self.ff = nn.Sequential(
nn.Conv1d(self.target_dim, self.target_dim*2, kernel_size=1),
nn.GELU(),
nn.Dropout(),
nn.Conv1d(self.target_dim*2, self.target_dim, kernel_size=1)
)

def forward(self, x):
x = x.transpose(1, 2) # (B, E, L)
x = self.ff(x)
return x.transpose(1, 2)


class WaveFormerBlock(nn.Module):
def __init__(self, emb_dim, seq_len, head_dim, mlp_dim):
super(WaveFormerBlock, self).__init__()
self.clca = CLCA(emb_dim, head_dim, mlp_dim)
self.pwff = PWFF(mlp_dim)
self.dclayer1 = nn.Conv1d(emb_dim, mlp_dim, kernel_size=1, stride=1)
self.dclayer2 = nn.Conv1d(seq_len, seq_len//2, kernel_size=2, stride=2)
self.relu = nn.ReLU()

def forward(self, x):
# X = B, E, S (E: embed dim, =Channel / S: sequence len, =L)
x1 = self.clca(x)
x1 = self.pwff(x1)

x2 = self.dclayer1(x) # B, E'*2, L
x2 = self.dclayer2(x2.transpose(1, 2)) # B, E', L/2
x2 = self.relu(x2) # B, S/2, E'
return (x1 + x2).transpose(1, 2) # B, E', S/2


class Waveformer(nn.Module):
def __init__(self, num_classes=12):
"""
Waveformer implemented, from https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=10546828
Expects input of shape (batch_size, 8192) or (batch_size, 1, 8192)
Note: WaveFormerBlock's mlp_dim should be twice of next layer's embed_dim
aram num_classes: number of classes to classify inputs, Google Speech Commands is set to 12
"""
super(Waveformer, self).__init__()
self.blocks = nn.Sequential(
WaveFormerBlock(1, 16384, 4, 8),
WaveFormerBlock(4, 8192, 4, 8),
WaveFormerBlock(4, 4096, 4, 16),
WaveFormerBlock(8, 2048, 8, 16),
WaveFormerBlock(8, 1024, 8, 32),
WaveFormerBlock(16, 512, 16, 32),
WaveFormerBlock(16, 256, 16, 64),
WaveFormerBlock(32, 128, 32, 64),
WaveFormerBlock(32, 64, 32, 128),
WaveFormerBlock(64, 32, 32, 128),
)
self.pool = nn.AdaptiveAvgPool1d(1)
self.head = nn.Linear(64, num_classes)

def forward(self, x):
if len(x.shape) == 2:
x.unsqueeze(1)

# X shape (B, E, S) (64, 1, 16384)
x = self.blocks(x) # 64, 64, 16
x = self.pool(x).squeeze(-1) # 64, 64
return self.head(x)

Continue reading...