让大文档摘要模型全程不截断训练，突破长文本处理瓶颈。
End-to-End Long Document Summarization using Gradient Caching
- 用滑动窗口分段处理文档，解码器融合信息，避免训练时截断。
- 通过梯度缓存与重计算，实现端到端训练，支持超50万词元输入。
- 无需额外参数，可直接扩展现有模型，适合长文本生成任务。
基于Transformer的编码器-解码器模型在长文档摘要任务中面临训练时二次内存消耗的问题。尽管已有方法可延长测试时输入长度，但训练仍需截断文档，导致训练与测试条件不一致。本文提出CachED（Gradient Caching for Encoder-Decoder models），使现有模型能无截断地端到端训练。具体地，对输入文档采用非重叠滑动窗口，解码器进行信息融合；反向传播时，梯度在解码器缓存，并通过重计算隐藏状态分块传回编码器，类似梯度检查点机制。在长文档摘要实验中，将BART扩展为CachED BART，训练时处理超过50万词元，性能更优且未引入额外参数。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。