Forward-step attention for TTS

ForwardStepAttention is a PyTorch attention module for Tacotron-style text-to-speech. It keeps content-based attention and adds a learned bias toward moving forward through the text.

Problem

A Tacotron-style decoder produces audio one step at a time. At each step, attention chooses which input characters the decoder reads. For speech, that choice should move through the text in order: the position at one step is at or just past the position at the previous step.

Content-based attention scores each character from the decoder state and the character's encoding. Nothing in the score encodes order, so alignment can fail in three ways (Figure 1).

Skip. Attention jumps ahead and some characters are never read.

Repeat. Attention returns to earlier characters and the speech repeats.

Stuck. Attention stays on one character and the decoder keeps producing audio for it.

Approach

The module adds a location prior to the content scores. The prior is built from the previous step's attention weights.

Step probabilities. A linear layer on the decoder state predicts three probabilities: stay, advance one character, advance two.

Shift. The previous attention weights are shifted by 0, 1 and 2 characters. Each copy is weighted by its probability, and the sum is the prior. Padding is masked out and the prior is renormalized.

Escape. A small fraction of uniform weight is mixed into the prior. Without it, a character far from the previous position has almost no prior, and content scores cannot recover it.

Scoring. Each character's score is its content score plus prior_strength times the log of the prior. The attention weights are the Boltzmann distribution over these scores. Adding a log is the same as multiplying probabilities. A strength of 0 gives plain content-based attention.

Table 1. Constructor parameters.

Implementation

import torch
from torch import nn
from torch.nn import functional as F


class ForwardStepAttention(nn.Module):
    def __init__(
        self,
        query_dim: int = 256,
        text_dim: int = 256,
        attention_dim: int = 128,
        prior_strength: float = 0.5,
        escape_probability: float = 0.02,
    ):
        super().__init__()
        self.query_projection = nn.Linear(query_dim, attention_dim, bias=False)
        self.text_projection = nn.Linear(text_dim, attention_dim, bias=False)
        self.score_projection = nn.Linear(attention_dim, 1, bias=False)

        self.step_projection = nn.Linear(query_dim, 3)  # stay, +1, +2

        self.prior_strength = prior_strength
        self.escape_probability = escape_probability

    def forward(
        self,
        query: torch.Tensor,  # [batch, query_dim]
        text_states: torch.Tensor,  # [batch, characters, text_dim]
        previous_alpha: torch.Tensor,  # [batch, characters]
        text_mask: torch.Tensor,  # [batch, characters], bool, True = real character
    ):
        batch, characters, _ = text_states.shape

        scores = self.score_projection(
            torch.tanh(
                self.query_projection(query)[:, None, :]
                + self.text_projection(text_states)
            )
        ).squeeze(-1)

        step_probs = F.softmax(self.step_projection(query), dim=-1)
        prior = torch.zeros_like(previous_alpha)

        for step in range(3):
            prior[:, step:] += (
                previous_alpha[:, :characters - step]
                * step_probs[:, step, None]
            )

        prior = prior * text_mask
        prior = prior / prior.sum(dim=-1, keepdim=True).clamp_min(1e-8)

        uniform = text_mask.to(prior.dtype)
        uniform = uniform / uniform.sum(dim=-1, keepdim=True)
        prior = (
            (1 - self.escape_probability) * prior
            + self.escape_probability * uniform
        )

        scores = scores + self.prior_strength * torch.log(prior.clamp_min(1e-8))
        scores = scores.masked_fill(~text_mask, torch.finfo(scores.dtype).min)

        alpha = F.softmax(scores, dim=-1)
        context = torch.bmm(alpha[:, None, :], text_states).squeeze(1)

        return context, alpha

In the decoder, previous_alpha starts with all weight on the first character. Each new alpha is fed into the next step.

previous_alpha = F.one_hot(
    torch.zeros(batch_size, dtype=torch.long, device=text_states.device),
    num_classes=text_states.size(1),
).to(text_states.dtype)

# at each decoder step:
context, alpha = attention(
    query=attention_gru_state,
    text_states=text_states,
    previous_alpha=previous_alpha,
    text_mask=text_mask,
)
previous_alpha = alpha

decoder_input = torch.cat([attention_gru_state, context], dim=-1)