End-to-End Long Document Summarization using Gradient Caching

要約

長文要約のための変換器ベースのエンコーダ・デコーダモデルの学習は、学習時の二次関数的なメモリ消費のため、大きな課題となっている。テスト時に入力長を拡張するアプローチがいくつか提案されているが、これらのアプローチによる学習は依然として困難であり、入力文書の切り捨てを必要とし、学習条件とテスト条件のミスマッチを引き起こす。本研究では、既存の変換器ベースのエンコーダ・デコーダモデルのエンドツーエンドの学習を、文書を切り捨てることなく文書全体を使って可能にするアプローチであるCachED (Gradient $textbf{Cach}$ing for $textbf{E}$ncoder-$textbf{D}$ecoder models) を提案する。具体的には、非重複スライディングウィンドウを入力文書に適用し、その後デコーダで融合する。バックプロパゲーションの間、勾配はデコーダでキャッシュされ、勾配チェックポインティングと同様に、隠れベクトルを再計算することにより、チャンク単位でエンコーダを通過する。長文要約の実験では、BARTをCachED BARTに拡張し、学習中に500K以上のトークンを処理し、追加のパラメータを使用することなく優れた性能を達成した。

要約(オリジナル)

Training transformer-based encoder-decoder models for long document summarization poses a significant challenge due to the quadratic memory consumption during training. Several approaches have been proposed to extend the input length at test time, but training with these approaches is still difficult, requiring truncation of input documents and causing a mismatch between training and test conditions. In this work, we propose CachED (Gradient $\textbf{Cach}$ing for $\textbf{E}$ncoder-$\textbf{D}$ecoder models), an approach that enables end-to-end training of existing transformer-based encoder-decoder models, using the entire document without truncation. Specifically, we apply non-overlapping sliding windows to input documents, followed by fusion in decoder. During backpropagation, the gradients are cached at the decoder and are passed through the encoder in chunks by re-computing the hidden vectors, similar to gradient checkpointing. In the experiments on long document summarization, we extend BART to CachED BART, processing more than 500K tokens during training and achieving superior performance without using any additional parameters.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL