Longformer: The Long-Document Transformer

Iz Beltagy, Matthew E. Peters, Arman Cohan

Статья: https://arxiv.org/abs/2004.05150

Код: https://github.com/allenai/longformer

Longformer также на подходе (Work in Progress) в общеизвестной библиотечке от Huggingface: https://huggingface.co/transformers/model_doc/longformer.html



Свежие трансформеры. На этот раз от Allen Institute for AI.



Общеизвестная проблема трансформера — квадратичная относительно размера входа сложность механизма внимания. Из-за этого, в частности, нет нормальной возможности работать с длинными документами, которые не влезают целиком в attention span трансформера (обычно не более 512 токенов). Приходится исхитряться, и наиболее частый подход -- резать на окна фиксированного размера (обычно с перекрытием) и бежать по всей последовательности, как-то потом агрегируя активации каждого из окон. Хочется уметь делать это естественнее.



Проблему уже пытались решить по-разному, например, через оптимизацию внимания или через внедрение какой-либо памяти.



Из подходов первого типа, которые наиболее на слуху, можно вспомнить, пожалуй, Sparse Transformer от OpenAI (https://arxiv.org/abs/1904.10509, https://t.me/gonzo_ML/65) или Reformer от Гугла (https://arxiv.org/abs/2001.04451, https://t.me/gonzo_ML/176). В первом были кастомные разреженные ядра, во втором приближённое вычисление внимания через locality-sensitive hashing. Из интересных был ещё также Adaptive attention span (https://arxiv.org/abs/1905.07799, https://t.me/gonzo_ML/99).



Из второго типа можно вспомнить Transformer-XL (https://arxiv.org/abs/1901.02860, https://t.me/gonzo_ML/62), а также недавний Compressive Transformer от DeepMind (https://arxiv.org/abs/1911.05507, https://t.me/gonzo_ML/165).



В текущей работе делается очередной подход к разреженному вниманию, чтобы можно было работать с длинными документами.



Предлагается следующее: вместо полного n^2 внимания делаем более гибкие варианты, акцентирующиеся на локальном контексте (который в разных попытках изучения Берта показал свою важность), а также добавляем когда надо элементы глобального контекста.



Локальный контекст добавляется через внимание скользящим окном (разрешаем self-attention только внутри окна фиксированного размера), возможно также делая это окно разреженным (dilated). Это всё уже совсем похоже на свёртки, только не с фиксированным ядром, а с вычисляемым по данным. Такая работа в природе была, и хоть её представляли на ICLR 2019, всё равно она, кажется, несколько недооценена (https://arxiv.org/abs/1901.10430, Pay Less Attention with Lightweight and Dynamic Convolutions).



К локальному контексту добавляется глобальный для предопределённых входных позиций. В случае аналогичных берту задач классификации это позиция [CLS], или, например, позиции токенов вопроса для QA задач.



Соответственно в модель вводятся отдельные Q, K, V (если эти термины непонятны, то рекомендую лучшую статью по трансформеру, что я видел http://jalammar.github.io/illustrated-transformer/) для скользящего окна и для глобального внимания.



Полученные механизмы внимания скейлятся линейно относительно входа. Профит!



В этом месте есть инженерная проблема. Наивная реализация таких вариантов внимания слишком медленная, требуются кастомные ядра для CUDA. Это сделано с помощью Tensor Virtual Machine (TVM, https://tvm.apache.org/, https://arxiv.org/abs/1802.04799) с помощью которой можно описать функцию на сравнительно высокоуровневом питоноподобном языке (https://github.com/allenai/longformer/blob/master/longformer/diagonaled_mm_tvm.py#L52), а затем скомпилировать в целевую архитектуру, например, CUDA. Таким образом написали ядро, которое вполне сносно работает (но потенциал ускорения ещё есть).