Longformer

Longformer là một transformer dành cho tài liệu dài, thay cơ chế self-attention toàn phần bằng attention cửa sổ trượt quanh mỗi token, kết hợp attention toàn cục trên một số token được chọn theo tác vụ; nhờ vậy, chi phí tăng tuyến tính theo độ dài đầu vào thay vì tăng theo bình phương. Iz Beltagy, Matthew Peters và Arman Cohan tại Allen Institute for AI giới thiệu model này trong bài báo năm 2020 Longformer: Transformer cho tài liệu dài. Các encoder như BERT và RoBERTa chỉ đọc tối đa 512 token, nên văn bản dài hơn phải được chia thành các đoạn, làm mất liên kết giữa chúng. Các model Longformer được phát hành có thể đọc 4.096 token trong một lượt, với cửa sổ ngữ cảnh dài gấp tám lần.

Cơ chế attention của Longformer hoạt động ra sao#

Mỗi lớp kết hợp các mẫu cục bộ và toàn cục trong cùng một cơ chế attention:

• Attention cửa sổ trượt: Mỗi token chú ý đến một cửa sổ cố định gồm các token lân cận, mỗi bên 256 token trong các model được phát hành.
• Cửa sổ trượt giãn: Khoảng cách giữa các vị trí được attention đến giúp một head nhìn xa hơn với cùng chi phí. Các tác giả dùng độ giãn trên một số head ở các lớp trên của model ngôn ngữ cấp ký tự.
• Attention toàn cục: Các token được chọn chú ý đến mọi vị trí và mọi vị trí chú ý đến chúng, thông qua các phép chiếu query, key và value riêng biệt. Tác vụ quyết định token nào được chọn: token [CLS] cho phân loại và các token câu hỏi cho tác vụ trả lời câu hỏi.

Xếp chồng các lớp làm tăng phạm vi của cửa sổ, tương tự cách các lớp tích chập xếp chồng mở rộng trường tiếp nhận của một mạng nơ-ron tích chập (CNN): l lớp với cửa sổ w bao phủ l × w token. Mỗi lớp tính khoảng n × w điểm attention thay vì n × n. Đếm các cặp trong mask Longformer sẽ cho thấy rõ sự khác biệt:

import torch

window = 512  # Longformer-base: 256 tokens on each side of every token

for n in (1024, 2048, 4096):
    pos = torch.arange(n)
    # Sliding window: each token attends to neighbors within window / 2 positions
    mask = (pos[:, None] - pos[None, :]).abs() <= window // 2
    # Global attention on token 0 ([CLS]): it sees every token and every token sees it
    mask[0, :] = True
    mask[:, 0] = True
    print(f"{n} tokens: full attention {n * n / 1e6:.1f}M scores, Longformer {mask.sum().item() / 1e6:.1f}M")

Lệnh này in ra 0.5M, 1.0M và 2.0M điểm Longformer, so với 1.0M, 4.2M và 16.8M của attention toàn phần: mỗi lần tăng gấp đôi đầu vào, lượng tính toán attention của Longformer tăng gấp đôi, còn attention toàn phần tăng gấp bốn. Các triển khai thực tế chỉ tính dải điểm, nhờ đó tiết kiệm bộ nhớ.

Huấn luyện và kết quả#

Longformer bắt đầu từ checkpoint RoBERTa. Các tác giả sao chép tám lần 512 embedding vị trí đã học để bao phủ 4.096 vị trí, rồi tiếp tục pretraining model ngôn ngữ được che trên các tài liệu dài. Sau khi fine-tune, Longformer luôn vượt RoBERTa trên các tác vụ tài liệu dài và đạt kết quả tiên tiến nhất vào thời điểm đó trên các benchmark hỏi đáp WikiHop và TriviaQA; trong khi đó, model ngôn ngữ cấp ký tự language model lập kết quả mới trên text8 và enwik8.

Longformer-Encoder-Decoder (LED) bổ sung decoder cho các tác vụ tóm tắt văn bản và sinh văn bản khác, được minh họa bằng việc tóm tắt toàn bộ bài báo khoa học từ arXiv. Hugging Face Transformers triển khai cả hai model và nhận các token toàn cục dưới dạng global_attention_mask. Tiếp đó là các phiên bản theo miền, chẳng hạn Clinical-Longformer, được tiền huấn luyện trên các kho ngữ liệu lâm sàng lớn và đánh giá trên 10 tác vụ, trong đó có nhận dạng thực thể có tên, trả lời câu hỏi và phân loại tài liệu.

Longformer và các model liên quan#

• Transformer-XL: Đọc văn bản dài theo từng đoạn và cho phép mỗi đoạn attention đến các trạng thái ẩn đã lưu trong bộ nhớ đệm của đoạn trước, mở rộng ngữ cảnh thông qua hồi quy để sinh văn bản từ trái sang phải.
• BigBird: Model của Google ra mắt cùng năm, bổ sung các liên kết attention ngẫu nhiên bên cạnh attention cửa sổ và toàn cục, đồng thời cũng đọc được 4.096 token.
• Attention thưa: Họ mô hình rộng hơn gồm các mẫu attention cố định mà Longformer thuộc về.
• Flash attention: Tính attention toàn phần chính xác bằng các tile bộ nhớ trên chip, nhờ đó bộ nhớ tăng tuyến tính theo độ dài trong khi số lượng điểm vẫn tăng theo bình phương.

Attention cục bộ và toàn cục trong thị giác#

Ảnh độ phân giải cao cũng gặp vấn đề tương tự, vì một transformer thị giác (ViT) chuyển ảnh thành hàng nghìn token patch. Swin Transformer giải quyết vấn đề bằng các cửa sổ cục bộ dịch chuyển; Segment Anything Model (SAM) của Meta, chạy trong package Ultralytics, kết hợp attention cửa sổ và toàn cục qua từng block để phân đoạn ảnh: encoder ảnh ViT-B attention trong các cửa sổ không chồng lấp kích thước 14 × 14 trên lưới patch 64 × 64 ở 8 trong 12 block, đồng thời attention toàn cục trên cả 4.096 patch ở các block 2, 5, 8 và 11.