Vì sao Vision Transformer (ViT) cần register token

• Token có norm cao trong feature map của ViT chủ yếu xuất hiện ở các vùng nền ít thông tin, và được diễn giải là hiện tượng patch token bị tái sử dụng như không gian lưu trữ cho tính toán nội bộ
• Cách giải quyết là thêm các register token tách biệt với patch ảnh vào chuỗi đầu vào, để tách riêng vai trò tính toán nội bộ sang các token chuyên dụng
• Ở cả ViT học có giám sát và tự giám sát, artifact giảm đi, feature map và attention map trở nên mượt hơn, có lợi cho các tác vụ thị giác downstream
• Có báo cáo cải thiện ở object discovery và dense prediction, nhưng với OpenCLIP thì hiệu năng object localization giảm nhẹ, nên cách diễn đạt “cải thiện trên mọi mô hình” đã được chỉnh sửa
• Khi thêm 4~16 register token sẽ phát sinh 2~6% chi phí tính toán, và việc tăng số lượng token không phải lúc nào cũng làm hiệu năng tác vụ dense tăng lên

---

Artifact lộ ra trong feature map của ViT

• Vision Transformer là công cụ mạnh trong học biểu diễn thị giác, nhưng artifact đã được quan sát trong feature map của cả ViT học có giám sát lẫn tự giám sát
• Artifact được quan sát dưới dạng token có norm cao xuất hiện trong quá trình suy luận
  • Chủ yếu xảy ra ở các vùng nền ít thông tin của ảnh
  • Có vẻ như mạng đang tái sử dụng các token đó cho tính toán nội bộ
• Các token này xuất phát từ những patch nền giống với các patch xung quanh, và được đánh giá là chứa ít thông tin patch cục bộ hơn nhưng mang nhiều thông tin ảnh toàn cục hơn
• Trong DINOv2 và các ViT khác, việc giảm độ phân giải không gian và độ chính xác của attention map có thể ảnh hưởng đến hiệu năng object discovery và localization ở downstream

Cách làm với register token

• Giải pháp được đề xuất là thêm các register token tách biệt với patch ảnh vào chuỗi đầu vào của Vision Transformer
• Register được mô tả là token tách khỏi patch ảnh, tương tự CLS token, nhưng không được dùng trực tiếp để tính loss
• Mục tiêu là để token chuyên dụng đảm nhận vai trò đó, thay vì để patch token hiện có bị dùng như “không gian lưu trữ” cho tính toán nội bộ
• Cách làm này có liên quan đến Memory Transformer trước đây, nhưng hướng tiếp cận phân tích artifact của feature rồi thêm token mới để sửa nó là điểm đặc thù trong phân tích của bài báo này

Thay đổi ở feature map và các tác vụ downstream

• Báo cáo cho thấy khi thêm register, artifact nói trên được loại bỏ ở cả mô hình học có giám sát và tự giám sát
• Feature map và attention map trở nên mượt hơn, giúp ích cho xử lý thị giác downstream
• Được tổng kết là đã đạt state of the art mới trên các dense visual prediction task của mô hình thị giác tự giám sát
• Nhờ đó cũng có thể áp dụng phương pháp object discovery cho các mô hình lớn hơn
• Trên DINOv2 và DeiT-III, hiệu năng object discovery tăng lên, còn hiệu năng ở các tác vụ mức ảnh nhìn chung được giữ nguyên

Ngoại lệ theo từng mô hình và phân tích bổ sung

• Với OpenCLIP, thêm register không cải thiện kết quả object localization, và theo các con số trong bảng thì còn có mức suy giảm hiệu năng nhẹ
  • Phản hồi của tác giả thừa nhận rằng cách diễn đạt “hiệu năng object discovery cải thiện trên mọi mô hình” không được các số liệu hậu thuẫn
  • Nội dung đó đã được sửa ở Section 3.3
• Ở OpenCLIP, Figure 7 cho thấy norm giảm, nhưng điều này không dẫn tới cải thiện object localization trong Table 3, và sự không nhất quán đó đã được thảo luận
  • Trong đánh giá LOST, với từng mô hình sẽ chọn embedding có hiệu năng tốt nhất trong keys, queries và values
  • Với CLIP, values là tốt nhất, còn artifact xuất hiện ở keys hoặc queries nhưng không xuất hiện ở values, nên điều này nhất quán với Table 3 và phân tích định tính
• Với MAE, có vẻ các “outlier patches” tương tự không xuất hiện
  • Có thể là vì MAE được huấn luyện chỉ với local loss nên không cần hoặc ít cần gom thông tin toàn cục
  • Độ chính xác phân loại ImageNet bằng linear probing của MAE-ViT-Large là 75%, thấp hơn các phương pháp SSL khác

Hiệu năng, chi phí và giới hạn

• Register là một phương pháp đơn giản, nhưng làm tăng chi phí tính toán do có thêm token
  • Khi thêm 4~16 register token, mức tăng chi phí tính toán được báo cáo là 2~6%
• Việc tăng số lượng register không phải lúc nào cũng dẫn đến cải thiện ở tác vụ dense
  • Khi tăng từ 0 lên 1 register, artifact biến mất và segmentation cùng depth prediction được cải thiện
  • Sau 1 register, dù thêm nhiều register hơn thì feature map đã sạch rồi nên hiệu năng segmentation và depth prediction không tăng nhiều
  • Hiệu năng classification có xu hướng được cải thiện thêm khi số lượng register tăng, nhưng hiện vẫn chưa có trực giác rõ ràng về lý do
• Trên bộ dữ liệu độ sâu NYU, khi tăng từ 8 lên 16 register thì RMSE tăng nhẹ
  • Có thể do khi độ dài chuỗi tăng, hyperparameter tối ưu cũng thay đổi
  • Chênh lệch RMSE 0.03 khi đi từ 8 lên 16 là nhỏ so với chênh lệch 0.1 khi đi từ 0 lên 1, nên cũng có thể chỉ là nhiễu đơn thuần

OpenReview chấp nhận và các điểm tranh luận trong review

• Quyết định trên OpenReview là Accept (oral)
• Meta review cho biết cả bốn reviewer đều chấm 8 điểm, và đóng góp kỹ thuật của bài báo được công nhận
• Các reviewer bày tỏ lo ngại về sự không nhất quán trong hiệu năng, mối quan hệ giữa dataset bias và mô hình SSL, lựa chọn thiết lập tối ưu hóa, cũng như ảnh hưởng của gradient clipping
• Trong quá trình phản biện, các tác giả đã bổ sung thêm thí nghiệm và phân tích trong appendix
• Các thảo luận sau đó cũng tăng cường phân tích cơ chế nội tại, chẳng hạn register có tiếp quản hành vi của outlier hay không, image token có trở lại mang tính local hơn hay không, và norm của CLS token thay đổi như thế nào

[GN⁺]

Ý kiến trên Hacker News

• Theo bài báo, register là các token học được bổ sung vào chuỗi đầu vào trong quá trình huấn luyện mô hình Vision Transformer
  Chúng được thêm vào sau lớp patch embedding như token [CLS] dưới dạng các giá trị có thể học được, và ở cuối Vision Transformer thì các token register bị loại bỏ, chỉ token [CLS] và các token patch được dùng làm biểu diễn ảnh
  Token này cung cấp cho mô hình một không gian để lưu trữ, xử lý và truy xuất thông tin toàn cục trong quá trình lan truyền thuận, nhờ đó không cần phải ép các token patch làm nhiệm vụ đó
  Kết quả là các artifact xuất hiện trong feature map của Vision Transformer đã huấn luyện và các outlier token có norm cao biến mất, feature map trở nên mượt hơn, đồng thời hiệu năng dense prediction và khả năng phát hiện đối tượng không giám sát cũng được cải thiện
  Chi phí suy luận chỉ tăng 2% nhưng cải thiện hiệu năng ViT khá đáng kể, nên đây có vẻ là một kết quả gọn gàng gần như bữa trưa miễn phí

  • Xu hướng cứ thêm token vào khắp nơi rồi mọi thứ hoạt động như có phép màu này trông khá đáng ngờ, và nếu còn lôi cả xác suất vào thì chắc những người làm toán rời rạc khó mà xem nhẹ được
    Biểu thức chính quy xác suất cũng đã là một chuyện, còn cái này là một cấp độ phóng túng toán học hoàn toàn khác
    Dù vậy, việc nó thực sự hoạt động vẫn rất đáng kinh ngạc
  • Cũng từng có thử nghiệm thêm nhiều token [CLS] vào BERT, nhưng kết quả kém ấn tượng hơn: https://arxiv.org/pdf/2210.05043
  • Tò mò không biết liệu Transformer văn bản thông thường có điều gì để học từ đây không
    Nghe khá giống attention sink hoặc pause token
  • Muốn dùng loại token này thì phải đưa kết quả của tầng cuối trở lại thành token trước, rồi embedding lại ở lượt tiếp theo
    Tò mò không biết có ai đã nghiên cứu cách chuyển thẳng đầu ra thô của một lượt thành đầu vào cho lượt kế tiếp chưa

• Có vẻ liên quan: "Let's Think Dot by Dot: Hidden Computation in Transformer Language Models" https://arxiv.org/abs/2404.15758
  Bài báo này cho rằng vẫn chưa rõ việc cải thiện hiệu năng chain-of-thought của mô hình ngôn ngữ đến từ kiểu phân rã nhiệm vụ giống con người, hay chỉ là từ lượng tính toán bổ sung mà các token thêm vào mang lại
  Họ cho thấy Transformer chỉ với các token đệm vô nghĩa, ví dụ ......, cũng có thể giải được các bài toán thuật toán khó mà trước đó không giải nổi khi phải trả lời không có token trung gian, nhưng để huấn luyện việc dùng các token đệm như vậy thì cần giám sát cụ thể và dày đặc
  Vì vậy, cũng tự hỏi liệu việc thêm các token tính toán/register đặc biệt vào LLM văn bản có đem lại lợi ích không
  Thảo luận thêm: https://news.ycombinator.com/item?id=40182695, https://www.reddit.com/r/LocalLLaMA/comments/1cf2w5a/transfo...

• Trong vài năm qua đã thử nghiệm hàng chục, có lẽ hơn 100 bài báo và ý tưởng trong lĩnh vực thị giác máy tính và nhận thức đa phương thức, và đây là một trong số hiếm hoi những trường hợp mà mọi thứ đều hoạt động tốt
  Ví dụ, mô hình này được huấn luyện cùng 4 token register và loss kiểu Matryoshka, đồng thời tạo ra embedding 64 chiều rất nhỏ có thể hữu ích cho những ai đang tìm lựa chọn thay thế CLIP: https://huggingface.co/unum-cloud/uform3-image-text-english-...

• Đã xem ở ICLR, và đây là một trong những bài báo hay nhất năm nay; điều đó cũng thể hiện rất rõ ở phiên poster
  Xin chúc mừng các tác giả

• Tôi đã dùng DinoV2 suốt vài tháng nay, và cũng đã thử mô hình có 4 token register cùng với [CLS] + token patch
  Vì nhiều token embedding từ mô hình cũ không có register đã là một phần của giải pháp hiện tại, nên token [CLS] giữa mô hình 0 register và mô hình 4 register không khớp nhau, vì thế tôi chưa áp dụng mô hình register mới
  Sẽ rất tốt nếu [CLS] và các token patch có thể được căn chỉnh theo cách nào đó giữa hai mô hình

• Trông giống attention sink cho mô hình thị giác