Repository: CUNY-CL/yoyodyne

Issue:
GRU support
It would be nice if our RNN encoders and decoders, which are currently LSTMs, could be replaced with GRUs.

A simple CLI option would be something like this. Rename the architectures `Rnn` and `AttentiveRnn` and modify the architecture options to reflect the change. By default, make the RNN type default to LSTM, but add a flag (`--rnn_type`? just `--rnn`) to select a GRU instead.

Alternatively, but only if necessary, one could project this choice down into the architecture level so that one can choose between an `AttentiveLstm` and an `AttentiveGru` etc. 

I am assigning this to @bonham79 since I know it's an interest of his.

Discussion:

Solution patch:
diff --git a/README.md b/README.md
index 759520ab..63ed3747 100644
--- a/README.md
+++ b/README.md
@@ -26,11 +26,10 @@ Yoyodyne is inspired by [FairSeq](https://github.com/facebookresearch/fairseq)
 -   It is for small-vocabulary sequence-to-sequence generation, and therefore
     includes no affordances for machine translation or language modeling.
     Because of this:
-    -   It has no plugin interface and the architectures provided are intended
-        to be reasonably exhaustive.
+    -   The architectures provided are intended to be reasonably exhaustive.
     -   There is little need for data preprocessing; it works with TSV files.
 -   It has support for using features to condition decoding, with
-    architecture-specific code handling feature information.
+    architecture-specific code for handling feature information.
 -   It supports the use of validation accuracy (not loss) for model selection
     and early stopping.
 -   Releases are made regularly.
@@ -190,54 +189,65 @@ not an attention mechanism is present. This flag also specifies a default
 architecture for the encoder(s), but it is possible to override this with
 additional flags. Supported values for `--arch` are:
 
--   `attentive_lstm`: This is an LSTM decoder with LSTM encoders (by default)
-    and an attention mechanism. The initial hidden state is treated as a learned
+-   `attentive_gru`: This is an GRU decoder with GRU encoders (by default) and
+    an attention mechanism. The initial hidden state is treated as a learned
     parameter.
--   `hard_attention_lstm`: This is an LSTM encoder/decoder modeling generation
-    as a Markov process. By default, it assumes a non-monotonic progression over
+-   `attentive_lstm`: This is similar to the `attentive_gru` but instead uses an
+    LSTM decoder and encoder (by default).
+-   `gru`: This is an GRU decoder with GRU encoders (by default); in lieu of an
+    attention mechanism, the last non-padding hidden state of the encoder is
+    concatenated with the decoder hidden state.
+-   `hard_attention_gru`: This is an GRU encoder/decoder modeling generation as
+    a Markov process. By default, it assumes a non-monotonic progression over
     the source string, but with `--enforce_monotonic` the model must progress
     over each source character in order. A non-zero value of
     `--attention_context` (default: `0`) widens the context window for
     conditioning state transitions to include one or more previous states.
--   `lstm`: This is an LSTM decoder with LSTM encoders (by default); in lieu of
-    an attention mechanism, the last non-padding hidden state of the encoder is
-    concatenated with the decoder hidden state.
--   `pointer_generator_lstm`: This is an LSTM decoder with LSTM encoders (by
+-   `hard_attention_lstm`: This is similar to the `hard_attention_gru` but
+    instead uses an LSTM decoder and encoder (by deafult). `--attention_context`
+    (default: `0`) widens the context window for conditioning state transitions
+    to include one or more previous states.
+-   `lstm`: This is similar to the `gru` but instead uses an LSTM decoder and
+    encoder (by default).
+-   `pointer_generator_gru`: This is an GRU decoder with GRU encoders (by
     default) and a pointer-generator mechanism. Since this model contains a copy
-    mechanism, it may be superior to an ordinary attentive LSTM when the source
+    mechanism, it may be superior to an ordinary attentive GRU when the source
     and target vocabularies overlap significantly. Note that this model requires
     that the number of `--encoder_layers` and `--decoder_layers` match.
--   `pointer_generator_transformer`: This is a transformer decoder with
-    transformer encoders (by default) and a pointer-generator mechanism. Like
-    `pointer_generator_lstm`, it may be superior to an ordinary transformer when
-    the source and target vocabularies overlap significantly. When using
-    features, the user may wish to specify the n umber of features attention
-    heads (with `--features_attention_heads`; default: `1`).
--   `transducer`: This is an LSTM decoder with LSTM encoders (by default) and a
-    neural transducer mechanism. On model creation, expectation maximization is
-    used to learn a sequence of edit operations, and imitation learning is used
-    to train the model to implement the oracle policy, with roll-in controlled
-    by the `--oracle_factor` flag (default: `1`). Since this model assumes
-    monotonic alignment, it may be superior to attentive models when the
+-   `pointer_generator_lstm`: This is similar to the `pointer_generator_gru` but
+    instead uses an LSTM decoder and encoder (by default).
+-   `pointer_generator_transformer`: This is similar to the
+    `pointer_generator_gru` and `pointer_generator_lstm` but instead uses a
+    transformer decoder and encoder (by default). When using features, the user
+    may wish to specify the number of features attention heads (with
+    `--features_attention_heads`).
+-   `transducer_gru`: This is an GRU decoder with GRU encoders (by default) and
+    a neural transducer mechanism. On model creation, expectation maximization
+    is used to learn a sequence of edit operations, and imitation learning is
+    used to train the model to implement the oracle policy, with roll-in
+    controlled by the `--oracle_factor` flag (default: `1`). Since this model
+    assumes monotonic alignment, it may be superior to attentive models when the
     alignment between input and output is roughly monotonic and when input and
     output vocabularies overlap significantly.
+-   `transducer_lstm`: This is similar to the `transducer_gru` but instead uses
+    an LSTM decoder and encoder (by default).
 -   `transformer`: This is a transformer decoder with transformer encoders (by
     default). Sinusodial positional encodings and layer normalization are used.
     The user may wish to specify the number of attention heads (with
     `--source_attention_heads`; default: `4`).
 
-The user can override the default encoder architectures. One can override the
-source encoder using the `--source_encoder_arch` flag:
-
--   `feature_invariant_transformer`: This is a variant of the transformer
-    encoder used with features; it concatenates source and features and uses a
-    learned embedding to distinguish between source and features symbols.
--   `linear`: This is a linear encoder.
--   `lstm`: This is a LSTM encoder.
--   `transformer`: This is a transformer encoder.
+The `--arch` flag specifies the decoder type; the user can override default
+encoder types using the `--source_encoder_arch` flag and, when features are
+present, the `--features_encoder_arch` flag. Valid values are:
 
-When using features, the user can also specify a non-default features encoder
-using the `--features_encoder_arch` flag (`linear`, `lstm`, `transformer`).
+-   `feature_invariant_transformer` (`--source_encoder_arch` only): a variant of
+    the transformer encoder used with features; it concatenates source and
+    features and uses a learned embedding to distinguish between source and
+    features symbols.
+-   `linear`: a linear encoder.
+-   `gru`: a GRU encoder.
+-   `lstm`: a LSTM encoder.
+-   `transformer`: a transformer encoder.
 
 For all models, the user may also wish to specify:
 
@@ -246,8 +256,8 @@ For all models, the user may also wish to specify:
 -   `--encoder_layers` (default: `1`): number of encoder layers
 -   `--hidden_size` (default: `512`): hidden layer size
 
-By default, LSTM encoders are bidirectional. One can disable this with the
-`--no_bidirectional` flag.
+By default, RNN-backed (i.e., GRU and LSTM) encoders are bidirectional. One can
+disable this with the `--no_bidirectional` flag.
 
 ## Training options
 
diff --git a/examples/wandb_sweeps/best_hyperparameters.py b/examples/wandb_sweeps/best_hyperparameters.py
index 9a762fdc..16f987ed 100755
--- a/examples/wandb_sweeps/best_hyperparameters.py
+++ b/examples/wandb_sweeps/best_hyperparameters.py
@@ -54,7 +54,9 @@ def main(args: argparse.Namespace) -> None:
     logging.basicConfig(format="%(levelname)s: %(message)s", level="INFO")
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
-        "--entity", required=True, help="The entity scope for the project."
+        "--entity",
+        required=True,
+        help="The entity scope for the project.",
     )
     parser.add_argument(
         "--project", required=True, help="The project of the sweep."
diff --git a/yoyodyne/__init__.py b/yoyodyne/__init__.py
index c83cbb34..ed5e6fd2 100644
--- a/yoyodyne/__init__.py
+++ b/yoyodyne/__init__.py
@@ -4,9 +4,11 @@
 
 # Silences irrelevant warnings; these are more like "Did you know?"s.
 warnings.filterwarnings(
-    "ignore", ".*does not have many workers which may be a bottleneck.*"
+    "ignore",
+    ".*does not have many workers which may be a bottleneck.*",
 )
 warnings.filterwarnings(
-    "ignore", ".*option adds dropout after all but last recurrent layer.*"
+    "ignore",
+    ".*option adds dropout after all but last recurrent layer.*",
 )
 warnings.filterwarnings("ignore", ".*is a wandb run already in progress.*")
diff --git a/yoyodyne/evaluators.py b/yoyodyne/evaluators.py
index 1f14a287..cc8059cc 100644
--- a/yoyodyne/evaluators.py
+++ b/yoyodyne/evaluators.py
@@ -126,7 +126,7 @@ def finalize_predictions(
     ) -> torch.Tensor:
         """Finalizes predictions.
 
-        Cuts off tensors at the first end_idx, and replaces the rest of the
+        Cuts off tensors at the first END_IDX, and replaces the rest of the
         predictions with PAD_IDX, as these are erroneously decoded while the
         rest of the batch is finishing decoding.
 
@@ -198,8 +198,8 @@ def _finalize_tensor(
         """Finalizes each tensor.
 
         Truncates at EOS for each prediction and returns a List of predictions.
-        This does basically the same as util.pad_after_eos, but does
-        not actually pad since we do not need to return a well-formed tensor.
+        This does basically the same as util.pad_after_eos, but does not
+        actually pad since we do not need to return a well-formed tensor.
 
         Args:
             tensor (torch.Tensor).
diff --git a/yoyodyne/models/__init__.py b/yoyodyne/models/__init__.py
index a1f9b544..09b64b8f 100644
--- a/yoyodyne/models/__init__.py
+++ b/yoyodyne/models/__init__.py
@@ -3,28 +3,40 @@
 import argparse
 
 from .. import defaults
-from .base import BaseEncoderDecoder
-from .hard_attention import HardAttentionLSTM
-from .lstm import AttentiveLSTMEncoderDecoder, LSTMEncoderDecoder
-from .pointer_generator import (
-    PointerGeneratorLSTMEncoderDecoder,
-    PointerGeneratorTransformerEncoderDecoder,
-)
-from .transducer import TransducerEncoderDecoder
-from .transformer import TransformerEncoderDecoder
+from .base import BaseModel
+
+from .hard_attention import HardAttentionGRUModel
+from .hard_attention import HardAttentionLSTMModel
+from .hard_attention import HardAttentionRNNModel  # noqa: F401
+from .pointer_generator import PointerGeneratorGRUModel
+from .pointer_generator import PointerGeneratorLSTMModel
+from .pointer_generator import PointerGeneratorTransformerModel
+from .rnn import AttentiveGRUModel
+from .rnn import AttentiveLSTMModel
+from .rnn import GRUModel
+from .rnn import LSTMModel
+from .rnn import RNNModel  # noqa: F401
+
+from .transducer import TransducerGRUModel, TransducerLSTMModel  # noqa: F401
+from .transformer import TransformerModel
 
 _model_fac = {
-    "attentive_lstm": AttentiveLSTMEncoderDecoder,
-    "hard_attention_lstm": HardAttentionLSTM,
-    "lstm": LSTMEncoderDecoder,
-    "pointer_generator_lstm": PointerGeneratorLSTMEncoderDecoder,
-    "pointer_generator_transformer": PointerGeneratorTransformerEncoderDecoder,  # noqa: 501
-    "transducer": TransducerEncoderDecoder,
-    "transformer": TransformerEncoderDecoder,
+    "attentive_gru": AttentiveGRUModel,
+    "attentive_lstm": AttentiveLSTMModel,
+    "hard_attention_gru": HardAttentionGRUModel,
+    "hard_attention_lstm": HardAttentionLSTMModel,
+    "pointer_generator_gru": PointerGeneratorGRUModel,
+    "pointer_generator_lstm": PointerGeneratorLSTMModel,
+    "pointer_generator_transformer": PointerGeneratorTransformerModel,  # noqa: E501
+    "gru": GRUModel,
+    "lstm": LSTMModel,
+    "transducer_gru": TransducerGRUModel,
+    "transducer_lstm": TransducerLSTMModel,
+    "transformer": TransformerModel,
 }
 
 
-def get_model_cls(arch: str) -> BaseEncoderDecoder:
+def get_model_cls(arch: str) -> BaseModel:
     """Model factory.
 
     Args:
@@ -35,7 +47,7 @@ def get_model_cls(arch: str) -> BaseEncoderDecoder:
         NotImplementedError: Architecture not found.
 
     Returns:
-        BaseEncoderDecoder.
+        BaseModel.
     """
     try:
         return _model_fac[arch]
@@ -45,14 +57,14 @@ def get_model_cls(arch: str) -> BaseEncoderDecoder:
 
 def get_model_cls_from_argparse_args(
     args: argparse.Namespace,
-) -> BaseEncoderDecoder:
+) -> BaseModel:
     """Creates a model from CLI arguments.
 
     Args:
         args (argparse.Namespace).
 
     Returns:
-        BaseEncoderDecoder.
+        BaseModel.
     """
     return get_model_cls(args.arch)
 
@@ -69,7 +81,7 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--arch",
         choices=_model_fac.keys(),
-        default="attentive_lstm",
+        default=defaults.ARCH,
         help="Model architecture. Default: %(default)s.",
     )
     parser.add_argument(
diff --git a/yoyodyne/models/base.py b/yoyodyne/models/base.py
index 72354622..13cc4a80 100644
--- a/yoyodyne/models/base.py
+++ b/yoyodyne/models/base.py
@@ -19,7 +19,7 @@
 from . import modules
 
 
-class BaseEncoderDecoder(lightning.LightningModule):
+class BaseModel(lightning.LightningModule):
     """Base class, handling Lightning integration."""
 
     #  TODO: clean up type checking here.
@@ -56,29 +56,30 @@ class BaseEncoderDecoder(lightning.LightningModule):
     def __init__(
         self,
         *,
-        vocab_size,
+        beta1=defaults.BETA1,
+        beta2=defaults.BETA2,
         features_vocab_size,
-        target_vocab_size,
         source_encoder_cls,
+        target_vocab_size,
+        vocab_size,
+        # All of these have keyword defaults.
+        beam_width=defaults.BEAM_WIDTH,
+        decoder_layers=defaults.DECODER_LAYERS,
+        dropout=defaults.DROPOUT,
+        embedding_size=defaults.EMBEDDING_SIZE,
+        encoder_layers=defaults.ENCODER_LAYERS,
         eval_metrics=defaults.EVAL_METRICS,
         features_encoder_cls=None,
-        beta1=defaults.BETA1,
-        beta2=defaults.BETA2,
+        hidden_size=defaults.HIDDEN_SIZE,
+        label_smoothing=defaults.LABEL_SMOOTHING,
         learning_rate=defaults.LEARNING_RATE,
+        max_source_length=defaults.MAX_SOURCE_LENGTH,
+        max_target_length=defaults.MAX_TARGET_LENGTH,
         optimizer=defaults.OPTIMIZER,
         scheduler=None,
         scheduler_kwargs=None,
-        dropout=defaults.DROPOUT,
-        label_smoothing=defaults.LABEL_SMOOTHING,
         teacher_forcing=defaults.TEACHER_FORCING,
-        beam_width=defaults.BEAM_WIDTH,
-        max_source_length=defaults.MAX_SOURCE_LENGTH,
-        max_target_length=defaults.MAX_TARGET_LENGTH,
-        encoder_layers=defaults.ENCODER_LAYERS,
-        decoder_layers=defaults.DECODER_LAYERS,
-        embedding_size=defaults.EMBEDDING_SIZE,
-        hidden_size=defaults.HIDDEN_SIZE,
-        **kwargs,  # Ignored.
+        **kwargs,
     ):
         super().__init__()
         # Symbol processing.
@@ -118,25 +119,25 @@ def __init__(
         )
         # Instantiates encoders class.
         self.source_encoder = source_encoder_cls(
-            embeddings=self.embeddings,
-            embedding_size=self.embedding_size,
-            num_embeddings=self.vocab_size,
             dropout=self.dropout,
-            layers=self.encoder_layers,
-            hidden_size=self.hidden_size,
+            embedding_size=self.embedding_size,
+            embeddings=self.embeddings,
             features_vocab_size=features_vocab_size,
+            hidden_size=self.hidden_size,
+            layers=self.encoder_layers,
             max_source_length=max_source_length,
+            num_embeddings=self.vocab_size,
             **kwargs,
         )
         self.features_encoder = (
             features_encoder_cls(
-                embeddings=self.embeddings,
-                embedding_size=self.embedding_size,
-                num_embeddings=self.vocab_size,
                 dropout=self.dropout,
-                layers=self.encoder_layers,
+                embedding_size=self.embedding_size,
+                embeddings=self.embeddings,
                 hidden_size=self.hidden_size,
+                layers=self.encoder_layers,
                 max_source_length=max_source_length,
+                num_embeddings=self.vocab_size,
                 **kwargs,
             )
             if features_encoder_cls is not None
@@ -145,7 +146,7 @@ def __init__(
         self.decoder = self.get_decoder()
         # Saves hyperparameters for PL checkpointing.
         self.save_hyperparameters(
-            ignore=["source_encoder", "decoder", "features_encoder"]
+            ignore=["source_encoder", "decoder", "features_encoder"],
         )
         # Logs the module names.
         util.log_info(f"Model: {self.name}")
@@ -197,7 +198,7 @@ def beam_decode(
                 log-likelihood of each prediction.
         """
         raise NotImplementedError(
-            f"Beam search not implemented for {self.name} model."
+            f"Beam search not implemented for {self.name} model"
         )
 
     @property
diff --git a/yoyodyne/models/expert.py b/yoyodyne/models/expert.py
index 13187093..00662fd6 100644
--- a/yoyodyne/models/expert.py
+++ b/yoyodyne/models/expert.py
@@ -10,7 +10,16 @@
 import abc
 import argparse
 import dataclasses
-from typing import Any, Dict, Iterable, Iterator, List, Sequence, Set, Tuple
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Sequence,
+    Set,
+    Tuple,
+)
 
 import numpy
 from maxwell import actions, sed
@@ -125,7 +134,9 @@ def to_i2w(self) -> List[str]:
         return self.i2w[len(self.start_vocab_idx) :]  # noqa: E203
 
     @property
-    def substitutions(self) -> List[Tuple[int, actions.ConditionalEdit]]:
+    def substitutions(
+        self,
+    ) -> List[Tuple[int, actions.ConditionalEdit]]:
         return [
             i
             for i, a in enumerate(self.i2w)
diff --git a/yoyodyne/models/hard_attention.py b/yoyodyne/models/hard_attention.py
index 1f4641fd..708a237a 100644
--- a/yoyodyne/models/hard_attention.py
+++ b/yoyodyne/models/hard_attention.py
@@ -7,11 +7,11 @@
 from torch import nn
 
 from .. import data, defaults, special
-from . import lstm, modules
+from . import modules, rnn
 
 
-class HardAttentionLSTM(lstm.LSTMEncoderDecoder):
-    """Hard attention transducer.
+class HardAttentionRNNModel(rnn.RNNModel):
+    """Base class for hard attention models.
 
     Learns probability distribution of target string by modeling transduction
     of source string to target string as Markov process. Assumes each character
@@ -123,17 +123,13 @@ def forward(
                 [encoder_out, encoder_features_out], dim=-1
             )
         if self.training:
-            output = self.decode(
+            return self.decode(
                 encoder_out,
                 batch.source.mask,
                 batch.target.padded,
             )
         else:
-            output = self.greedy_decode(
-                encoder_out,
-                batch.source.mask,
-            )
-        return output
+            return self.greedy_decode(encoder_out, batch.source.mask)
 
     def decode(
         self,
@@ -179,51 +175,6 @@ def decode(
             all_transition_probs.append(transition_probs)
         return torch.stack(all_log_probs), torch.stack(all_transition_probs)
 
-    def decode_step(
-        self,
-        tgt_symbol: torch.Tensor,
-        decoder_hiddens: Tuple[torch.Tensor, torch.Tensor],
-        encoder_out: torch.Tensor,
-        encoder_mask: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        """Performs a single decoding step for current state of decoder.
-
-        Args:
-            tgt_symbol (torch.Tensor): tgt symbol for current state.
-            decoder_hiddens (Tuple[torch.Tensor, torch.Tensor]): the last
-                hidden states from the decoder of shapes 1 x B x decoder_dim
-                and 1 x B x decoder_dim.
-            encoder_out (torch.Tensor): batch of encoded input symbols
-                of shape batch_size x src_len x (encoder_hidden *
-                num_directions).
-            encoder_mask (torch.Tensor): mask for the batch of encoded
-                input symbols of shape batch_size x src_len.
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor,
-                torch.Tensor]: emission probabilities for each transition
-                state (target symbol), transition probabilities for each
-                transition state (target symbol), and the last hidden states
-                from the decoder, of shape (1 x B x decoder_dim, 1 x B x
-                decoder_dim).
-        """
-        decoded = self.decoder(
-            tgt_symbol,
-            decoder_hiddens,
-            encoder_out,
-            encoder_mask,
-        )
-        output, transition_prob, decoder_hiddens = (
-            decoded.output,
-            decoded.embeddings,
-            decoded.hiddens,
-        )
-        logits = self.classifier(output)
-        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
-        # Expands matrix for all time steps.
-        if self.enforce_monotonic:
-            transition_prob = self._apply_mono_mask(transition_prob)
-        return log_probs, transition_prob, decoder_hiddens
-
     def greedy_decode(
         self,
         encoder_out: torch.Tensor,
@@ -246,7 +197,7 @@ def greedy_decode(
                 x pred_seq_len, final likelihoods per prediction step of shape
                 batch_size x 1 x src_len.
         """
-        batch_size, _ = encoder_mask.shape
+        batch_size = encoder_mask.shape[0]
         log_probs, transition_prob, decoder_hiddens = self.init_decoding(
             encoder_out, encoder_mask
         )
@@ -331,7 +282,9 @@ def _gather_at_idx(prob: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
         return output * pad_mask
 
     @staticmethod
-    def _apply_mono_mask(transition_prob: torch.Tensor) -> torch.Tensor:
+    def _apply_mono_mask(
+        transition_prob: torch.Tensor,
+    ) -> torch.Tensor:
         """Applies monotonic attention mask to transition probabilities.
 
         Enforces a 0 log-probability values for all non-monotonic relations
@@ -387,10 +340,10 @@ def validation_step(self, batch: data.PaddedBatch, batch_idx: int) -> Dict:
         # Processes for accuracy calculation.
         val_eval_items_dict = {}
         for evaluator in self.evaluators:
-            predictions = evaluator.finalize_predictions(predictions)
-            golds = evaluator.finalize_golds(batch.target.padded)
+            final_predictions = evaluator.finalize_predictions(predictions)
+            final_golds = evaluator.finalize_golds(batch.target.padded)
             val_eval_items_dict[evaluator.name] = evaluator.get_eval_item(
-                predictions, golds
+                final_predictions, final_golds
             )
         val_eval_items_dict.update({"val_loss": -likelihood.mean()})
         return val_eval_items_dict
@@ -399,41 +352,6 @@ def predict_step(self, batch: data.PaddedBatch, batch_idx: int) -> Dict:
         predictions, _ = self(batch)
         return predictions
 
-    def get_decoder(self) -> modules.lstm.HardAttentionLSTMDecoder:
-        if self.attention_context > 0:
-            return modules.lstm.ContextHardAttentionLSTMDecoder(
-                attention_context=self.attention_context,
-                bidirectional=False,
-                decoder_input_size=(
-                    self.source_encoder.output_size
-                    + self.features_encoder.output_size
-                    if self.has_features_encoder
-                    else self.source_encoder.output_size
-                ),
-                dropout=self.dropout,
-                embeddings=self.embeddings,
-                embedding_size=self.embedding_size,
-                hidden_size=self.hidden_size,
-                layers=self.decoder_layers,
-                num_embeddings=self.target_vocab_size,
-            )
-        else:
-            return modules.lstm.HardAttentionLSTMDecoder(
-                decoder_input_size=(
-                    self.source_encoder.output_size
-                    + self.features_encoder.output_size
-                    if self.has_features_encoder
-                    else self.source_encoder.output_size
-                ),
-                bidirectional=False,
-                embeddings=self.embeddings,
-                embedding_size=self.embedding_size,
-                dropout=self.dropout,
-                layers=self.decoder_layers,
-                hidden_size=self.hidden_size,
-                num_embeddings=self.target_vocab_size,
-            )
-
     def _get_loss_func(
         self,
     ) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]:
@@ -466,6 +384,13 @@ def _loss(
         loss = -torch.logsumexp(fwd, dim=-1).mean() / target.shape[1]
         return loss
 
+    def get_decoder(self):
+        raise NotImplementedError
+
+    @property
+    def name(self) -> str:
+        raise NotImplementedError
+
     @staticmethod
     def add_argparse_args(parser: argparse.ArgumentParser) -> None:
         """Adds HMM configuration options to the argument parser.
@@ -492,3 +417,177 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
             help="Width of attention context "
             "(hard attention architectures only). Default: %(default)s.",
         )
+
+
+class HardAttentionGRUModel(HardAttentionRNNModel, rnn.GRUModel):
+    """Hard attention with GRU backend."""
+
+    def get_decoder(self):
+        if self.attention_context > 0:
+            return modules.ContextHardAttentionGRUDecoder(
+                attention_context=self.attention_context,
+                bidirectional=False,
+                decoder_input_size=(
+                    self.source_encoder.output_size
+                    + self.features_encoder.output_size
+                    if self.has_features_encoder
+                    else self.source_encoder.output_size
+                ),
+                dropout=self.dropout,
+                embeddings=self.embeddings,
+                embedding_size=self.embedding_size,
+                hidden_size=self.hidden_size,
+                layers=self.decoder_layers,
+                num_embeddings=self.target_vocab_size,
+            )
+        else:
+            return modules.HardAttentionGRUDecoder(
+                bidirectional=False,
+                decoder_input_size=(
+                    self.source_encoder.output_size
+                    + self.features_encoder.output_size
+                    if self.has_features_encoder
+                    else self.source_encoder.output_size
+                ),
+                dropout=self.dropout,
+                embedding_size=self.embedding_size,
+                embeddings=self.embeddings,
+                hidden_size=self.hidden_size,
+                layers=self.decoder_layers,
+                num_embeddings=self.target_vocab_size,
+            )
+
+    def decode_step(
+        self,
+        tgt_symbol: torch.Tensor,
+        decoder_hiddens: torch.Tensor,
+        encoder_out: torch.Tensor,
+        encoder_mask: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Performs a single decoding step for current state of decoder.
+
+        Args:
+            tgt_symbol (torch.Tensor): tgt symbol for current state.
+            decoder_hiddens (torch.Tensor): the last
+                hidden states from the decoder of shape 1 x B x decoder_dim.
+            encoder_out (torch.Tensor): batch of encoded input symbols
+                of shape batch_size x src_len x (encoder_hidden *
+                num_directions).
+            encoder_mask (torch.Tensor): mask for the batch of encoded
+                input symbols of shape batch_size x src_len.
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: emission
+                probabilities for each transition state (target symbol),
+                transition probabilities for each transition state (target
+                symbol), and the last hidden states from the decoder of
+                shape (1 x B x decoder_dim, 1 x B x decoder_dim).
+        """
+        decoded = self.decoder(
+            tgt_symbol,
+            decoder_hiddens,
+            encoder_out,
+            encoder_mask,
+        )
+        output, transition_prob, decoder_hiddens = (
+            decoded.output,
+            decoded.embeddings,
+            decoded.hiddens,
+        )
+        logits = self.classifier(output)
+        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+        # Expands matrix for all time steps.
+        if self.enforce_monotonic:
+            transition_prob = self._apply_mono_mask(transition_prob)
+        return log_probs, transition_prob, decoder_hiddens
+
+    @property
+    def name(self) -> str:
+        return "hard attention GRU"
+
+
+class HardAttentionLSTMModel(HardAttentionRNNModel, rnn.LSTMModel):
+    """Hard attention with LSTM backend."""
+
+    def get_decoder(self):
+        if self.attention_context > 0:
+            return modules.ContextHardAttentionLSTMDecoder(
+                attention_context=self.attention_context,
+                bidirectional=False,
+                decoder_input_size=(
+                    self.source_encoder.output_size
+                    + self.features_encoder.output_size
+                    if self.has_features_encoder
+                    else self.source_encoder.output_size
+                ),
+                dropout=self.dropout,
+                hidden_size=self.hidden_size,
+                embeddings=self.embeddings,
+                embedding_size=self.embedding_size,
+                layers=self.decoder_layers,
+                num_embeddings=self.target_vocab_size,
+            )
+        else:
+            return modules.HardAttentionLSTMDecoder(
+                bidirectional=False,
+                decoder_input_size=(
+                    self.source_encoder.output_size
+                    + self.features_encoder.output_size
+                    if self.has_features_encoder
+                    else self.source_encoder.output_size
+                ),
+                dropout=self.dropout,
+                embeddings=self.embeddings,
+                embedding_size=self.embedding_size,
+                hidden_size=self.hidden_size,
+                layers=self.decoder_layers,
+                num_embeddings=self.target_vocab_size,
+            )
+
+    def decode_step(
+        self,
+        tgt_symbol: torch.Tensor,
+        decoder_hiddens: Tuple[torch.Tensor, torch.Tensor],
+        encoder_out: torch.Tensor,
+        encoder_mask: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """Performs a single decoding step for current state of decoder.
+
+        Args:
+            tgt_symbol (torch.Tensor): tgt symbol for current state.
+            decoder_hiddens (Tuple[torch.Tensor, torch.Tensor]): the last
+                hidden states from the decoder of shapes 1 x B x decoder_dim
+                and 1 x B x decoder_dim.
+            encoder_out (torch.Tensor): batch of encoded input symbols
+                of shape batch_size x src_len x (encoder_hidden *
+                num_directions).
+            encoder_mask (torch.Tensor): mask for the batch of encoded
+                input symbols of shape batch_size x src_len.
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor,
+                torch.Tensor]: emission probabilities for each transition
+                state (target symbol), transition probabilities for each
+                transition state (target symbol), and the last hidden states
+                from the decoder, of shape (1 x B x decoder_dim, 1 x B x
+                decoder_dim).
+        """
+        decoded = self.decoder(
+            tgt_symbol,
+            decoder_hiddens,
+            encoder_out,
+            encoder_mask,
+        )
+        output, transition_prob, decoder_hiddens = (
+            decoded.output,
+            decoded.embeddings,
+            decoded.hiddens,
+        )
+        logits = self.classifier(output)
+        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+        # Expands matrix for all time steps.
+        if self.enforce_monotonic:
+            transition_prob = self._apply_mono_mask(transition_prob)
+        return log_probs, transition_prob, decoder_hiddens
+
+    @property
+    def name(self) -> str:
+        return "hard attention LSTM"
diff --git a/yoyodyne/models/modules/__init__.py b/yoyodyne/models/modules/__init__.py
index b54624e2..ccb8c1bb 100644
--- a/yoyodyne/models/modules/__init__.py
+++ b/yoyodyne/models/modules/__init__.py
@@ -4,9 +4,21 @@
 
 from .base import BaseModule
 from .linear import LinearEncoder
-from .lstm import LSTMAttentiveDecoder, LSTMDecoder, LSTMEncoder  # noqa: F401
-from .transformer import TransformerDecoder  # noqa F401
-from .transformer import FeatureInvariantTransformerEncoder, TransformerEncoder
+from .rnn import AttentiveGRUDecoder  # noqa: F401
+from .rnn import AttentiveLSTMDecoder  # noqa: F401
+from .rnn import ContextHardAttentionGRUDecoder  # noqa: F401
+from .rnn import ContextHardAttentionLSTMDecoder  # noqa: F401
+from .rnn import GRUDecoder  # noqa: F401
+from .rnn import GRUEncoder
+from .rnn import HardAttentionGRUDecoder  # noqa: F401
+from .rnn import HardAttentionLSTMDecoder  # noqa: F401
+from .rnn import LSTMDecoder  # noqa: F401
+from .rnn import LSTMEncoder
+from .transformer import TransformerDecoder  # noqa: F401
+from .transformer import (
+    FeatureInvariantTransformerEncoder,
+    TransformerEncoder,
+)
 
 
 class Error(Exception):
@@ -19,18 +31,24 @@ class EncoderMismatchError(Error):
 
 _encoder_fac = {
     "feature_invariant_transformer": FeatureInvariantTransformerEncoder,
+    "gru": GRUEncoder,
     "linear": LinearEncoder,
     "lstm": LSTMEncoder,
     "transformer": TransformerEncoder,
 }
 _model_to_encoder_fac = {
+    "attentive_gru": GRUEncoder,
     "attentive_lstm": LSTMEncoder,
+    "gru": GRUEncoder,
+    "hard_attention_gru": GRUEncoder,
+    "hard_attention_lstm": LSTMEncoder,
     "lstm": LSTMEncoder,
+    "pointer_generator_gru": GRUEncoder,
     "pointer_generator_lstm": LSTMEncoder,
     "pointer_generator_transformer": TransformerEncoder,
-    "transducer": LSTMEncoder,
+    "transducer_gru": GRUEncoder,
+    "transducer_lstm": LSTMEncoder,
     "transformer": TransformerEncoder,
-    "hard_attention_lstm": LSTMEncoder,
 }
 
 
@@ -86,7 +104,7 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
     )
     parser.add_argument(
         "--features_encoder_arch",
-        choices=["linear", "lstm", "transformer"],
+        choices=["gru", "linear", "lstm", "transformer"],
         help="Model architecture to use for the features encoder.",
     )
 
diff --git a/yoyodyne/models/modules/attention.py b/yoyodyne/models/modules/attention.py
index 349d4bd3..2e49a44d 100644
--- a/yoyodyne/models/modules/attention.py
+++ b/yoyodyne/models/modules/attention.py
@@ -90,6 +90,6 @@ def score(
         # -> B x seq_len x hidden_size
         m = self.M(concat)
         # -> B x seq_len x 1.
-        scores = self.V(torch.tanh((m)))
+        scores = self.V(torch.tanh(m))
         # -> B x seq_len.
         return scores.squeeze(2)
diff --git a/yoyodyne/models/modules/base.py b/yoyodyne/models/modules/base.py
index 0eb91f2f..8dee8c86 100644
--- a/yoyodyne/models/modules/base.py
+++ b/yoyodyne/models/modules/base.py
@@ -1,8 +1,7 @@
-"""Base module class, with PL integration.
-"""
+"""Base module class with PL integration."""
 
 import dataclasses
-from typing import Optional, Tuple
+from typing import Optional, Union, Tuple
 
 import lightning
 import torch
@@ -13,10 +12,12 @@
 
 @dataclasses.dataclass
 class ModuleOutput:
-    """For tracking outputs of forward passes over varying architectures."""
+    """Output for forward passes."""
 
     output: torch.Tensor
-    hiddens: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    hiddens: Optional[
+        Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
+    ] = None
     embeddings: Optional[torch.Tensor] = None
 
     @property
@@ -75,4 +76,5 @@ def embed(self, symbols: torch.Tensor) -> torch.Tensor:
         return self.dropout_layer(embedded)
 
     @property
-    def output_size(self) -> int: ...
+    def output_size(self) -> int:
+        raise NotImplementedError
diff --git a/yoyodyne/models/modules/linear.py b/yoyodyne/models/modules/linear.py
index f81fcb97..0ab6179c 100644
--- a/yoyodyne/models/modules/linear.py
+++ b/yoyodyne/models/modules/linear.py
@@ -1,23 +1,11 @@
-"""Linear model classes."""
-
-from typing import Tuple
-
-import torch
+"""Linear module class."""
 
 from ... import data
 from . import base
 
 
-class LinearModule(base.BaseModule):
-    """Simple linear embedding module."""
-
-    pass
-
-
-class LinearEncoder(LinearModule):
-    def forward(
-        self, source: data.PaddedTensor
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+class LinearEncoder(base.BaseModule):
+    def forward(self, source: data.PaddedTensor) -> base.ModuleOutput:
         """Encodes the input.
 
         Args:
@@ -25,8 +13,7 @@ def forward(
                 for source, of shape B x seq_len x 1.
 
         Returns:
-            Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-                encoded timesteps, and the LSTM h0 and c0 cells.
+            base.ModuleOutput.
         """
         return base.ModuleOutput(self.embed(source.padded))
 
diff --git a/yoyodyne/models/modules/lstm.py b/yoyodyne/models/modules/rnn.py
similarity index 68%
rename from yoyodyne/models/modules/lstm.py
rename to yoyodyne/models/modules/rnn.py
index c9e6b805..be2e983b 100644
--- a/yoyodyne/models/modules/lstm.py
+++ b/yoyodyne/models/modules/rnn.py
@@ -1,4 +1,4 @@
-"""LSTM model classes."""
+"""RNN module classes."""
 
 from typing import Tuple
 
@@ -9,19 +9,16 @@
 from . import attention, base
 
 
-class LSTMModule(base.BaseModule):
-    """Base encoder for LSTM.
+class RNNModule(base.BaseModule):
+    """Base class for RNN modules.
 
     Args:
-        *args: passed to superclass.
         bidirectional (bool).
+        *args: passed to superclass.
         **kwargs: passed to superclass.
     """
 
-    # Model arguments.
     bidirectional: bool
-    # Constructed inside __init__.
-    module: nn.LSTM
 
     def __init__(
         self,
@@ -37,9 +34,16 @@ def __init__(
     def num_directions(self) -> int:
         return 2 if self.bidirectional else 1
 
+    def get_module(self):
+        raise NotImplementedError
 
-class LSTMEncoder(LSTMModule):
-    """LSTM encoder."""
+    @property
+    def name(self) -> str:
+        raise NotImplementedError
+
+
+class RNNEncoder(RNNModule):
+    """Base class for RNN encoders."""
 
     def forward(self, source: data.PaddedTensor) -> base.ModuleOutput:
         """Encodes the input.
@@ -49,46 +53,71 @@ def forward(self, source: data.PaddedTensor) -> base.ModuleOutput:
                 for source, of shape B x seq_len x 1.
 
         Returns:
-            base.ModuleOutput: encoded timesteps, and LSTM h0 and c0 cells.
+            base.ModuleOutput.
         """
         embedded = self.embed(source.padded)
         # Packs embedded source symbols into a PackedSequence.
         packed = nn.utils.rnn.pack_padded_sequence(
-            embedded, source.lengths(), batch_first=True, enforce_sorted=False
+            embedded,
+            source.lengths(),
+            batch_first=True,
+            enforce_sorted=False,
         )
-        # -> B x seq_len x encoder_dim, (h0, c0).
-        packed_outs, (H, C) = self.module(packed)
+        # -> B x seq_len x encoder_dim, hiddens
+        packed_outs, hiddens = self.module(packed)
         encoded, _ = nn.utils.rnn.pad_packed_sequence(
             packed_outs,
             batch_first=True,
             padding_value=special.PAD_IDX,
             total_length=None,
         )
-        return base.ModuleOutput(encoded, hiddens=(H, C))
+        return base.ModuleOutput(encoded, hiddens)
 
-    def get_module(self) -> nn.LSTM:
-        return nn.LSTM(
+    @property
+    def output_size(self) -> int:
+        return self.hidden_size * self.num_directions
+
+
+class GRUEncoder(RNNEncoder):
+    """GRU encoder."""
+
+    def get_module(self) -> nn.GRU:
+        return nn.GRU(
             self.embedding_size,
             self.hidden_size,
-            num_layers=self.layers,
-            dropout=self.dropout,
             batch_first=True,
             bidirectional=self.bidirectional,
+            dropout=self.dropout,
+            num_layers=self.layers,
         )
 
     @property
-    def output_size(self) -> int:
-        return self.hidden_size * self.num_directions
+    def name(self) -> str:
+        return "GRU"
+
+
+class LSTMEncoder(RNNEncoder):
+    """LSTM encoder."""
+
+    def get_module(self) -> nn.LSTM:
+        return nn.LSTM(
+            self.embedding_size,
+            self.hidden_size,
+            batch_first=True,
+            bidirectional=self.bidirectional,
+            dropout=self.dropout,
+            num_layers=self.layers,
+        )
 
     @property
     def name(self) -> str:
         return "LSTM"
 
 
-class LSTMDecoder(LSTMModule):
-    """LSTM decoder."""
+class RNNDecoder(RNNModule):
+    """Base class for RNN decoders."""
 
-    def __init__(self, *args, decoder_input_size, **kwargs):
+    def __init__(self, decoder_input_size, *args, **kwargs):
         self.decoder_input_size = decoder_input_size
         super().__init__(*args, **kwargs)
 
@@ -112,8 +141,8 @@ def forward(
                 shape B x seq_len.
 
         Returns:
-            base.ModuleOutput:  decoder output, and the previous hidden states
-                from the decoder LSTM.
+            base.ModuleOutput: decoder output, and the previous hidden states
+                from the decoder RNN.
         """
         embedded = self.embed(symbol)
         # -> 1 x B x decoder_dim.
@@ -130,16 +159,24 @@ def forward(
         )
         # -> B x 1 x encoder_dim.
         last_encoder_out = torch.gather(encoder_out, 1, last_encoder_out_idxs)
-        # The input to decoder LSTM is the embedding concatenated to the
+        # The input to decoder RNN is the embedding concatenated to the
         # weighted, encoded, inputs.
         output, hiddens = self.module(
             torch.cat((embedded, last_encoder_out), 2), last_hiddens
         )
         output = self.dropout_layer(output)
-        return base.ModuleOutput(output, hiddens=hiddens)
+        return base.ModuleOutput(output, hiddens)
 
-    def get_module(self) -> nn.LSTM:
-        return nn.LSTM(
+    @property
+    def output_size(self) -> int:
+        return self.hidden_size
+
+
+class GRUDecoder(RNNDecoder):
+    """GRU decoder."""
+
+    def get_module(self) -> nn.GRU:
+        return nn.GRU(
             self.decoder_input_size + self.embedding_size,
             self.hidden_size,
             num_layers=self.layers,
@@ -149,26 +186,81 @@ def get_module(self) -> nn.LSTM:
         )
 
     @property
-    def output_size(self) -> int:
-        return self.hidden_size
+    def name(self) -> str:
+        return "GRU"
+
+
+class LSTMDecoder(RNNDecoder):
+    """LSTM decoder."""
+
+    def get_module(self) -> nn.LSTM:
+        return nn.LSTM(
+            self.decoder_input_size + self.embedding_size,
+            self.hidden_size,
+            num_layers=self.layers,
+            dropout=self.dropout,
+            batch_first=True,
+            bidirectional=self.bidirectional,
+        )
 
     @property
     def name(self) -> str:
         return "LSTM"
 
 
-class LSTMAttentiveDecoder(LSTMDecoder):
-    """Attentive LSTM decoder."""
-
-    attention_input_size: int
+class AttentiveRNNDecoder(RNNDecoder):
+    """Base class for attentive RNN decoders."""
 
-    def __init__(self, *args, attention_input_size, **kwargs):
+    def __init__(self, attention_input_size, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.attention_input_size = attention_input_size
         self.attention = attention.Attention(
-            self.attention_input_size, self.hidden_size
+            attention_input_size, self.hidden_size
         )
 
+
+class AttentiveGRUDecoder(AttentiveRNNDecoder, GRUDecoder):
+    """Attentive GRU decoder."""
+
+    def forward(
+        self,
+        symbol: torch.Tensor,
+        last_hiddens: torch.Tensor,
+        encoder_out: torch.Tensor,
+        encoder_mask: torch.Tensor,
+    ) -> base.ModuleOutput:
+        """Single decode pass.
+
+        Args:
+            symbol (torch.Tensor): previously decoded symbol of shape B x 1.
+            last_hiddens (torch.Tensor): last hidden states from the decoder
+                of shape 1 x B x decoder_dim.
+            encoder_out (torch.Tensor): encoded input sequence of shape
+                B x seq_len x encoder_dim.
+            encoder_mask (torch.Tensor): mask for the encoded input batch of
+                shape B x seq_len.
+
+        Returns:
+            base.ModuleOutput: decoder output, and the previous hidden states
+                from the decoder RNN.
+        """
+        embedded = self.embed(symbol)
+        context, _ = self.attention(
+            last_hiddens.transpose(0, 1), encoder_out, encoder_mask
+        )
+        output, hiddens = self.module(
+            torch.cat((embedded, context), 2), last_hiddens
+        )
+        output = self.dropout_layer(output)
+        return base.ModuleOutput(output, hiddens)
+
+    @property
+    def name(self) -> str:
+        return "attentive GRU"
+
+
+class AttentiveLSTMDecoder(AttentiveRNNDecoder, LSTMDecoder):
+    """Attentive LSTM decoder."""
+
     def forward(
         self,
         symbol: torch.Tensor,
@@ -181,36 +273,36 @@ def forward(
         Args:
             symbol (torch.Tensor): previously decoded symbol of shape B x 1.
             last_hiddens (Tuple[torch.Tensor, torch.Tensor]): last hidden
-                states from the decoder of shape
-                (1 x B x decoder_dim, 1 x B x decoder_dim).
+                and cell state from the decoder of shape 1 x B x decoder_dim.
             encoder_out (torch.Tensor): encoded input sequence of shape
                 B x seq_len x encoder_dim.
             encoder_mask (torch.Tensor): mask for the encoded input batch of
                 shape B x seq_len.
 
         Returns:
-            base.ModuleOutput:  decoder output, and the previous hidden states
-                from the decoder LSTM.
+            base.ModuleOutput: decoder output, and the previous hidden states
+                from the decoder RNN.
         """
+        # The last hiddens includes the cell state, which isn't needed for the
+        # forward pass.
         embedded = self.embed(symbol)
-        # -> 1 x B x decoder_dim.
         last_h0, last_c0 = last_hiddens
-        context, attention_weights = self.attention(
+        context, _ = self.attention(
             last_h0.transpose(0, 1), encoder_out, encoder_mask
         )
         output, hiddens = self.module(
             torch.cat((embedded, context), 2), last_hiddens
         )
         output = self.dropout_layer(output)
-        return base.ModuleOutput(output, hiddens=hiddens)
+        return base.ModuleOutput(output, hiddens)
 
     @property
     def name(self) -> str:
         return "attentive LSTM"
 
 
-class HardAttentionLSTMDecoder(LSTMDecoder):
-    """Zeroth-order HMM hard attention LSTM decoder."""
+class HardAttentionRNNDecoder(RNNDecoder):
+    """Base module for zeroth-order HMM hard attention RNN decoders."""
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -304,33 +396,53 @@ def forward(
         output = self.output_proj(output)
         # Gets transition probabilities (alignment) for current states.
         alignment = self._alignment_step(decoded, encoder_out, encoder_mask)
-        return base.ModuleOutput(
-            output=output, embeddings=alignment, hiddens=hiddens
+        return base.ModuleOutput(output, hiddens, embeddings=alignment)
+
+    @property
+    def output_size(self) -> int:
+        return self.decoder_input_size + self.hidden_size
+
+
+class HardAttentionGRUDecoder(HardAttentionRNNDecoder):
+    """Zeroth-order HMM hard attention GRU decoder."""
+
+    def get_module(self) -> nn.GRU:
+        return nn.GRU(
+            self.embedding_size,
+            self.hidden_size,
+            batch_first=True,
+            bidirectional=self.bidirectional,
+            dropout=self.dropout,
+            num_layers=self.layers,
         )
 
+    @property
+    def name(self) -> str:
+        return "hard attention GRU"
+
+
+class HardAttentionLSTMDecoder(HardAttentionRNNDecoder):
+    """Zeroth-order HMM hard attention LSTM decoder."""
+
     def get_module(self) -> nn.LSTM:
         return nn.LSTM(
             self.embedding_size,
             self.hidden_size,
-            num_layers=self.layers,
-            dropout=self.dropout,
             batch_first=True,
             bidirectional=self.bidirectional,
+            dropout=self.dropout,
+            num_layers=self.layers,
         )
 
-    @property
-    def output_size(self) -> int:
-        return self.decoder_input_size + self.hidden_size
-
     @property
     def name(self) -> str:
         return "hard attention LSTM"
 
 
-class ContextHardAttentionLSTMDecoder(HardAttentionLSTMDecoder):
-    """First-order HMM hard attention LSTM."""
+class ContextHardAttentionRNNDecoder(HardAttentionRNNDecoder):
+    """Base module for first-order HMM hard attention RNN decoder."""
 
-    def __init__(self, *args, attention_context, **kwargs):
+    def __init__(self, attention_context, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.delta = attention_context
         # Window size must include center and both sides.
@@ -369,8 +481,7 @@ def _alignment_step(
             ],
             dim=1,
         )
-        # Gets probability of alignments.
-        # Masks padding.
+        # Gets probability of alignments, masking padding.
         alignment_probs = (
             alignment_probs * (~encoder_mask).unsqueeze(1) + defaults.EPSILON
         )
@@ -380,6 +491,22 @@ def _alignment_step(
         # Log probs for quicker computation.
         return alignment_probs.log()
 
+
+class ContextHardAttentionGRUDecoder(
+    ContextHardAttentionRNNDecoder, HardAttentionGRUDecoder
+):
+    """First-order HMM hard attention GRU decoder."""
+
+    @property
+    def name(self) -> str:
+        return "contextual hard attention GRU"
+
+
+class ContextHardAttentionLSTMDecoder(
+    ContextHardAttentionRNNDecoder, HardAttentionLSTMDecoder
+):
+    """First-order HMM hard attention LSTM decoder."""
+
     @property
     def name(self) -> str:
         return "contextual hard attention LSTM"
diff --git a/yoyodyne/models/modules/transformer.py b/yoyodyne/models/modules/transformer.py
index 5c0090a6..782e47da 100644
--- a/yoyodyne/models/modules/transformer.py
+++ b/yoyodyne/models/modules/transformer.py
@@ -1,4 +1,4 @@
-"""Transformer model classes."""
+"""Transformer module classes."""
 
 from typing import List, Optional, Tuple
 
@@ -46,7 +46,9 @@ def __init__(
         self.register_buffer("positional_encoding", positional_encoding)
 
     def forward(
-        self, symbols: torch.Tensor, mask: Optional[torch.Tensor] = None
+        self,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Computes the positional encoding.
 
@@ -616,7 +618,11 @@ class TransformerPointerDecoder(TransformerDecoder):
     """
 
     def __init__(
-        self, *args, separate_features, features_attention_heads, **kwargs
+        self,
+        *args,
+        separate_features,
+        features_attention_heads,
+        **kwargs,
     ):
         self.separate_features = separate_features
         self.features_attention_heads = features_attention_heads
diff --git a/yoyodyne/models/pointer_generator.py b/yoyodyne/models/pointer_generator.py
index f02420f5..38ff2007 100644
--- a/yoyodyne/models/pointer_generator.py
+++ b/yoyodyne/models/pointer_generator.py
@@ -7,7 +7,7 @@
 from torch import nn
 
 from .. import data, special
-from . import lstm, modules, transformer
+from . import modules, rnn, transformer
 
 
 class Error(Exception):
@@ -32,7 +32,10 @@ class GenerationProbability(nn.Module):
     bias: nn.Parameter
 
     def __init__(
-        self, embedding_size: int, hidden_size: int, attention_size: int
+        self,
+        embedding_size: int,
+        hidden_size: int,
+        attention_size: int,
     ):
         super().__init__()
         self.W_emb = nn.Linear(embedding_size, 1, bias=False)
@@ -88,7 +91,7 @@ def _get_loss_func(
         """Returns the actual function used to compute loss.
 
         This overrides the loss function behavior in
-        models.base.BaseEncoderDecoder because we need to use NLLLoss in
+        models.base.BaseModel because we need to use NLLLoss in
         order to factor the addition of two separate probability
         distributions. An NLLLoss-compatible implementation of label smoothing
         is also provided when label smoothing is requested.
@@ -146,10 +149,8 @@ def _smooth_nllloss(
         return loss
 
 
-class PointerGeneratorLSTMEncoderDecoder(
-    lstm.LSTMEncoderDecoder, PointerGenerator
-):
-    """Pointer-generator model with an LSTM backend.
+class PointerGeneratorRNNModel(rnn.RNNModel, PointerGenerator):
+    """Base class for pointer-generator models with an RNN backend.
 
     After:
         See, A., Liu, P. J., and Manning, C. D. 2017. Get to the point:
@@ -160,7 +161,12 @@ class PointerGeneratorLSTMEncoderDecoder(
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._check_layer_sizes()
+        if self.encoder_layers != self.decoder_layers:
+            raise Error(
+                "The number of encoder and decoder layers must match "
+                f"({self.encoder_layers} != {self.decoder_layers})"
+            )
+
         # We use the inherited defaults for the source embeddings/encoder.
         # Overrides classifier to take larger input.
         if not self.has_features_encoder:
@@ -168,14 +174,13 @@ def __init__(self, *args, **kwargs):
                 self.hidden_size + self.source_encoder.output_size,
                 self.target_vocab_size,
             )
-            self.generation_probability = GenerationProbability(  # noqa: E501
+            self.generation_probability = GenerationProbability(
                 self.embedding_size,
                 self.hidden_size,
                 self.source_encoder.output_size,
             )
         else:
             self.merge_h = nn.Linear(2 * self.hidden_size, self.hidden_size)
-            self.merge_c = nn.Linear(2 * self.hidden_size, self.hidden_size)
             self.features_attention = modules.attention.Attention(
                 self.features_encoder.output_size, self.hidden_size
             )
@@ -192,24 +197,6 @@ def __init__(self, *args, **kwargs):
                 + self.features_encoder.output_size,
             )
 
-    def get_decoder(self) -> modules.lstm.LSTMAttentiveDecoder:
-        return modules.lstm.LSTMAttentiveDecoder(
-            decoder_input_size=(
-                self.source_encoder.output_size
-                + self.features_encoder.output_size
-                if self.has_features_encoder
-                else self.source_encoder.output_size
-            ),
-            embeddings=self.embeddings,
-            embedding_size=self.embedding_size,
-            num_embeddings=self.vocab_size,
-            dropout=self.dropout,
-            bidirectional=False,
-            layers=self.decoder_layers,
-            hidden_size=self.hidden_size,
-            attention_input_size=self.source_encoder.output_size,
-        )
-
     def _check_layer_sizes(self) -> None:
         """Checks that encoder and decoder layers are the same number.
 
@@ -375,6 +362,257 @@ def decode(
         predictions = torch.stack(predictions).transpose(0, 1)
         return predictions
 
+
+class PointerGeneratorGRUModel(PointerGeneratorRNNModel):
+    """Pointer-generator model with an GRU backend."""
+
+    def get_decoder(self) -> modules.AttentiveGRUDecoder:
+        return modules.AttentiveGRUDecoder(
+            attention_input_size=self.source_encoder.output_size,
+            bidirectional=False,
+            decoder_input_size=(
+                self.source_encoder.output_size
+                + self.features_encoder.output_size
+                if self.has_features_encoder
+                else self.source_encoder.output_size
+            ),
+            dropout=self.dropout,
+            embeddings=self.embeddings,
+            embedding_size=self.embedding_size,
+            hidden_size=self.hidden_size,
+            layers=self.decoder_layers,
+            num_embeddings=self.vocab_size,
+        )
+
+    def decode_step(
+        self,
+        symbol: torch.Tensor,
+        last_hiddens: torch.Tensor,
+        source_indices: torch.Tensor,
+        source_enc: torch.Tensor,
+        source_mask: torch.Tensor,
+        features_enc: Optional[torch.Tensor] = None,
+        features_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Runs a single step of the decoder.
+
+        This predicts a distribution for one symbol.
+
+        Args:
+            symbol (torch.Tensor).
+            last_hiddens (torch.Tensor).
+            source_indices (torch.Tensor).
+            source_enc (torch.Tensor).
+            source_mask (torch.Tensor).
+            features_enc (Optional[torch.Tensor]).
+            features_mask (Optional[torch.Tensor]).
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor].
+        """
+        embedded = self.decoder.embed(symbol)
+        source_context, attention_weights = self.decoder.attention(
+            last_hiddens.transpose(0, 1), source_enc, source_mask
+        )
+        if self.has_features_encoder:
+            features_context, _ = self.features_attention(
+                last_hiddens.transpose(0, 1), features_enc, features_mask
+            )
+            # -> B x 1 x 4*hidden_size.
+            context = torch.cat([source_context, features_context], dim=2)
+        else:
+            context = source_context
+        _, h = self.decoder.module(
+            torch.cat((embedded, context), 2), last_hiddens
+        )
+        # -> B x 1 x hidden_size
+        hidden = h[-1, :, :].unsqueeze(1)
+        output_dist = self.classifier(torch.cat([hidden, context], dim=2))
+        output_dist = nn.functional.softmax(output_dist, dim=2)
+        # -> B x 1 x target_vocab_size.
+        ptr_dist = torch.zeros(
+            symbol.size(0),
+            self.target_vocab_size,
+            device=self.device,
+            dtype=attention_weights.dtype,
+        ).unsqueeze(1)
+        # Gets the attentions to the source in terms of the output generations.
+        # These are the "pointer" distribution.
+        # -> B x 1 x target_vocab_size.
+        ptr_dist.scatter_add_(
+            2, source_indices.unsqueeze(1), attention_weights
+        )
+        # Probability of generating (from output_dist).
+        gen_probs = self.generation_probability(context, hidden, embedded)
+        scaled_ptr_dist = ptr_dist * (1 - gen_probs)
+        scaled_output_dist = output_dist * gen_probs
+        return torch.log(scaled_output_dist + scaled_ptr_dist), h
+
+    def forward(
+        self,
+        batch: data.PaddedBatch,
+    ) -> torch.Tensor:
+        """Runs the encoder-decoder.
+
+        Args:
+            batch (data.PaddedBatch).
+
+        Returns:
+            torch.Tensor.
+        """
+        encoder_output = self.source_encoder(batch.source)
+        source_encoded = encoder_output.output
+        if encoder_output.has_hiddens:
+            last_hiddens = self._reshape_hiddens(
+                encoder_output.hiddens,
+                self.source_encoder.layers,
+                self.source_encoder.num_directions,
+            )
+        else:
+            last_hiddens = self.init_hiddens(
+                len(batch), self.source_encoder.layers
+            )
+        if not self.has_features_encoder:
+            if self.beam_width is not None and self.beam_width > 1:
+                raise NotImplementedError
+            else:
+                predictions = self.decode(
+                    source_encoded,
+                    batch.source.mask,
+                    batch.source.padded,
+                    last_hiddens,
+                    self.teacher_forcing if self.training else False,
+                    target=(batch.target.padded if batch.target else None),
+                )
+        else:
+            features_encoder_output = self.features_encoder(batch.features)
+            features_encoded = features_encoder_output.output
+            if features_encoder_output.has_hiddens:
+                last_hiddens = self._reshape_hiddens(
+                    features_encoder_output.hiddens,
+                    self.features_encoder.layers,
+                    self.features_encoder.num_directions,
+                )
+            else:
+                last_hiddens = self.init_hiddens(
+                    len(batch), self.source_encoder.layers
+                )
+            predictions = self.decode(
+                source_encoded,
+                batch.source.mask,
+                batch.source.padded,
+                last_hiddens,
+                self.teacher_forcing if self.training else False,
+                features_enc=features_encoded,
+                features_mask=batch.features.mask,
+                target=batch.target.padded if batch.target else None,
+            )
+        return predictions
+
+    @staticmethod
+    def _reshape_hiddens(
+        h: torch.Tensor,
+        layers: int,
+        num_directions: int,
+    ) -> torch.Tensor:
+        return h.view(layers, num_directions, h.size(1), h.size(2)).sum(axis=1)
+
+    @property
+    def name(self) -> str:
+        return "pointer-generator GRU"
+
+
+class PointerGeneratorLSTMModel(PointerGeneratorRNNModel):
+    """Pointer-generator model with an LSTM backend."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.has_features_encoder:
+            self.merge_c = nn.Linear(2 * self.hidden_size, self.hidden_size)
+
+    def get_decoder(self) -> modules.AttentiveLSTMDecoder:
+        return modules.AttentiveLSTMDecoder(
+            attention_input_size=self.source_encoder.output_size,
+            bidirectional=False,
+            decoder_input_size=(
+                self.source_encoder.output_size
+                + self.features_encoder.output_size
+                if self.has_features_encoder
+                else self.source_encoder.output_size
+            ),
+            dropout=self.dropout,
+            embeddings=self.embeddings,
+            embedding_size=self.embedding_size,
+            hidden_size=self.hidden_size,
+            layers=self.decoder_layers,
+            num_embeddings=self.vocab_size,
+        )
+
+    def decode_step(
+        self,
+        symbol: torch.Tensor,
+        last_hiddens: Tuple[torch.Tensor, torch.Tensor],
+        source_indices: torch.Tensor,
+        source_enc: torch.Tensor,
+        source_mask: torch.Tensor,
+        features_enc: Optional[torch.Tensor] = None,
+        features_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """Runs a single step of the decoder.
+
+        This predicts a distribution for one symbol.
+
+        Args:
+            symbol (torch.Tensor).
+            last_hiddens (Tuple[torch.Tensor, torch.Tensor]).
+            source_indices (torch.Tensor).
+            source_enc (torch.Tensor).
+            source_mask (torch.Tensor).
+            features_enc (Optional[torch.Tensor]).
+            features_mask (Optional[torch.Tensor]).
+
+        Returns:
+            Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]].
+        """
+        embedded = self.decoder.embed(symbol)
+        last_h0, last_c0 = last_hiddens
+        source_context, attention_weights = self.decoder.attention(
+            last_h0.transpose(0, 1), source_enc, source_mask
+        )
+        if self.has_features_encoder:
+            features_context, _ = self.features_attention(
+                last_h0.transpose(0, 1), features_enc, features_mask
+            )
+            # -> B x 1 x 4*hidden_size.
+            context = torch.cat([source_context, features_context], dim=2)
+        else:
+            context = source_context
+        _, (h, c) = self.decoder.module(
+            torch.cat((embedded, context), 2), (last_h0, last_c0)
+        )
+        # -> B x 1 x hidden_size
+        hidden = h[-1, :, :].unsqueeze(1)
+        output_dist = self.classifier(torch.cat([hidden, context], dim=2))
+        output_dist = nn.functional.softmax(output_dist, dim=2)
+        # -> B x 1 x target_vocab_size.
+        ptr_dist = torch.zeros(
+            symbol.size(0),
+            self.target_vocab_size,
+            device=self.device,
+            dtype=attention_weights.dtype,
+        ).unsqueeze(1)
+        # Gets the attentions to the source in terms of the output generations.
+        # These are the "pointer" distribution.
+        # -> B x 1 x target_vocab_size.
+        ptr_dist.scatter_add_(
+            2, source_indices.unsqueeze(1), attention_weights
+        )
+        # Probability of generating (from output_dist).
+        gen_probs = self.generation_probability(context, hidden, embedded)
+        scaled_ptr_dist = ptr_dist * (1 - gen_probs)
+        scaled_output_dist = output_dist * gen_probs
+        return torch.log(scaled_output_dist + scaled_ptr_dist), (h, c)
+
     def forward(
         self,
         batch: data.PaddedBatch,
@@ -402,7 +640,7 @@ def forward(
         if not self.has_features_encoder:
             if self.beam_width > 1:
                 raise NotImplementedError(
-                    f"Beam search not implemented for {self.name} model."
+                    f"Beam search not implemented for {self.name} model"
                 )
             else:
                 predictions = self.decode(
@@ -411,7 +649,7 @@ def forward(
                     batch.source.padded,
                     last_hiddens,
                     self.teacher_forcing if self.training else False,
-                    target=batch.target.padded if batch.target else None,
+                    target=(batch.target.padded if batch.target else None),
                 )
         else:
             features_encoder_output = self.features_encoder(batch.features)
@@ -429,7 +667,7 @@ def forward(
             if self.beam_width > 1:
                 # LSTM beam search does not work with pointer generator LSTM.
                 raise NotImplementedError(
-                    f"Beam search not implemented for {self.name} model."
+                    f"Beam search not implemented for {self.name} model"
                 )
             else:
                 predictions = self.decode(
@@ -446,7 +684,10 @@ def forward(
 
     @staticmethod
     def _reshape_hiddens(
-        h: torch.Tensor, c: torch.Tensor, layers: int, num_directions: int
+        h: torch.Tensor,
+        c: torch.Tensor,
+        layers: int,
+        num_directions: int,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         h = h.view(layers, num_directions, h.size(1), h.size(2)).sum(axis=1)
         c = c.view(layers, num_directions, c.size(1), c.size(2)).sum(axis=1)
@@ -457,9 +698,9 @@ def name(self) -> str:
         return "pointer-generator LSTM"
 
 
-class PointerGeneratorTransformerEncoderDecoder(
+class PointerGeneratorTransformerModel(
     PointerGenerator,
-    transformer.TransformerEncoderDecoder,
+    transformer.TransformerModel,
 ):
     """Pointer-generator model with a transformer backend.
 
@@ -491,19 +732,19 @@ def __init__(self, *args, features_attention_heads, **kwargs):
                 self.embedding_size,
             )
 
-    def get_decoder(self):
+    def get_decoder(self) -> modules.transformer.TransformerPointerDecoder:
         return modules.transformer.TransformerPointerDecoder(
+            decoder_input_size=self.source_encoder.output_size,
+            dropout=self.dropout,
             embeddings=self.embeddings,
             embedding_size=self.embedding_size,
+            hidden_size=self.hidden_size,
+            features_attention_heads=self.features_attention_heads,
+            layers=self.decoder_layers,
+            max_source_length=self.max_source_length,
             num_embeddings=self.vocab_size,
-            dropout=self.dropout,
-            decoder_input_size=self.source_encoder.output_size,
             source_attention_heads=self.source_attention_heads,
             separate_features=self.has_features_encoder,
-            features_attention_heads=self.features_attention_heads,
-            max_source_length=self.max_source_length,
-            layers=self.decoder_layers,
-            hidden_size=self.hidden_size,
         )
 
     def decode_step(
@@ -668,7 +909,6 @@ def forward(
             torch.Tensor.
         """
         source_encoded = self.source_encoder(batch.source).output
-
         if self.training and self.teacher_forcing:
             assert (
                 batch.target.padded is not None
diff --git a/yoyodyne/models/lstm.py b/yoyodyne/models/rnn.py
similarity index 82%
rename from yoyodyne/models/lstm.py
rename to yoyodyne/models/rnn.py
index 5bddd52d..742b3a2e 100644
--- a/yoyodyne/models/lstm.py
+++ b/yoyodyne/models/rnn.py
@@ -1,4 +1,4 @@
-"""LSTM model classes."""
+"""RNN model classes."""
 
 import argparse
 import heapq
@@ -11,32 +11,21 @@
 from . import base, embeddings, modules
 
 
-class LSTMEncoderDecoder(base.BaseEncoderDecoder):
-    """LSTM encoder-decoder without attention.
+class RNNModel(base.BaseModel):
+    """Base class for RNN models.
 
-    # TODO: Evaluate if this blurb is still correct.
-    We achieve this by concatenating the last (non-padding) hidden state of
+    In lieu of attention, we concatenate the last (non-padding) hidden state of
     the encoder to the decoder hidden state.
-
-    Args:
-        *args: passed to superclass.
-        **kwargs: passed to superclass.
     """
 
     # Constructed inside __init__.
+    classifier: nn.Linear
     h0: nn.Parameter
-    c0: nn.Parameter
 
-    def __init__(
-        self,
-        *args,
-        **kwargs,
-    ):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Initial hidden state whose parameters are shared across all examples.
-        self.h0 = nn.Parameter(torch.rand(self.hidden_size))
-        self.c0 = nn.Parameter(torch.rand(self.hidden_size))
         self.classifier = nn.Linear(self.hidden_size, self.target_vocab_size)
+        self.h0 = nn.Parameter(torch.rand(self.hidden_size))
 
     def init_embeddings(
         self,
@@ -54,37 +43,6 @@ def init_embeddings(
         """
         return embeddings.normal_embedding(num_embeddings, embedding_size)
 
-    def get_decoder(self) -> modules.lstm.LSTMDecoder:
-        return modules.lstm.LSTMDecoder(
-            decoder_input_size=self.source_encoder.output_size,
-            embeddings=self.embeddings,
-            embedding_size=self.embedding_size,
-            num_embeddings=self.vocab_size,
-            dropout=self.dropout,
-            bidirectional=False,
-            layers=self.decoder_layers,
-            hidden_size=self.hidden_size,
-        )
-
-    def init_hiddens(
-        self, batch_size: int
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Initializes the hidden state to pass to the LSTM.
-
-        Note that we learn the initial state h0 as a parameter of the model.
-
-        Args:
-            batch_size (int).
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor]: hidden cells for LSTM
-                initialization.
-        """
-        return (
-            self.h0.repeat(self.decoder_layers, batch_size, 1),
-            self.c0.repeat(self.decoder_layers, batch_size, 1),
-        )
-
     def decode(
         self,
         encoder_out: torch.Tensor,
@@ -322,10 +280,6 @@ def forward(
             predictions = predictions.transpose(0, 1)
             return predictions
 
-    @property
-    def name(self) -> str:
-        return "LSTM"
-
     @staticmethod
     def add_argparse_args(parser: argparse.ArgumentParser) -> None:
         """Adds LSTM configuration options to the argument parser.
@@ -337,7 +291,7 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
             "--bidirectional",
             action="store_true",
             default=defaults.BIDIRECTIONAL,
-            help="Uses a bidirectional encoder (LSTM-backed architectures "
+            help="Uses a bidirectional encoder (RNN-backed architectures "
             "only. Default: enabled.",
         )
         parser.add_argument(
@@ -346,21 +300,133 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
             dest="bidirectional",
         )
 
+    def get_decoder(self):
+        raise NotImplementedError
 
-class AttentiveLSTMEncoderDecoder(LSTMEncoderDecoder):
-    """LSTM encoder-decoder with attention."""
+    def init_hiddens(self, batch_size: int):
+        raise NotImplementedError
 
-    def get_decoder(self):
-        return modules.lstm.LSTMAttentiveDecoder(
+    @property
+    def name(self) -> str:
+        raise NotImplementedError
+
+
+class GRUModel(RNNModel):
+    """GRU encoder-decoder without attention."""
+
+    def get_decoder(self) -> modules.GRUDecoder:
+        return modules.GRUDecoder(
+            bidirectional=False,
             decoder_input_size=self.source_encoder.output_size,
-            embeddings=self.embeddings,
+            dropout=self.dropout,
             embedding_size=self.embedding_size,
+            embeddings=self.embeddings,
+            hidden_size=self.hidden_size,
+            layers=self.decoder_layers,
             num_embeddings=self.vocab_size,
-            dropout=self.dropout,
+        )
+
+    def init_hiddens(self, batch_size: int) -> torch.Tensor:
+        """Initializes the hidden state to pass to the RNN.
+
+        We treat the initial value as a model parameter.
+
+        Args:
+            batch_size (int).
+
+        Returns:
+            torch.Tensor: hidden state for initialization.
+        """
+        return self.h0.repeat(self.decoder_layers, batch_size, 1)
+
+    @property
+    def name(self) -> str:
+        return "GRU"
+
+
+class LSTMModel(RNNModel):
+    """LSTM encoder-decoder without attention.
+
+    Args:
+        *args: passed to superclass.
+        **kwargs: passed to superclass.
+    """
+
+    # This also needs an initial cell state parameter.
+    c0: nn.Parameter
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.c0 = nn.Parameter(torch.rand(self.hidden_size))
+
+    def get_decoder(self) -> modules.LSTMDecoder:
+        return modules.LSTMDecoder(
             bidirectional=False,
+            decoder_input_size=self.source_encoder.output_size,
+            dropout=self.dropout,
+            embedding_size=self.embedding_size,
+            embeddings=self.embeddings,
+            hidden_size=self.hidden_size,
             layers=self.decoder_layers,
+            num_embeddings=self.vocab_size,
+        )
+
+    def init_hiddens(self, batch_size: int) -> torch.Tensor:
+        """Initializes the hidden state to pass to the RNN.
+
+        We treat the initial value as a model parameter.
+
+        Args:
+            batch_size (int).
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor].
+        """
+        return (
+            self.h0.repeat(self.decoder_layers, batch_size, 1),
+            self.c0.repeat(self.decoder_layers, batch_size, 1),
+        )
+
+    @property
+    def name(self) -> str:
+        return "LSTM"
+
+
+class AttentiveGRUModel(GRUModel):
+    """GRU encoder-decoder with attention."""
+
+    def get_decoder(self) -> modules.AttentiveGRUDecoder:
+        return modules.AttentiveGRUDecoder(
+            attention_input_size=self.source_encoder.output_size,
+            bidirectional=False,
+            decoder_input_size=self.source_encoder.output_size,
+            dropout=self.dropout,
+            embeddings=self.embeddings,
+            embedding_size=self.embedding_size,
             hidden_size=self.hidden_size,
+            layers=self.decoder_layers,
+            num_embeddings=self.vocab_size,
+        )
+
+    @property
+    def name(self) -> str:
+        return "attentive GRU"
+
+
+class AttentiveLSTMModel(LSTMModel):
+    """LSTM encoder-decoder with attention."""
+
+    def get_decoder(self) -> modules.AttentiveLSTMDecoder:
+        return modules.AttentiveLSTMDecoder(
             attention_input_size=self.source_encoder.output_size,
+            bidirectional=False,
+            decoder_input_size=self.source_encoder.output_size,
+            dropout=self.dropout,
+            embeddings=self.embeddings,
+            embedding_size=self.embedding_size,
+            hidden_size=self.hidden_size,
+            layers=self.decoder_layers,
+            num_embeddings=self.vocab_size,
         )
 
     @property
diff --git a/yoyodyne/models/transducer.py b/yoyodyne/models/transducer.py
index ed2a58a1..853eeb94 100644
--- a/yoyodyne/models/transducer.py
+++ b/yoyodyne/models/transducer.py
@@ -8,11 +8,11 @@
 from torch import nn
 
 from .. import data, defaults, special, util
-from . import expert, lstm, modules
+from . import expert, modules, rnn
 
 
-class TransducerEncoderDecoder(lstm.LSTMEncoderDecoder):
-    """Transducer model with an LSTM backend.
+class TransducerRNNModel(rnn.RNNModel):
+    """Base class for transducer models.
 
     This uses a trained oracle for imitation learning edits.
 
@@ -51,23 +51,6 @@ def __init__(
         self.substitutions = self.actions.substitutions
         self.insertions = self.actions.insertions
 
-    def get_decoder(self) -> modules.lstm.LSTMDecoder:
-        return modules.lstm.LSTMDecoder(
-            decoder_input_size=(
-                self.source_encoder.output_size
-                + self.features_encoder.output_size
-                if self.has_features_encoder
-                else self.source_encoder.output_size
-            ),
-            embeddings=self.embeddings,
-            embedding_size=self.embedding_size,
-            num_embeddings=self.vocab_size,
-            dropout=self.dropout,
-            bidirectional=False,
-            layers=self.decoder_layers,
-            hidden_size=self.hidden_size,
-        )
-
     def forward(
         self,
         batch: data.PaddedBatch,
@@ -195,7 +178,9 @@ def decode(
                 break
             # Proceeds to make new edit; new action for all current decoding.
             action_count = torch.where(
-                not_complete.to(self.device), action_count + 1, action_count
+                not_complete.to(self.device),
+                action_count + 1,
+                action_count,
             )
             # Decoding.
             # We offset the action idx by the symbol vocab size so that we
@@ -204,7 +189,7 @@ def decode(
                 last_action.unsqueeze(dim=1) + self.vocab_offset,
                 last_hiddens,
                 encoder_out,
-                # To accomodate LSTMDecoder. See encoder_mask behavior.
+                # Accomodates RNNDecoder; see encoder_mask behavior.
                 ~(alignment.unsqueeze(1) + 1),
             )
             decoded, last_hiddens = (
@@ -215,7 +200,11 @@ def decode(
             # If given targets, asks expert for optimal actions.
             optim_actions = (
                 self.batch_expert_rollout(
-                    source, target, alignment, prediction, not_complete
+                    source,
+                    target,
+                    alignment,
+                    prediction,
+                    not_complete,
                 )
                 if target is not None
                 else None
@@ -225,7 +214,7 @@ def decode(
                 alignment,
                 input_length,
                 not_complete,
-                optim_actions=optim_actions if teacher_forcing else None,
+                optim_actions=(optim_actions if teacher_forcing else None),
             )
             alignment = self.update_prediction(
                 last_action, source, alignment, prediction
@@ -267,7 +256,7 @@ def decode_action_step(
         # Finds valid actions given remaining input length.
         end_of_input = (input_length - alignment) <= 1  # 1 -> Last char.
         valid_actions = [
-            self.compute_valid_actions(eoi) if nc else [self.actions.end_idx]
+            (self.compute_valid_actions(eoi) if nc else [self.actions.end_idx])
             for eoi, nc in zip(end_of_input, not_complete)
         ]
         # Masks invalid actions.
@@ -329,7 +318,7 @@ def choose_action(
         if optim_actions is None:
             # Argmax decoding.
             next_action = [
-                torch.argmax(probs, dim=0) if nc else self.actions.end_idx
+                (torch.argmax(probs, dim=0) if nc else self.actions.end_idx)
                 for probs, nc in zip(log_probs, not_complete)
             ]
         else:
@@ -338,7 +327,7 @@ def choose_action(
                 if self.expert.explore():
                     # Action is picked by random exploration.
                     next_action = [
-                        self.sample(probs) if nc else self.actions.end_idx
+                        (self.sample(probs) if nc else self.actions.end_idx)
                         for probs, nc in zip(log_probs, not_complete)
                     ]
                 else:
@@ -539,7 +528,7 @@ def training_step(self, batch: data.PaddedBatch, batch_idx: int) -> Dict:
     def validation_step(self, batch: data.PaddedBatch, batch_idx: int) -> Dict:
         predictions, loss = self(batch)
         # Evaluation requires prediction as a tensor.
-        predictions = self.convert_prediction(predictions)
+        predictions = self.convert_predictions(predictions)
         # Gets a dict of all eval metrics for this batch.
         val_eval_items_dict = {
             evaluator.name: evaluator.evaluate(
@@ -557,24 +546,23 @@ def predict_step(self, batch: Tuple[torch.tensor], batch_idx: int) -> Dict:
             batch,
         )
         # Evaluation requires prediction tensor.
-        return self.convert_prediction(predictions)
+        return self.convert_predictions(predictions)
 
-    def convert_prediction(self, prediction: List[List[int]]) -> torch.Tensor:
+    def convert_predictions(
+        self, predictions: List[List[int]]
+    ) -> torch.Tensor:
         """Converts prediction values to tensor for evaluator compatibility."""
         # FIXME: the two steps below may be partially redundant.
         # TODO: Clean this up and make it more efficient.
-        max_len = len(max(prediction, key=len))
-        for i, pred in enumerate(prediction):
+        max_len = len(max(predictions, key=len))
+        for i, pred in enumerate(predictions):
             pad = [self.actions.end_idx] * (max_len - len(pred))
             pred.extend(pad)
-            prediction[i] = torch.tensor(pred, dtype=torch.int)
-        prediction = torch.stack(prediction)
-        # Uses the same util that all other models use.
-        # This turns all symbols after the first EOS into PADs
-        # so prediction tensors match gold tensors.
-        return util.pad_tensor_after_eos(
-            torch.tensor(prediction),
-        )
+            predictions[i] = torch.tensor(pred, dtype=torch.int)
+        predictions = torch.stack(predictions)
+        # This turns all symbols after the first EOS into PAD so prediction
+        # tensors match gold tensors.
+        return util.pad_tensor_after_eos(predictions)
 
     def on_train_epoch_start(self) -> None:
         """Scheduler for oracle."""
@@ -591,9 +579,201 @@ def sample(log_probs: torch.Tensor) -> torch.Tensor:
                 break
         return action
 
+    def get_decoder(self):
+        raise NotImplementedError
+
     @property
     def name(self) -> str:
-        return "transducer"
+        raise NotImplementedError
 
 
 # TODO: Implement beam decoding.
+
+
+class TransducerGRUModel(TransducerRNNModel, rnn.GRUModel):
+    """Transducer with GRU backend."""
+
+    # h0: nn.Parameter
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # self.h0 = nn.Parameter(torch.rand(self.hidden_size))
+
+    def forward(
+        self,
+        batch: data.PaddedBatch,
+    ) -> Tuple[List[List[int]], torch.Tensor]:
+        """Runs the encoder-decoder model.
+
+        Args:
+            batch (data.PaddedBatch).
+
+        Returns:
+            Tuple[List[List[int]], torch.Tensor]: encoded prediction values
+                and loss tensor; due to transducer setup, prediction is
+                performed during training, so these are returned.
+        """
+        encoder_out = self.source_encoder(batch.source)
+        encoded = encoder_out.output[:, 1:, :]  # Ignores start symbol.
+        source_padded = batch.source.padded[:, 1:]
+        source_mask = batch.source.mask[:, 1:]
+        # Start of decoding.
+        if self.has_features_encoder:
+            features_encoder_out = self.features_encoder(batch.features)
+            features_encoded = features_encoder_out.output
+            if features_encoder_out.has_hiddens:
+                h_features = features_encoder_out.hiddens
+                last_hiddens = h_features.mean(dim=0, keepdim=True).expand(
+                    self.decoder_layers, -1, -1
+                )
+                last_hiddens = h_features
+            else:
+                last_hiddens = self.init_hiddens(source_mask.shape[0])
+            features_encoded = features_encoded.mean(dim=1, keepdim=True)
+            encoded = torch.cat(
+                (
+                    encoded,
+                    features_encoded.expand(-1, encoded.shape[1], -1),
+                ),
+                dim=2,
+            )
+        else:
+            last_hiddens = self.init_hiddens(source_mask.shape[0])
+        prediction, loss = self.decode(
+            encoded,
+            last_hiddens,
+            source_padded,
+            source_mask,
+            teacher_forcing=(self.teacher_forcing if self.training else False),
+            target=batch.target.padded if batch.target else None,
+            target_mask=batch.target.mask if batch.target else None,
+        )
+        return prediction, loss
+
+    def get_decoder(self) -> modules.GRUDecoder:
+        return modules.GRUDecoder(
+            bidirectional=False,
+            decoder_input_size=(
+                self.source_encoder.output_size
+                + self.features_encoder.output_size
+                if self.has_features_encoder
+                else self.source_encoder.output_size
+            ),
+            dropout=self.dropout,
+            embeddings=self.embeddings,
+            embedding_size=self.embedding_size,
+            layers=self.decoder_layers,
+            hidden_size=self.hidden_size,
+            num_embeddings=self.vocab_size,
+        )
+
+    @property
+    def name(self) -> str:
+        return "transducer GRU"
+
+
+class TransducerLSTMModel(TransducerRNNModel):
+    """Transducer with LSTM backend."""
+
+    h0: nn.Parameter
+    c0: nn.Parameter
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.h0 = nn.Parameter(torch.rand(self.hidden_size))
+        self.c0 = nn.Parameter(torch.rand(self.hidden_size))
+
+    def forward(
+        self,
+        batch: data.PaddedBatch,
+    ) -> Tuple[List[List[int]], torch.Tensor]:
+        """Runs the encoder-decoder model.
+
+        Args:
+            batch (data.PaddedBatch).
+
+        Returns:
+            Tuple[List[List[int]], torch.Tensor]: encoded prediction values
+                and loss tensor; due to transducer setup, prediction is
+                performed during training, so these are returned.
+        """
+        encoder_out = self.source_encoder(batch.source)
+        encoded = encoder_out.output[:, 1:, :]  # Ignores start symbol.
+        source_padded = batch.source.padded[:, 1:]
+        source_mask = batch.source.mask[:, 1:]
+        # Start of decoding.
+        if self.has_features_encoder:
+            features_encoder_out = self.features_encoder(batch.features)
+            features_encoded = features_encoder_out.output
+            if features_encoder_out.has_hiddens:
+                h_features, c_features = features_encoder_out.hiddens
+                h_features = h_features.mean(dim=0, keepdim=True).expand(
+                    self.decoder_layers, -1, -1
+                )
+                c_features = c_features.mean(dim=0, keepdim=True).expand(
+                    self.decoder_layers, -1, -1
+                )
+                last_hiddens = h_features, c_features
+            else:
+                last_hiddens = self.init_hiddens(
+                    source_mask.shape[0], self.decoder_layers
+                )
+            features_encoded = features_encoded.mean(dim=1, keepdim=True)
+            encoded = torch.cat(
+                (
+                    encoded,
+                    features_encoded.expand(-1, encoded.shape[1], -1),
+                ),
+                dim=2,
+            )
+        else:
+            last_hiddens = self.init_hiddens(source_mask.shape[0])
+        prediction, loss = self.decode(
+            encoded,
+            last_hiddens,
+            source_padded,
+            source_mask,
+            teacher_forcing=(self.teacher_forcing if self.training else False),
+            target=batch.target.padded if batch.target else None,
+            target_mask=batch.target.mask if batch.target else None,
+        )
+        return prediction, loss
+
+    def init_hiddens(
+        self, batch_size: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Initializes the hidden state to pass to the RNN.
+
+        We treat the initial value as a model parameter.
+
+        Args:
+            batch_size (int).
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor].
+        """
+        return (
+            self.h0.repeat(self.decoder_layers, batch_size, 1),
+            self.c0.repeat(self.decoder_layers, batch_size, 1),
+        )
+
+    def get_decoder(self) -> modules.LSTMDecoder:
+        return modules.LSTMDecoder(
+            bidirectional=False,
+            decoder_input_size=(
+                self.source_encoder.output_size
+                + self.features_encoder.output_size
+                if self.has_features_encoder
+                else self.source_encoder.output_size
+            ),
+            dropout=self.dropout,
+            embeddings=self.embeddings,
+            embedding_size=self.embedding_size,
+            layers=self.decoder_layers,
+            hidden_size=self.hidden_size,
+            num_embeddings=self.vocab_size,
+        )
+
+    @property
+    def name(self) -> str:
+        return "transducer LSTM"
diff --git a/yoyodyne/models/transformer.py b/yoyodyne/models/transformer.py
index 68252b55..74577d13 100644
--- a/yoyodyne/models/transformer.py
+++ b/yoyodyne/models/transformer.py
@@ -10,13 +10,12 @@
 from . import base, embeddings, modules
 
 
-class TransformerEncoderDecoder(base.BaseEncoderDecoder):
-    """Transformer encoder-decoder.
+class TransformerModel(base.BaseModel):
+    """Base class for transformer models.
 
     Args:
         source_attention_heads (int).
         max_source_length (int).
-        *args: passed to superclass.
         **kwargs: passed to superclass.
     """
 
@@ -33,7 +32,9 @@ def __init__(
     ):
         self.source_attention_heads = source_attention_heads
         super().__init__(
-            *args, source_attention_heads=source_attention_heads, **kwargs
+            *args,
+            source_attention_heads=source_attention_heads,
+            **kwargs,
         )
         self.classifier = nn.Linear(
             self.embedding_size, self.target_vocab_size
@@ -53,17 +54,17 @@ def init_embeddings(
         """
         return embeddings.xavier_embedding(num_embeddings, embedding_size)
 
-    def get_decoder(self):
+    def get_decoder(self) -> modules.transformer.TransformerDecoder:
         return modules.transformer.TransformerDecoder(
             decoder_input_size=self.source_encoder.output_size,
+            dropout=self.dropout,
             embeddings=self.embeddings,
             embedding_size=self.embedding_size,
+            hidden_size=self.hidden_size,
+            layers=self.decoder_layers,
+            max_source_length=self.max_source_length,
             num_embeddings=self.vocab_size,
-            dropout=self.dropout,
             source_attention_heads=self.source_attention_heads,
-            max_source_length=self.max_source_length,
-            layers=self.decoder_layers,
-            hidden_size=self.hidden_size,
         )
 
     def _decode_greedy(
@@ -103,7 +104,10 @@ def _decode_greedy(
             target_mask = torch.ones_like(target_tensor, dtype=torch.float)
             target_mask = target_mask == 0
             decoder_output = self.decoder(
-                encoder_hidden, source_mask, target_tensor, target_mask
+                encoder_hidden,
+                source_mask,
+                target_tensor,
+                target_mask,
             ).output
             logits = self.classifier(decoder_output)
             last_output = logits[:, -1, :]  # Ignores EOS.
@@ -144,7 +148,9 @@ def forward(
             # Initializes the start symbol for decoding.
             starts = (
                 torch.tensor(
-                    [special.START_IDX], device=self.device, dtype=torch.long
+                    [special.START_IDX],
+                    device=self.device,
+                    dtype=torch.long,
                 )
                 .repeat(batch.target.padded.size(0))
                 .unsqueeze(1)
@@ -155,7 +161,10 @@ def forward(
             )
             encoder_output = self.source_encoder(batch.source).output
             decoder_output = self.decoder(
-                encoder_output, batch.source.mask, target_padded, target_mask
+                encoder_output,
+                batch.source.mask,
+                target_padded,
+                target_mask,
             ).output
             logits = self.classifier(decoder_output)
             output = logits[:, :-1, :]  # Ignore EOS.
diff --git a/yoyodyne/predict.py b/yoyodyne/predict.py
index 4aaf308b..4b057b68 100644
--- a/yoyodyne/predict.py
+++ b/yoyodyne/predict.py
@@ -36,7 +36,7 @@ def get_datamodule_from_argparse_args(
         data.DataModule.
     """
     separate_features = args.features_col != 0 and args.arch in [
-        "pointer_generator_lstm",
+        "pointer_generator_rnn",
         "pointer_generator_transformer",
         "transducer",
     ]
@@ -59,14 +59,14 @@ def get_datamodule_from_argparse_args(
 
 def get_model_from_argparse_args(
     args: argparse.Namespace,
-) -> models.BaseEncoderDecoder:
+) -> models.BaseModel:
     """Creates the model from CLI arguments.
 
     Args:
         args (argparse.Namespace).
 
     Returns:
-        models.BaseEncoderDecoder.
+        models.BaseModel.
     """
     model_cls = models.get_model_cls_from_argparse_args(args)
     kwargs = {}
@@ -89,7 +89,7 @@ def _mkdir(output: str) -> None:
 
 def predict(
     trainer: lightning.Trainer,
-    model: models.BaseEncoderDecoder,
+    model: models.BaseModel,
     datamodule: data.DataModule,
     output: str,
 ) -> None:
@@ -131,7 +131,9 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
     """
     # Path arguments.
     parser.add_argument(
-        "--checkpoint", required=True, help="Path to checkpoint (.ckpt)."
+        "--checkpoint",
+        required=True,
+        help="Path to checkpoint (.ckpt).",
     )
     parser.add_argument(
         "--model_dir",
diff --git a/yoyodyne/sizing.py b/yoyodyne/sizing.py
index 451d68af..0cdc3f82 100644
--- a/yoyodyne/sizing.py
+++ b/yoyodyne/sizing.py
@@ -51,7 +51,7 @@ def _optimal_batch_size(
 def find_batch_size(
     method: str,
     trainer: lightning.Trainer,
-    model: models.BaseEncoderDecoder,
+    model: models.BaseModel,
     datamodule: data.DataModule,
     *,
     steps_per_trial: int = defaults.FIND_BATCH_SIZE_STEPS_PER_TRIAL,
@@ -73,7 +73,7 @@ def find_batch_size(
             the "optimal" batch size, using the gradient accumulation trick if
             necessary.)
         trainer (lightning.Trainer).
-        model (models.BaseEncoderDecoder).
+        model (models.BaseModel).
         datamodule (data.DataModule).
         steps_per_trial (int, optional): number of steps to run with a given
             batch size.
diff --git a/yoyodyne/train.py b/yoyodyne/train.py
index e5305b01..08fd5f1d 100644
--- a/yoyodyne/train.py
+++ b/yoyodyne/train.py
@@ -141,10 +141,13 @@ def get_datamodule_from_argparse_args(
         data.DataModule.
     """
     separate_features = args.features_col != 0 and args.arch in [
+        "hard_attention_gru",
         "hard_attention_lstm",
+        "pointer_generator_gru",
         "pointer_generator_lstm",
         "pointer_generator_transformer",
-        "transducer",
+        "transducer_grm",
+        "transducer_lstm",
     ]
     datamodule = data.DataModule(
         train=args.train,
@@ -171,7 +174,7 @@ def get_datamodule_from_argparse_args(
 def get_model_from_argparse_args(
     args: argparse.Namespace,
     datamodule: data.DataModule,
-) -> models.BaseEncoderDecoder:
+) -> models.BaseModel:
     """Creates the model.
 
     Args:
@@ -179,16 +182,17 @@ def get_model_from_argparse_args(
         datamodule (data.DataModule).
 
     Returns:
-        models.BaseEncoderDecoder.
+        models.BaseModel.
     """
     model_cls = models.get_model_cls(args.arch)
     # TODO(#156): add callback interface to check this.
     if (
         args.arch
         in [
-            "pointer_generator_lstm",
+            "pointer_generator_rnn",
             "pointer_generator_transformer",
-            "transducer",
+            "transducer_gru",
+            "transducer_lstm",
         ]
         and not args.tie_embeddings
     ):
@@ -201,22 +205,18 @@ def get_model_from_argparse_args(
     )
     # Instantiates expert, if needed.
     expert = None
-    if args.arch in ["transducer"]:
+    if args.arch in ["transducer_gru", "transducer_lstm"]:
         sed_params_paths = (
             args.sed_params
             if args.sed_params
             else f"{args.model_dir}/{args.experiment}/sed.pkl"
         )
-        expert = (
-            models.expert.get_expert(
-                datamodule.train_dataloader().dataset,
-                epochs=args.oracle_em_epochs,
-                oracle_factor=args.oracle_factor,
-                sed_params_path=sed_params_paths,
-                read_from_file=os.path.isfile(sed_params_paths),
-            )
-            if args.arch in ["transducer"]
-            else None
+        expert = models.expert.get_expert(
+            datamodule.train_dataloader().dataset,
+            epochs=args.oracle_em_epochs,
+            oracle_factor=args.oracle_factor,
+            sed_params_path=sed_params_paths,
+            read_from_file=os.path.isfile(sed_params_paths),
         )
     scheduler_kwargs = schedulers.get_scheduler_kwargs_from_argparse_args(args)
     # We use a separate features encoder if the datamodule has features, and
@@ -225,20 +225,24 @@ def get_model_from_argparse_args(
     #   the requested module), or
     # - no specific features encoder module is requested, but the model
     #   requires that we use a separate features encoder (in which case we use
-    #   the same module as the source encoder).
+    #   the same type of module as the source encoder).
     features_encoder_cls = (
         models.modules.get_encoder_cls(
-            encoder_arch=args.features_encoder_arch, model_arch=args.arch
+            encoder_arch=args.features_encoder_arch,
+            model_arch=args.arch,
         )
         if datamodule.has_features
         and (
             args.features_encoder_arch
             or args.arch
             in [
+                "hard_attention_gru",
                 "hard_attention_lstm",
+                "pointer_generator_gru",
                 "pointer_generator_lstm",
                 "pointer_generator_transformer",
-                "transducer",
+                "transducer_gru",
+                "transducer_lstm",
             ]
         )
         else None
@@ -279,7 +283,6 @@ def get_model_from_argparse_args(
         scheduler_kwargs=scheduler_kwargs,
         source_attention_heads=args.source_attention_heads,
         source_encoder_cls=source_encoder_cls,
-        start_idx=datamodule.index.start_idx,
         target_vocab_size=(
             len(expert.actions)
             if expert is not None
@@ -353,7 +356,8 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
         help="Path to input validation data TSV.",
     )
     parser.add_argument(
-        "--train_from", help="Path to checkpoint used to resume training."
+        "--train_from",
+        help="Path to checkpoint used to resume training.",
     )
     # Other training arguments.
     parser.add_argument(
@@ -386,12 +390,12 @@ def add_argparse_args(parser: argparse.ArgumentParser) -> None:
     evaluators.add_argparse_args(parser)
     metrics.add_argparse_args(parser)
     models.add_argparse_args(parser)
-    models.modules.add_argparse_args(parser)
-    models.BaseEncoderDecoder.add_argparse_args(parser)
-    models.HardAttentionLSTM.add_argparse_args(parser)
-    models.LSTMEncoderDecoder.add_argparse_args(parser)
-    models.TransformerEncoderDecoder.add_argparse_args(parser)
     models.expert.add_argparse_args(parser)
+    models.modules.add_argparse_args(parser)
+    models.BaseModel.add_argparse_args(parser)
+    models.HardAttentionRNNModel.add_argparse_args(parser)
+    models.RNNModel.add_argparse_args(parser)
+    models.TransformerModel.add_argparse_args(parser)
     schedulers.add_argparse_args(parser)
     sizing.add_argparse_args(parser)
     # Trainer arguments.
diff --git a/yoyodyne/util.py b/yoyodyne/util.py
index ef958ffa..5dc70d83 100644
--- a/yoyodyne/util.py
+++ b/yoyodyne/util.py
@@ -32,7 +32,7 @@ def pad_tensor_after_eos(
 ) -> torch.Tensor:
     """Replaces everything after an EOS token with PADs.
 
-    Cuts off tensors at the first end_idx, and replaces the rest of the
+    Cuts off tensors at the first END_IDX, and replaces the rest of the
     predictions with PAD_IDX, as these can be erroneously decoded while the
     rest of the batch is finishing decoding.
 


Test patch:
diff --git a/tests/model_test.py b/tests/model_test.py
index 6bd09bdb..70c96316 100644
--- a/tests/model_test.py
+++ b/tests/model_test.py
@@ -6,15 +6,22 @@
 @pytest.mark.parametrize(
     ["arch", "expected_cls"],
     [
-        ("attentive_lstm", models.AttentiveLSTMEncoderDecoder),
-        ("hard_attention_lstm", models.HardAttentionLSTM),
-        ("lstm", models.LSTMEncoderDecoder),
+        ("attentive_gru", models.AttentiveGRUModel),
+        ("attentive_lstm", models.AttentiveLSTMModel),
+        ("gru", models.GRUModel),
+        ("hard_attention_lstm", models.HardAttentionLSTMModel),
+        ("lstm", models.LSTMModel),
+        (
+            "pointer_generator_gru",
+            models.PointerGeneratorGRUModel,
+        ),
         (
             "pointer_generator_lstm",
-            models.PointerGeneratorLSTMEncoderDecoder,
+            models.PointerGeneratorLSTMModel,
         ),
-        ("transducer", models.TransducerEncoderDecoder),
-        ("transformer", models.TransformerEncoderDecoder),
+        ("transducer_gru", models.TransducerGRUModel),
+        ("transducer_lstm", models.TransducerLSTMModel),
+        ("transformer", models.TransformerModel),
     ],
 )
 def test_get_model_cls(arch, expected_cls):