Sequence-to-Sequence and Beam Search

Seq2Seq uses an encoder LSTM to compress input into a context vector $\mathbf{c} = \mathbf{h}_T^{\text{enc}}$, then a decoder LSTM generates output autoregressively:

$$\mathbf{h}_t^{\text{dec}} = \text{LSTM}(\mathbf{y}_{t-1}, \mathbf{h}_{t-1}^{\text{dec}}), \quad \mathbf{y}_t = \text{Linear}(\mathbf{h}_t^{\text{dec}})$$

class Seq2Seq(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.encoder = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.decoder = nn.LSTM(output_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

Beam search maintains $B$ candidate sequences (beams) at each step, expanding each with all possible next tokens, then keeping the top-$B$ by cumulative log-probability:

$$\text{score}(\mathbf{y}_{1:t}) = \sum_{i=1}^{t} \log p(y_i | y_{<i})$$

Beam width $B=1$ is greedy search. Typical $B \in [3, 10]$.