Giải thích Attention Residuals: Nghĩ lại về độ sâu của Transformer

Các LLM hiện đại ngày càng sâu, rộng và tiêu tốn nhiều tính toán hơn, tuy nhiên việc chồng thêm nhiều lớp Transformer không phải lúc nào cũng mang lại mức tăng tương xứng. Một lý do là các kết nối residual tiêu chuẩn gộp đầu ra các lớp với trọng số đơn vị cố định, khiến mỗi lớp thừa hưởng tổng đồng đều của mọi thứ đến trước nó. Điều này có thể làm loãng các biểu diễn sớm, khuếch đại độ lớn của trạng thái ẩn và khiến mạng khó tái sử dụng có chọn lọc những đặc trưng trung gian hữu ích nhất.

Thay vì coi độ sâu như một phép lặp cộng cố định, bài báo của nhóm Kimi về Attention Residuals cho phép mỗi lớp "attend" lên đầu ra các lớp trước đó bằng trọng số softmax học được. Trong bài viết này, chúng ta sẽ hiểu vì sao phép gộp residual tiêu chuẩn trở thành nút thắt, Attention Residuals hoạt động thế nào, vì sao biến thể theo khối quan trọng, và các kết quả thực nghiệm gợi ý gì về việc mở rộng độ sâu của mô hình ngôn ngữ.

Nếu bạn muốn tìm hiểu thêm về các ý tưởng đằng sau kiến trúc Transformer hiện đại, tôi khuyến nghị xem khóa học của DataCamp Hướng dẫn về Mô hình Transformer với PyTorch.

Vấn đề: Tín hiệu bị “chìm”

Kết nối residual là nền tảng của học sâu suốt nhiều năm. Trong Transformer, cập nhật residual tiêu chuẩn là:

hl=hl−1​+fl−1​(hl−1​)

Điều này giúp gradient lan truyền qua các mạng rất sâu. Nhưng residual không chỉ là mẹo cho gradient. Chúng còn xác định cách thông tin được gộp theo chiều sâu. Nếu khai triển phép lặp, ta có:

hl​=h1​+∑i=1l−1​fi​(hi​)

Điều này có nghĩa trạng thái ẩn tại lớp l chỉ là embedding cộng với tổng có trọng số đồng đều của mọi đầu ra lớp trước đó. Vậy nên, mỗi đóng góp thực chất bị gán cùng một trọng số.

Vấn đề xuất hiện ở quy mô lớn. Bài báo cho thấy trong kiến trúc PreNorm, việc cộng dồn residual không trọng số khiến độ lớn trạng thái ẩn tăng theo độ sâu, xấp xỉ O(L). 

Khi dòng residual phình to, các đầu ra lớp sớm bị loãng dần trong một tổng cộng dồn lớn hơn. Một khi tín hiệu từ lớp sớm đã hòa trộn vào tích lũy này, các lớp sâu hơn không thể phục hồi chọn lọc nó. Vì thế, chúng chỉ vận hành trên trạng thái đã gộp.

Điều này dẫn tới hiệu ứng mà bài báo gọi là “tín hiệu bị chìm”. Bằng chứng thực nghiệm mạnh mẽ cho sự kém hiệu quả này đến từ các nghiên cứu cắt tỉa lớp: một tỷ lệ đáng kể số lớp trong mô hình đã huấn luyện thường có thể bị loại bỏ mà ít ảnh hưởng đến hiệu năng. Điều này gợi ý rằng dù chúng ta tiếp tục tăng độ sâu, các mô hình thiếu cơ chế hiệu quả để tận dụng triệt để nó. Thay vì hình thành một chuỗi suy luận phân cấp, mạng hành xử giống như một đường chuyền thừa thãi, nơi tín hiệu sớm bị loãng dần.

Cũng có đánh đổi trong lượt truyền xuôi. Khi dòng residual tăng, các lớp sau có thể cần tạo đầu ra có độ lớn cao hơn để tác động đáng kể lên trạng thái tích lũy. Bài báo liên hệ việc này với hành vi PreNorm, nơi độ lớn trạng thái ẩn tăng đơn điệu theo độ sâu dưới phép gộp residual tiêu chuẩn.

Kết nối residual giúp gradient lan truyền, và gradient đối với một trạng thái ẩn trung gian có thể được viết như sau:

Thành phần đồng nhất ở đây bảo toàn đường đi gradient trực tiếp. Nhưng residual vẫn buộc đường gộp tiến tới phải coi mọi lớp trước có trọng số cố định 1.0. Đó là hạn chế cấu trúc mà AttnRes tìm cách khắc phục.

Từ residual cộng sang attention theo chiều sâu

Bước đi khái niệm thú vị nhất của bài báo là tính đối ngẫu thời gian–độ sâu. Kết nối residual nén thông tin theo chiều sâu, giống như mạng hồi quy nén thông tin theo thời gian. Trong mô hình chuỗi, attention thay thế hồi quy bằng cách cho phép mỗi vị trí truy cập chọn lọc các vị trí trước. AttnRes áp dụng chuyển đổi tương tự cho chiều sâu của mạng. Thay vì định nghĩa trạng thái ẩn kế tiếp là tổng cố định trên các lớp trước, AttnRes cho phép mỗi lớp "attend" lên đầu ra các lớp trước đó:

trong đó các trọng số αi→l là các trọng số attention softmax theo chiều sâu và có tổng bằng 1. Các trọng số đó được tính như sau:

với bài báo sử dụng:

Mỗi lớp có một vector giả-truy vấn học được wl, và truy vấn đó "attend" lên các khóa và giá trị được xây từ đầu ra các lớp trước. RMSNorm được áp dụng lên các khóa để các lớp có độ lớn đầu ra tự nhiên lớn hơn không chi phối softmax chỉ vì chúng có thang đo lớn hơn.

Chuyển dịch chủ chốt của bài báo là coi residual tiêu chuẩn như một dạng attention tuyến tính theo chiều sâu, trong khi AttnRes nâng cấp thành attention softmax theo chiều sâu. Thay vì cộng dồn đồng đều, ta có truy xuất chọn lọc theo chiều sâu.

Một chi tiết triển khai nhỏ nhưng quan trọng là khởi tạo. Tác giả gợi ý khởi tạo tất cả vector giả-truy vấn bằng 0, đảm bảo trọng số attention ban đầu là đồng đều, để mô hình bắt đầu huấn luyện như một phép trung bình trọng số bằng nhau thay vì một cơ chế attention thiên lệch ngẫu nhiên.

Hình 1: Tổng quan về Attention Residuals (bài báo Attention Residuals)

Attention Residuals đầy đủ (Full)

Trong biến thể Full Attention Residuals, mỗi lớp "attend" lên tất cả đầu ra của các lớp trước đó, cho phép gộp thông tin chọn lọc xuyên suốt độ sâu. Điều này mang lại linh hoạt tối đa, để lớp sâu hơn có thể nhấn mạnh lớp liền trước, embedding gốc, hoặc bất kỳ lớp sớm nào nếu đó là nơi chứa tín hiệu hữu ích nhất.

Trong huấn luyện tiêu chuẩn, chi phí bộ nhớ của Full AttnRes nhỏ hơn tưởng tượng ban đầu, vì nhiều đầu ra lớp vốn đã được giữ lại để lan truyền ngược. Nhưng ở huấn luyện quy mô lớn, bức tranh thay đổi. Khi bạn dùng tính toán lại kích hoạt và pipeline song song, các đầu ra sớm đó phải được lưu giữ và truyền đi tường minh để các lớp sau "attend" lên chúng, và điều đó trở nên tốn kém.

Vì vậy, Full AttnRes là cách tốt nhất để hiểu ý tưởng cốt lõi, nhưng không phải phiên bản mà hầu hết các nhóm muốn triển khai ở quy mô lớn.

Block Attention Residuals

Bài báo giới thiệu Block Attention Residuals để giảm chi phí bộ nhớ và truyền thông. Các lớp của mô hình được chia thành khối sao cho trong mỗi khối, đầu ra được kết hợp bằng cộng dồn residual tiêu chuẩn, nhưng giữa các khối, mô hình thực hiện attention trên các tóm tắt cấp khối thay vì mọi lớp trước riêng lẻ.

Nếu Bn là tập các lớp trong khối n, thì biểu diễn khối được cho bởi:

Sau đó, mô hình "attend" lên embedding b0=h1 các tóm tắt khối trước đó, và tổng từng phần của khối hiện tại khi tính toán tiến triển. Điều này giảm bộ nhớ và truyền thông từ O(Ld) xuống O(Nd), với N là số khối.

Bài báo báo cáo rằng dùng khoảng tám khối khôi phục phần lớn lợi ích của phiên bản đầy đủ, và khoảng cách hiệu năng giữa Full AttnRes và Block AttnRes thu hẹp khi mở rộng quy mô. Các kích thước khối như S=2,4,8 đều sát với phiên bản đầy đủ, trong khi nhóm quá thô sẽ xu hướng quay về hành vi nền tảng ban đầu.

Điểm rút ra hữu ích nhất là bạn không cần attention theo chiều sâu trên mọi lớp để đạt phần lớn lợi ích. Trong khi Block AttnRes khiến attention theo chiều sâu khả thi về tính toán, triển khai hiệu quả ở quy mô vẫn đòi hỏi thiết kế hệ thống cẩn trọng. Bài báo cũng giới thiệu một số tối ưu giúp AttnRes thực tiễn cho huấn luyện và suy luận trong thế giới thực.

Hình 2: Ví dụ giao tiếp pipeline dựa trên bộ nhớ đệm với 4 hạng vật lý và 2 giai đoạn ảo mỗi hạng, trong đó các hộp gạch chéo biểu thị kết thúc các khối AttnRes (bài báo Attention Residuals)

Dưới đây là ba tối ưu hệ thống then chốt giúp AttnRes có thể thay thế trực tiếp residual tiêu chuẩn:

• Tính toán hai pha: Bằng cách tách biệt giả-truy vấn khỏi trạng thái ẩn, ta có thể gom lô truy vấn cho tất cả lớp trong một khối. Pha 1 tính attention liên-khối song song, trong khi Pha 2 xử lý phụ thuộc nội-khối tuần tự bằng Online Softmax Merge. Điều này giảm đáng kể I/O trên mỗi lớp so với các phương pháp đa luồng trước đây như mHC.
• Bộ nhớ đệm xuyên-giai đoạn: Trong pipeline song song, chúng tôi loại bỏ truyền dữ liệu dư thừa bằng cách đệm cục bộ các khối đã nhận trước đó. Điều này mang lại cải thiện V× (với V là số giai đoạn ảo như trong Hình 2) về chi phí đỉnh mỗi lần chuyển, cho phép truyền thông được chồng lấp hoàn toàn với tính toán.
• Tiền điền hiệu quả bộ nhớ: Bằng cách chia nhỏ biểu diễn khối trên các thiết bị song song theo tensor, dấu chân bộ nhớ cho chuỗi ngữ cảnh dài (128K+) giảm theo bậc độ lớn (ví dụ, từ 15GB xuống 0,3GB mỗi thiết bị).

Kết quả và phân tích hiệu năng

Bài báo kiểm chứng AttnRes ở nhiều cấp độ, như quy luật mở rộng, cắt bỏ thành phần, động lực huấn luyện và các benchmark hạ nguồn.

Hình 3: Đường cong quy luật mở rộng cho Attention Residuals. (bài báo Attention Residuals)

Tuy nhiên, kết quả nổi bật đến từ các thí nghiệm quy luật mở rộng. Trên năm kích thước mô hình, cả Full AttnRes và Block AttnRes đều đạt loss xác thực thấp hơn so với đường cơ sở. Dựa trên các đường cong khớp, Block AttnRes đạt cùng mức loss như mô hình nền tảng cần khoảng 1,25× tài nguyên tính toán. Điều này cho thấy AttnRes không chỉ tốt hơn về lý thuyết mà còn hiệu quả tính toán hơn.

Ở các tác vụ hạ nguồn, AttnRes cải thiện so với đường cơ sở trên mọi benchmark được đánh giá. Một số mức tăng lớn nhất nằm ở các tác vụ đòi hỏi suy luận nhiều bước, như thể hiện trong Bảng 1:

Bảng 1: So sánh hiệu năng của AttnRes với đường cơ sở, cả hai đều theo cùng công thức tiền huấn luyện (bài báo Attention Residuals)

Lợi ích đặc biệt mạnh ở suy luận nhiều bước, toán và mã, phù hợp với giả thuyết cốt lõi của bài báo rằng nếu các lớp sau có thể truy xuất chọn lọc các biểu diễn sớm thay vì thừa hưởng một tổng mờ, suy luận phân hợp sẽ được cải thiện.

So với đường cơ sở, AttnRes cho thấy loss xác thực thấp hơn trong suốt quá trình huấn luyện, độ lớn đầu ra được giới hạn hơn theo chiều sâu, và độ lớn gradient đồng đều hơn giữa các lớp. Đường cơ sở chịu mô hình loãng thường thấy của PreNorm, nơi độ lớn trạng thái ẩn tăng đơn điệu theo độ sâu, trong khi Block AttnRes tạo ra mẫu hình được kiểm soát và giới hạn hơn nhờ gộp chọn lọc tại ranh giới khối.

Cũng có một phát hiện cấp độ kiến trúc rộng hơn rằng, trong quét tham số cố định và tính toán cố định, cấu hình tối ưu cho AttnRes nghiêng về mô hình sâu hơn và hẹp hơn so với đường cơ sở. Điều này gợi ý AttnRes có thể khiến độ sâu bổ sung hữu ích hơn so với dưới phép gộp residual tiêu chuẩn.

Kết luận

Residual thường được xem như một mẹo huấn luyện cần thiết, tức là một lối tắt đồng nhất giữ cho gradient “sống”. AttnRes cho rằng cách nhìn này quá hẹp. Các đường residual cũng là cơ chế định tuyến thông tin theo chiều sâu, và phép cộng cố định có thể quá thô sơ cho các mô hình rất sâu.

Attention Residuals không hẳn là vá một vấn đề đã biết mà là nâng cấp một lựa chọn thiết kế bị bỏ quên. Mô hình chuỗi đã tiến hóa từ hồi quy sang attention vì hồi quy cố định quá hạn chế. AttnRes lập luận rằng chiều sâu giờ có thể sẵn sàng cho chuyển đổi tương tự.

Mô hình vẫn thể hiện tính cục bộ mạnh, với các lớp thường "attend" nhiều nhất tới các tiền nhiệm gần kề, nhưng nó cũng học các mẫu bỏ qua không tầm thường, duy trì trọng số bền vững trên embedding, và giữ hành vi khác biệt giữa lớp trước attention và trước MLP. 

Điểm rút ra thực tiễn không phải là mọi mô hình nên thay residual bằng attention theo chiều sâu đầy đủ. Thay vào đó, nó tái khung hoá phép gộp residual như một không gian thiết kế kiến trúc cốt lõi.