長文向けの文書要約モデル「top-down transformer」

こんにちは。今回は、こちらの論文にて提唱されている文書要約モデル「top-down transformer」に関して、記事を書きます。
LONG DOCUMENT SUMMARIZATION WITH TOP-DOWN AND BOTTOM-UP INFERENCE

※本記事にある画像は、当論文より引用しています。

背景

長文要約

改めて文書要約についてですが、これはある文書の内容を簡潔に記述した文書を作るタスクです。
文書要約の中でも、比較的長い文書を対象にしたものを、ここでは長文要約と呼ぶことにします。長文の具体的な数値的基準は定かではないですが、一例としては「書籍内の章１つ文」くらいのイメージです。

当然かもしれませんが、比較的短い文書の要約に比べると、長文要約は人間にとって難しくて時間のかかる作業です。そのため、長文要約は短文の要約に比べて、AIによる自動化の需要が大きいと言えます。

近年の研究の経緯1

2019年から2020年あたりにかけて、いくつかの有力な文書要約モデルが発表されました。その中でも有名なものは、BERTSum、BART、PEGASUS、GPT-2、T5などです（ただし、これらのモデルすべてが文書要約に特化したモデルというわけではなく、GPT-2やT5は他の言語タスクにもよく用いられます）。

しかし、transformerのモデル構造に基づくモデル全般について言えることですが、難点があります。それは、入力の文書が長いほど、計算量がその二乗規模で増えてしまう難点です。文書が長いとは、トークン（端的に言えば単語）の数が多いということです。

この規模の計算が発生する箇所は、各attention、特にフル(full)のself-attentionの計算です。self-attentionは、いわば文書内の各トークン間の関連度を格納した行列です。self-attentionは、各トークンから見て他のどのトークンに注目（attentionという名の由来）すればいいのかを示す情報のようなものです。
仮にトークンが5個なら、self-attentionは5*5で25組分の関連度を持つことになります。先ほどself-attentionに対して「フルの」と形容したのは、各トークン間の関連度をすべて加味することを表します。ちなみに、ここでいう「組」は「組み合わせ」とは異なり、順序が考慮されるものです。例えば、「トークンA、トークンB」と「トークンB、トークンA」は別物です。

近年の研究の経緯2

2020年に発表されたlongformerというモデルには、この難点を解消する工夫が入っています。
それは、self-attentionの計算の際、すべての「２つのトークンの組」についての関連度の計算をやめることです。代わりに関連度の計算の対象を、各トークンごとに、そのトークンに近いいくつかのトークンに絞ることです。このような形式のattentionは、window attentionと呼ばれます。イメージとしては次のような図になります。

この図の各マスは、２つのトークンの組に相当します。この図は、すべての「２つのトークンの組」のマトリックスになります。
色の付いているマスは、関連度が加味されることを意味します。色の中でも特に濃い緑（ダークグリーン）のマスは、同じトークン自身との組に相当します。この図の場合、各トークンは、そのトークン自身に加え、前後それぞれにつき直近２トークン分まで（つまり前後合わせて計４トークン分）、関連度が加味されることになります。

すると、例えば文書内の全トークン数をN、各トークンごとの関連度計算対象とするトークンの数をwとした際、attentionの計算対象数が「Nの二乗」から（約）「N*w」に削減されます。これによって計算コストを削減できます。
wはウインドウサイズと呼ばれ、ウインドウのスペルwindowのイニシャルに由来します。ウインドウサイズは、あらかじめ値を設定しておくことがよくあります。上の図でいうと、ウインドウサイズは４です。

他には、同じ2020年にBigBirdというモデルが発表されました。このモデルにも、計算コストを減らす工夫が入っています。

しかし、longformerもBigBirdも、フルのattentionに比べると離れたトークン間のattentionの情報が欠落しますから、離れたトークン間の関係性を捉えにくくなります。結局、計算コストと、attentionの情報の欠落の両面で、ジレンマとなります。

当論文の提唱内容

ここから、セグメントという概念が登場します。端的に言えば、セグメントは文に相当するものです。１つの入力文書は複数のセグメントから成り、さらに各セグメントは複数のトークンから成ります。

当論文の提唱するモデルtop-down transformerは、抽象型要約のモデルです。構造は次のようになっています。

大きく、bottom-up推論という処理とtop-down推論という処理から成ります。特にtop-down推論の新規性を強調したいとの理由で、top-downという単語がモデル名に使われたようです。

モデル構造の第一の要点

２つの要点があります。第一の要点は、下の赤枠で囲った箇所です。

この赤枠で囲った箇所は、入力の文書の各トークンの表現（特徴量表現）から、入力の文書の各セグメントの表現を作る処理を意味します。

この処理の流れを簡略化して述べると、まずは各トークン表現を、セグメントごとの表現になるようにプーリングします。単純なやり方としては、セグメントごとに、そのセグメント内の各トークンやその周辺のトークンの表現の平均値を取ることです。いわゆる平均プーリングです。
その後、セグメント単位での（フルの）self-attentionを計算し、それを用いて各セグメントの表現を作成します。

ただし、後述する実験から示唆されるように、平均プーリングよりも効果的なプーリングがあります。それはadaptiveプーリングと呼ばれるものです。adaptiveプーリングでは、比較的重要なトークンに重みが付けられ、重みを加味した加重平均によりプーリングされます。

モデル構造の第二の要点

第二の要点は、下の緑枠で囲った箇所です。

この箇所の処理の中心は、各セグメント表現と各トークン表現の両方を用いて、各セグメントと各トークンの横断のattentionを計算することです。transformerのデコーダのマルチヘッドattentionと似たような形になっており、今回の場合、KeyとValueには各セグメント表現が、Queryには各トークン表現が割り当てられます。

結果として、各トークンに対する新しい表現を得ることができます。この新しい各トークン表現は、各セグメントとの間の関連度の情報を含んでいることになると思います。なぜなら、各セグメント表現は各セグメント間のattentionに基づいており、かつ新しい各トークン表現の計算には各セグメント表現が用いられているからです。
各トークンは各セグメントとの間の関連度を加味できていますし、各セグメントはそのセグメントに対応する各トークンと関係しています。よって、これで離れたトークン間の関連度も加味されたことになると思います。それも、「各トークン間のフルのattentionを計算することなく」です。

それを例として図示するとこのような感じです。

この図は、任意の２つのトークンの関係を表しています。１つ目のトークンAは、２つ目のトークンBに対応するセグメントと関係しています。もちろん２つ目のトークンBも、２つ目のトークンBに対応するセグメントと関係しています。
これら２種類の関係はいずれもattentionとして計算されるものであり、そして新しく出来上がる各トークン表現はこれらのattentionが加味されたものです。離れたトークン間の関連度も加味されると述べたのは、そういう理由です。

最後に、新しい各トークン表現、およびデコーダによる既存の出力（トークンの系列）をデコーダに入力し、モデル全体として１つの出力を出します。この繰り返しで、複数のトークンが出力され、そのまとまりが要約結果になると見られます。

計算量

各attentionにおける関連度計算対象の組の数は、次のようになります。

attention	組の数
bottom-up推論の前半にあるattention(トークン単位のself-attention)	N*w
bottom-up推論の後半にあるattention(セグメント単位のself-attention)	Mの二乗
top-down推論にあるattention(セグメントとトークンのcross-attention)	M*N

最初のattentionは、longformerなどに用いられているwindow attentionと同様です。これら組の数３つ分を足し合わせても、トークン単位のフルのself-attentionにおける関連度計算対象の組の数であるNの二乗に比べれば、十分に小さいことが分かります。ふつう、Mもwも、Nに比べれば十分小さい値だからです。
これで、計算量を抑えられていることを確認できました。

実験

当論文には、著者たちによる文書要約の実験とその結果について書かれています。見てみると、既存のモデルに比べてtop-down transformerは良いスコアを出せており、top-down transformerの性能の高さが示唆されます。

以下、データセットごとの実験結果を載せたうえで、実験結果に対する考察を記述します。

データセットの一覧、およびデータセットごとの実験結果

実験は、次の各データセットごとに行われました。

次にデータセットごとの実験結果ですが、あらかじめ注記しておきます。スコアの字体が太字のものは１位、下線のものは２位を意味しますが、これはoracleというものを用いない場合の順位です。

また、用いられた評価指標はそれぞれROUGE-1、ROUGE-2、ROUGE-Lです。

データセットごとの性能検証結果は以下の通りです。

PubMedとarXiv

CNN-DM

TVMegasSiteとForeverDreaming

BookSum-Chapter-Level

BookSum-Book-Level

実験結果に対する考察

おおむね、top-down transformerのプーリング処理では、平均プーリングよりもadaptiveプーリングを用いたほうが良いスコアを出せています。そのため、top-down transformerについてはadaptiveプーリング（AdaPool）を用いる場合を前提として、他のモデルとの性能を比較します。

top-down transformerの優勢性が比較的顕著に表れたデータセットは、TVMegaSiteとForeverDreamingだと思います。次いで、PubMedとarXivでもそれなりの優位性を出せています。

BookSum-Chapter-Levelでは、興味深いことにlongformerやBigBirdのスコアが割と低く、それよりもBART、T5、PEGASUSのほうがハイスコアになりました。それでもtop-down transformerのほうが一枚上手でしたが、スコアはBART、T5、PEGASUSと少ししか変わりません。しかも、１文書あたりのトークン数は約5,000と、結構な量があるにも関わらずです。
そのため、比較的長めの文書でも、BART、T5、PEGASUSもまだまだ十分高い性能を誇りうると心得ておいたほうが良さそうです。

とはいうものの、さすがにBookSum-Book-Levelのデータセットになると、top-down transformerとBART、T5、PEGASUSのスコアの差が顕著に表れます。このデータセットは、１文書あたりのトークン数が約11万もあります。
ただし、GPT-3に基づくある強化学習のモデルがあり、スコアとしてtop-down transformerと同レベルとなっています。

CNN-DMについては、BARTを超えたもののあまりBARTと変わらないと結果となりました。CNN-DMの１文書あたりのトークン数は約900と少なめなので、top-down transformerの強みがあまり活きなかったのかもしれません。前述の実験結果からしても、やはり１文書あたりのトークン数は少なからずスコアに影響するように見えます。
しかし逆に言えば、短めの文書でもBARTなど既存の文書要約モデルと同等以上の性能を発揮しうることになります。

終わりに

以上のようにして、top-down transformerの仕組み、および長文要約の可能性が示唆される実験結果を確認することができました。
最後までお読みいただき、ありがとうございました。