Как определить морфологию слова исключительно по контексту?

Существует стандартный метод определения морфологии (pos-tagging), на основе lstm, которому на вход передается:

seq: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]  # vocab idx
out: [0, 0, 1, 2, 0, 1, 3, 2, 0, 1]   # target idx

где 0 — неизвестно и передано параметром в функцию потерь. Точность высокая. Но алгоритм существенно опирается на вектор тестируемого слова. Нужно же определять исключительно по контексту.

Переработал, чтобы исходные данные передавались как срез контекста, с заменой тестируемого слова на <pad> заглушку в seq и обнулением всех слов кроме целевого в out:

seq: [1, 2, 0, 4, 5]
out: [0, 0, 1, 0, 0]

seq: [1, 2, 3, 0, 5, 6]
out: [0, 0, 0, 2, 0, 0]

seq: [1, 2, 3, 4, 5, 0, 7, 8]
out: [0, 0, 0, 0, 0, 1, 0, 0]

...

seq: [5, 6, 7, 8, 9, 0]
out: [0, 0, 0, 0, 0, 1]

Но скорость обучения, соответственно существенно упала. Есть ли способ настроить обучение, чтобы модель обрабатывала предложение целиком, но игнорировала вектора собственно тестируемых слов?

class LSTMTagger(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, tagset_size):
        super(LSTMTagger, self).__init__()
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, bidirectional=True)
        self.hidden2tag = nn.Linear(hidden_dim*2, tagset_size) # *2 — bidirect

    def forward(self, embeds):
        length = len(embeds)
        lstm_out, _ = self.lstm(embeds.view(length, 1, -1))
        tag_space = self.hidden2tag(lstm_out.view(length, -1))
        tag_scores = F.log_softmax(tag_space, dim=1)
        return tag_scores

embedding = nn.Embedding.from_pretrained(weights, padding_idx=0)
model = LSTMTagger(embedding.embedding_dim, HIDDEN_DIM, TAGSET_SIZE)
loss_function = nn.NLLLoss(ignore_index=0)
optimizer = optim.SGD(model.parameters(), lr=0.1)

for source, target in corpus:
    embeds = embedding(source)
    model.zero_grad()
    tag_scores = model(embeds)
    loss = loss_function(tag_scores, target)
    loss.backward()
    optimizer.step()