Compared to decoder-only architectures commonly used for language modeling and generation,
an encoder-decoder architecture allows a bidirectional description representation (tokens at the
beginning of the description can attend to tokens at the end) and the extra ﬂexibility to untie the
encoder structure from the decoder. Because problem descriptions are on average twice as long as
their corresponding human solutions, we use an asymmetric architecture with 1536 tokens for the
encoder but only 768 tokens for the decoder. We further found that using a shallow encoder and a
deep decoder signiﬁcantly improves the eﬃciency of training without hurting problem solve rate. The
exact architectures for our models are listed in Table 3. The 9B and 41B models were trained using
model parallelism, with 1 key and value head per shard. We built our model using JAX (Bradbury
et al., 2018) and Haiku (Hennigan et al., 2020), and trained them on TPUv4 accelerators using
bﬂoat16 precision.
alphacode