Quá trình phát triển LLM phần 8 - Kỹ thuật self-attention có thể huấn luyện

• Trong Transformer chỉ dùng decoder kiểu GPT, self-attention có thể huấn luyện tính toán token nào trong đầu vào trước đó mà mỗi token cần chú ý tới để tạo ra vector ngữ cảnh
• Trọng tâm là scaled dot product attention, sử dụng ba ma trận có thể học Wq, Wk, Wv để đưa embedding đầu vào vào các không gian query, key, value
• Ma trận đầu vào X được biến đổi thành Q=XWq, K=XWk, V=XWv; sau đó Ω=QKᵀ được chia cho √c rồi áp dụng softmax theo từng hàng để thu được trọng số attention A
• Vector ngữ cảnh được tạo bằng một phép nhân ma trận duy nhất C=AV, và toàn bộ phép tính có thể áp dụng cho mọi token chỉ với 5 phép nhân ma trận và một phép chuyển vị
• Giai đoạn này vượt qua ví dụ đồ chơi dùng dot product trực tiếp giữa các embedding đầu vào, để tiến tới attention có thể huấn luyện có thể triển khai bằng PyTorch nn.Module và nn.Linear

---

Vị trí của self-attention trong luồng xử lý LLM

• LLM dựa trên Transformer chỉ dùng decoder kiểu GPT là kiến trúc dự đoán token tiếp theo dựa trên các token đã xuất hiện cho tới thời điểm hiện tại
• Luồng xử lý gồm tách chuỗi thành token, biến từng token thành token embedding, rồi cộng thêm positional embedding biểu diễn thông tin vị trí để tạo embedding đầu vào
• Self-attention tạo ra danh sách điểm attention biểu thị mức độ mỗi embedding đầu vào cần chú ý tới các token khác
  • Trong câu ví dụ "the fat cat sat on the mat", khi nhìn vào "cat" thì "fat" có thể là thông tin quan trọng
  • Khi nhìn vào "mat", mức độ quan trọng của "fat" có thể tương đối thấp hơn
• Các điểm attention đi qua softmax để trở thành trọng số attention có tổng bằng 1, và các embedding đầu vào được cộng có trọng số theo các giá trị này để tạo vector ngữ cảnh
• Vector ngữ cảnh được xem là vector biểu diễn ý nghĩa của từng token trong bối cảnh của toàn bộ đầu vào

Mục tiêu của self-attention có thể huấn luyện

• Ở các bước trước, mô hình dùng self-attention đồ chơi tính dot product trực tiếp giữa các embedding đầu vào
• Mục tiêu của bước này là xây dựng cơ chế attention có thể huấn luyện có khả năng tạo điểm attention từ các vector đầu vào
• Mục 3.4 của Sebastian Raschka trong Build a Large Language Model (from Scratch) triển khai điều này bằng scaled dot product attention
• Trọng tâm không phải là vì sao cấu trúc này hiệu quả, mà là nó hoạt động bằng những phép tính nào

Ma trận Query, Key, Value và phép chiếu không gian

• Gọi độ dài chuỗi đầu vào là n, số chiều embedding đầu vào là d, và số chiều vector ngữ cảnh là c
• Chuỗi embedding đầu vào được biểu diễn là x1, x2, x3, ... xn, và mỗi embedding đầu vào là một vector d chiều
• Định nghĩa ba ma trận trọng số có thể học
  • query weights matrix: Wq
  • key weights matrix: Wk
  • value weights matrix: Wv
• Mỗi ma trận có kích thước d×c và chiếu vector đầu vào d chiều sang không gian c chiều
• Phép tính đưa vector đầu vào xm vào không gian query là qm=xmWq
• Không gian key và value cũng chiếu embedding đầu vào sang các không gian c chiều khác nhau theo cùng cách đó

Cách nhìn ma trận như phép chiếu

• Ma trận có thể dùng cho biến đổi hình học, chẳng hạn như xoay một điểm
• Ma trận vuông thực hiện biến đổi trong cùng một số chiều, còn ma trận không vuông có thể đưa vector sang không gian có số chiều khác
• Ví dụ, ma trận 3×2 có thể biến các vector 3 chiều thành các vector 2 chiều
• Trong đồ họa 3D, ma trận frustum dùng để biến các điểm 3D thành các điểm trên màn hình 2D cũng là một ví dụ của kiểu phép chiếu này
• Self-attention đưa embedding đầu vào vào ba không gian chiếu khác nhau là query, key và value, rồi thực hiện tính toán trên các vector đã được chiếu
• Vì các ma trận chiếu này được học trong quá trình huấn luyện, nên xuất hiện một mức gián tiếp mà attention chỉ dùng dot product đơn giản không có

Tính toán điểm attention

• Khi xét một đầu vào cụ thể xm, điểm attention đối với một đầu vào khác xp được định nghĩa là dot product giữa phép chiếu query và phép chiếu key
• Công thức tính như sau
  • qm=xmWq
  • kp=xpWk
  • ωm,p=qm·kp
• Có thể xử lý phép tính này bằng vòng lặp cho mọi đầu vào, nhưng dùng phép nhân ma trận sẽ tính được tất cả cùng lúc
• Nếu đặt toàn bộ embedding đầu vào là ma trận X, thì X có kích thước n×d
• Ma trận key được tính một lần là K=XWk
  • Kết quả K có kích thước n×c
  • Mỗi hàng là vector thu được khi chiếu embedding đầu vào tương ứng sang không gian key
• Ma trận query cũng được tính tương tự bằng Q=XWq
• Dot product giữa mọi query và mọi key được lấy bằng QKᵀ
  • Q là n×c
  • Kᵀ là c×n
  • Kết quả Ω là n×n
• Ωm,p là điểm attention biểu thị mức độ cần chú ý tới xp khi tạo vector ngữ cảnh cho xm

Scaling và chuẩn hóa softmax

• Giống như ví dụ trước, các điểm attention được đưa qua softmax để biến thành các trọng số có tổng bằng 1
• Softmax làm các giá trị lớn nổi bật hơn và các giá trị nhỏ giảm đi, đồng thời điều chỉnh để tổng toàn bộ danh sách bằng 1
• Trong LLM thực tế, d và c có thể lên tới hàng nghìn, nên nếu chỉ dùng softmax thuần túy thì gradient nhỏ có thể xuất hiện
• Trong trường hợp đó, softmax có thể hoạt động “giống như step function”
  • Có thể hiểu là giá trị lớn nhất chi phối toàn bộ và các giá trị còn lại trở nên rất nhỏ
• Để giảm bớt hiện tượng này, các điểm attention được chia cho căn bậc hai của số chiều không gian chiếu c trước khi áp dụng softmax
• Biểu diễn ma trận như sau
  • A=softmax(Ω/√c, axis=1)
• axis=1 là ký pháp kiểu PyTorch, nghĩa là softmax được áp dụng theo từng hàng
• Kết quả A là các điểm attention đã được chuẩn hóa, tức ma trận trọng số attention

Tạo vector ngữ cảnh

• Phép chiếu sang không gian value được tính bằng V=XWv
• A là ma trận trọng số attention kích thước n×n
  • Am,p là trọng số attention áp dụng cho đầu vào p khi tạo vector ngữ cảnh của xm
• V có kích thước n×c, và mỗi hàng là vector thu được khi chiếu embedding đầu vào sang không gian value
• Ma trận vector ngữ cảnh được tính bằng C=AV
  • Kết quả C có kích thước n×c
  • Hàng thứ m của C là vector ngữ cảnh cho đầu vào xm
• Phép tính này thực hiện việc nhân các vector value với trọng số attention rồi cộng lại cho từng token chỉ bằng một phép nhân ma trận

Tóm tắt toàn bộ phép tính

• Ma trận đầu vào X chứa các embedding đầu vào của chuỗi token và có kích thước n×d
• Ba ma trận có thể học chiếu đầu vào lần lượt sang các không gian query, key, value
  • Q=XWq
  • K=XWk
  • V=XWv
• Điểm attention được tính từ dot product của query và key
  • Ω=QKᵀ
• Sau khi scaling điểm số, áp dụng softmax theo hàng để tạo trọng số attention
  • A=softmax(Ω/√c, axis=1)
• Nhân phép chiếu value với trọng số attention để tạo vector ngữ cảnh
  • C=AV
• Toàn bộ cơ chế self-attention có thể tạo vector ngữ cảnh cho mọi token đầu vào chỉ với 5 phép nhân ma trận và một phép chuyển vị

Cài đặt PyTorch và bước tiếp theo

• Mục 3.4 của cuốn sách triển khai các phép tính trên bằng mã PyTorch, đồng thời tạo một lớp con nn.Module đơn giản thực hiện cùng các phép toán ma trận đó
• Phiên bản đầu tiên dùng các đối tượng nn.Parameter thông thường cho ba ma trận trọng số
• Phiên bản thứ hai dùng nn.Linear để huấn luyện hiệu quả hơn
• Hai chủ đề sẽ được đề cập tiếp theo là
  • causal self-attention: cách không chú ý tới các token phía sau khi đang xét một token cụ thể
  • multi-head attention: được báo trước là không phức tạp như tưởng tượng ban đầu
• Xử lý batch vẫn là một vấn đề cần xem xét riêng
  • Ngay cả với một chuỗi đầu vào đơn lẻ, attention score matrix đã được sử dụng
  • Để xử lý song song nhiều chuỗi đầu vào, có thể cần tensor bậc cao hơn ma trận
• Bài tiếp theo nối tiếp tại Writing an LLM from scratch, part 9 -- causal attention