Transformer モデルの仕組みを JAX/Flax で実装しながら解説してみる(パート2)

パート1はこちら。enakai00.hatenablog.com なんの話かと言うと 最近、大規模言語モデルを用いたチャットシステムがよく話題になりますが、言語モデルの性能が大きく向上するきっかけとなったのが、下記の論文で公表された「Transformer」のアーキテクチャーです。arxiv.orgここでは、JAX/Flax を用いて Transformer を実…