Reformer

Reformer là một transformer tiết kiệm bộ nhớ do Nikita Kitaev, Łukasz Kaiser và Anselm Levskaya thuộc UC Berkeley và Google Research giới thiệu trong bài báo Reformer: Transformer hiệu quả tại ICLR 2020. Model này giải quyết hai loại chi phí giới hạn độ dài chuỗi. Cơ chế attention dùng hashing nhạy địa phương (LSH) thay thế ma trận attention đầy đủ kích thước n × n, với chi phí khoảng L log L cho chuỗi dài L; các lớp residual thuận nghịch lưu trữ các activation cần cho lan truyền ngược một lần cho toàn bộ mạng thay vì một lần cho mỗi lớp. Google cho biết Reformer xử lý được cửa sổ ngữ cảnh lên đến 1 triệu từ trên một accelerator duy nhất có bộ nhớ 16 GB, đồng thời xử lý toàn bộ tiểu thuyết Crime and Punishment như một mẫu huấn luyện duy nhất.

Attention LSH#

Sau softmax, attention của mỗi query chủ yếu tập trung vào một vài key giống nó nhất, vì vậy Reformer chỉ tìm những key đó. Model sử dụng một phép chiếu chung cho query và key, đồng thời băm mọi vector bằng hashing nhạy địa phương theo góc: nhân với ma trận ngẫu nhiên R rồi lấy bucket tại vị trí có giá trị lớn nhất trong [xR, −xR], nhờ đó các vector có hướng tương tự thường rơi vào cùng bucket. Sau đó, các token được sắp xếp theo bucket, chuỗi đã sắp xếp được chia thành các đoạn bằng nhau, và mỗi token chỉ attention trong đoạn của nó cùng đoạn liền trước. Nhiều vòng hashing chạy song song để bắt được các cặp tương tự mà một vòng duy nhất có thể tách rời.

LSH là phương pháp kinh điển cho tìm kiếm vector láng giềng gần đúng: phương pháp này tìm được phần lớn k láng giềng gần nhất của query với chi phí thấp hơn nhiều so với tìm kiếm chính xác. Bước phân nhóm vào bucket đủ ngắn để chạy độc lập:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

Vector gần, có độ tương đồng cosine 0,96 với query, nằm cùng bucket trong 72% trên 1.000 vòng; vector không liên quan, có độ tương đồng 0,04, nằm cùng bucket trong 9% số vòng. Vì vậy, attention giới hạn trong một bucket vẫn giữ lại hầu hết các cặp quan trọng.

Các lớp thuận nghịch và chia đoạn#

Một mạng nơ-ron tiêu chuẩn giữ lại activation của mọi lớp cho lượt lan truyền ngược, nên bộ nhớ tăng theo cả độ sâu lẫn độ dài. Reformer điều chỉnh mạng residual thuận nghịch (RevNet): chia activation thành hai nửa và cập nhật mỗi nửa bằng một hàm của nửa còn lại.

Trong lượt lan truyền ngược, mỗi lớp khôi phục đầu vào từ đầu ra, x2 = y2 − FeedForward(y1), rồi x1 = y1 − Attention(x2), nhờ đó chỉ cần lưu activation cuối cùng và tính toán lại mỗi lớp một lần. Gradient checkpointing thực hiện một sự đánh đổi tương tự bằng cách lưu activation tại một vài checkpoint và tính toán lại các lớp ở giữa. Reformer cũng chạy các lớp feed-forward theo từng đoạn vị trí, vì lớp ẩn của chúng rộng hơn model nhiều lần. Bộ nhớ activation là yếu tố giới hạn batch size khi huấn luyện Ultralytics YOLO26: với batch=-1, AutoBatch đo bộ nhớ GPU ở một số batch size, khớp đường thẳng và chọn batch được dự đoán sẽ sử dụng 60% bộ nhớ còn trống.

Kết quả và các bản triển khai#

Reformer đạt kết quả ngang với transformer tiêu chuẩn trên văn bản cấp ký tự từ enwik8 và tác vụ tạo ảnh từng pixel trên ImageNet-64, trong khi dùng ít bộ nhớ hơn nhiều và chạy nhanh hơn trên các chuỗi dài. Hugging Face Transformers triển khai model này dưới dạng ReformerModel, kèm các checkpoint đã phát hành được huấn luyện trên enwik8 và Crime and Punishment.

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

• Longformer: Xác định trước mẫu attention: cửa sổ trượt quanh mỗi token cộng thêm một vài token toàn cục.
• Attention thưa: Nhóm phương pháp chỉ chấm điểm một số cặp token. Attention LSH là một phương pháp thuộc nhóm này, dựa trên nội dung.
• Attention tuyến tính: Thay softmax bằng ánh xạ đặc trưng kernel để có thể tính attention trong thời gian tuyến tính.
• Flash attention: Tính attention chính xác trên các tile bộ nhớ on-chip của GPU mà không lưu ma trận điểm số, nhờ đó bộ nhớ tăng tuyến tính trong khi số lượng điểm số vẫn tăng bậc hai.