Một chút về Transformers (Phần 3): What the hell is Residual Stream

Nếu các bạn chưa tự tin về Transformers thì hãy đọc qua phần 1 và phần 2 nhá. Phần 1 ở đây, phần 2 ở đây.

Toàn bộ bài viết này là note của mình cho bài báo [] (cũng là bài báo favorite của mình).

1. Một góc nhìn tổng quát

Một mô hình transformers sẽ bắt đầu token embedding, đưa một token $t$ (one-hot vector) thành một vector embedding $x_0$ thông qua embedding matrix $W_E$, tiếp sau đó là một chuỗi các residual blocks, và cuối cùng là đến token unembedding, đưa một vector embedding $x_{-1}$ (layer cuối) sang vector logits $T(t)$ thông qua unembedding matrix $W_U$.

Mỗi residual block (hay transformers block) sẽ bao gồm một attention layer và sau đó là một MLP layer. Cả attention layer và MLP layer read (đọc) input từ residual stream (đường thẳng từ $x_i$ đến $x_{i+1}$ và tương tự $x_{i+1}$ đến $x_{i+2}$ như hình) và sau đó write (ghi) kết quả của chúng (các layer) vào residual stream bằng cách dùng residual connection $x_{i+1} = x_i + \text{layer}(x_i)$.

2. Residual Stream là một kênh liên lạc giữa các component Transformers với nhau

Nếu ta xem Transformers như một cái máy tính siêu phức tạp thì Residual stream là một trong những tính năng chính của Transformers (thật ra residual stream đã xuất hiện trong những model xưa hơn ví dụ cell memory của LSTM hay ResNet). Residual stream chỉ đơn giản là tổng của ouput layer trước đó (ví dụ là attention layer hoặc MLP layer) với embedding ban đầu (ví dụ là vector embedding $x_i$).

Residual stream có linear structure (cấu trúc tuyến tính) (khác với ResNet, ở mỗi residual connection của ResNet đều sử dụng thêm activation ví dụ như ReLU). Mỗi layer trong residual block thực hiện các linear transformation (ánh xạ tuyến tính) bất kỳ để read input từ residual stream. Ngoài ra, trước khi write outputs của layer ấy vào residual stream, layer cũng thực hiện các linear transformation.

Một trong những hệ quả của tính linear, addtive (tuyến tính và cộng) của residual stream chính là residual stream không có cơ sở ưu tiên (privileged basis) ¹. Vì vậy, ta có thể xoay (rotate) các ma trận mà tương tác với residual stream mà không làm thay đổi hành vi (behaviors) của model.

a. Cơ sở ưu tiên

Dựa theo định nghĩa của bài báo gốc ¹:

A privileged basis occurs when some aspect of a model’s architecture encourages neural network features to align with basis dimensions, for example because of a sparse activation function such as ReLU.

Bây giờ ta sẽ giải thích từng câu để làm rõ cơ sở ưu tiên:

Để làm rõ hơn, giờ ta lấy ví dụ một neural network (MLP) đơn giản được huấn luyện để nhận dạng các hình dạng trong ảnh. Ta sẽ tập trung vào một hidden layer, gọi là Layer $L$.

• Layer L có $N=4$ neuron. Vì vậy, không gian kích hoạt (activation space) của có 4 chiều.
• Các vector cơ sở cho không gian này là các cơ sở chuẩn:
  • $\mathbf{b}_1 = [1, 0, 0, 0]$ (trục của Neuron 1)
  • $\mathbf{b}_2 = [0, 1, 0, 0]$ (trục của Neuron 2)
  • $\mathbf{b}_3 = [0, 0, 1, 0]$ (trục của Neuron 3)
  • $\mathbf{b}_4 = [0, 0, 0, 1]$ (trục của Neuron 4)

Mục tiêu của chúng ta là phát hiện ba hình dạng cơ bản: hình vuông, hình tam giác và hình tròn. Sau khi huấn luyện, mạng có thể học được rằng:

• (1) Không có ReLU
  • Một hình vuông được biểu diễn bởi vector [1.2, -0.9, 0.8, 1.1]. Đây là một dense vector (ít các phần tử bằng $0$) chỉ theo một hướng bất kỳ trong không gian 4D.
  • Một hình tam giác được biểu diễn bởi [-0.8, -1.1, 1.3, -0.9].
  • Một hình tròn được biểu diễn bởi [1.4, 1.2, 0.1, -0.7]
  • Có thể thấy, nếu không có ReLU, mỗi feature là một tổ hợp tuyến tính cụ thể và phức tạp của cả bốn neuron. Ta không thể chỉ vào một neuron duy nhất và nói “đây là neuron chuyên phát hiện hình vuông”. Các feature không thẳng hàng (align) ² với các vector cơ sở. Đây là một cơ sở không ưu tiên.
• (2) Có sự xuất hiện của ReLU
  • Một hình vuông được biểu diễn bởi vector [1.2, -0.9, 0.8, 1.1] lúc này trở thành [1.5, 0, 0, 0.1]. Có thể thấy vector này sparse (thưa) và đặc biệt hơn là vector này có thể được xấp xỉ bằng $\alpha \mathbf{b}_1$ với $\alpha = 1.5$, do đó feature hình vuông này được “align” hoàn toàn với neuron đầu tiên.
  • Tương tự với hình tròn và hình tam giác, ta có các vector sau ReLU lần lượt là [1.4, 1.2, 0.1, 0] và [0, 0, 1.3, 0]. Feature hình tam giác hoàn toàn align với neuron 3. Trong khi đó feature hình tròn lại phức tạp hơn khi cần 2 neuron để biểu diễn là neuron 1 và neuron 2 (điều này còn cho thấy, neural network xem feature hình tròn rất gần với feature hình vuông).
• Có thể thấy nhờ sự xuất hiện của ReLU (hay theo định nghĩa gốc, “occurs when some aspect of a model’s architecture encourages neural network features to align with basis dimensions”), các feature đã được align với từng neuron, vì vậy ta có thể “tự tin” chỉ tay vào từng neuron và bảo rằng neuron này trách nhiệm cho phát hiện hình vuông, … Điều này làm cho biểu diễn có thể diễn giải được, và đó là những gì ta gọi là một cơ sở ưu tiên (privileged basis). Một điều nữa đó là tính không tuyến tính (non-linear) của ReLU và cả tính sparse của vector kích hoạt (sau ReLU).

b. Xoay cơ sở

Chúng ta xoay cơ sở để chúng ta, con người, có thể hiểu được mô hình.

Và một điểm đặc biệt của Residual Stream đó chính là xoay cơ sở không làm thay đổi hành vi (hay chính xác là output) của mô hình. Tại sao lại vậy ?

• Xét một ma trận xoay $R$, đầu tiên $R$ là một orthogonal matrix (ma trận trực giao), tức là $R^{-1} = R^T$ hay $R^TR = RR^T = I$. Giả sử vector embedding $x_i$ nằm trên residual stream được xoay bằng ma trận xoay $R$, khi đó ta có $x'_i = R x_i$.
• Mỗi component (attention, mlp, …) nếu muốn tương tác (read, write) với residual stream thì cần phải xoay ma trận chiếu tương ứng của nó.
• Xét Attention Component:
  • Input từ residual stream sẽ được read thông qua 3 ma trận $W_Q, W_K, W_V$, do $x_i$ đã rotate thông qua ma trận $R$, để không làm thay đổi behavior, ta cần $W'_Q, W'_K, W'_V$ mới sao cho:
  $$W_Q x_{i} = W'_{Q} x'_{i}; \quad W_{K} x_{i} = W'_{K} x'_{i}; \quad W_{V}x_{i} = W'_{V}x'_{i}$$
  • Có thể thấy $W'_Qx'_i = W'_{Q}Rx_{i} = (W'_{Q}R)x_{i} \implies W'_{Q}R = W_{Q} \implies W'_{Q} = W_{Q}R^T$. Vậy, khi mà input ở residual stream bị xoay bởi ma trận $R$ thì các phép chiếu tuyến tính (hay chính là nhân ma trận) trong component sẽ “un-rotate” $x'_i$ thông qua ma trận $R^T$, tức là $R^Tx'_i = (R^TR)x_i = x_i$. Điều này cho thấy, các cách tính toán ($W_Q, W_K, W_V$) bên trong component vẫn giữ nguyên không thay đổi.
  • Tiếp theo ở operation write, attention layer sẽ write vào residual stream thông qua công thức sau:
  $$x'_{i+1} = x'_{i} + W'_{O} \text{head}_{i}$$
  • Do các giá trị trên residual stream đều rotate bởi $R$, vì vậy:
  $$\begin{aligned} Rx_{i+1} &= Rx_{i} + W'_{O} \text{head}_{i} \\ \implies R(x_{i} + W_{O}\text{head}_{i}) &= Rx_{i} + W'_{O}\text{head}_{i} \\ \implies RW_{O}\text{head}_{i} &= W'_{O}\text{head}_{i} \end{aligned}$$
  • Có thể thấy nếu input xoay bởi ma trận $R$, thì cùng lúc đó ma trận output cũng xoay bởi ma trận $R$ ⁴.
• Với residual stream, ta có thể thấy việc rotate input thông qua ma trận $R$ bất kỳ sẽ không làm ảnh hưởng đến các việc tính toán của các component (các component vẫn tính toán bằng đúng các tính toán trước khi rotate, ví dụ như: $W_{Q}x_{i} = W_Q(R^TR)x_i = W'_Qx'_i$). Vì vậy residual stream không phụ thuộc vào 1 basis bất kì nào (việc rotate sẽ thay đổi basis => ta nói residual stream là basis-free).

3. Virtual Weights (Trọng số “ảo”)

Xét component (kí hiệu là $C$) thứ $1$, ta có $x_2 = x_1 + W_O^1 C_1(W_I^1x_{1})$, đặt $y_i$ là output của component thứ $i$ (trước khi chiếu qua trọng số output), do đó $x_2 = x_1 + W_O^1 y_1$. Tổng quát lên, với mỗi component thứ $j$, ta có:

Khi đó, ở component tiếp theo $j+2$, component này sẽ read information từ residual stream, do đó input của $j+1$ sẽ là:

Có thể thấy, ở component layer sau, component này read information ở layer phía trước thông qua virtual weight $W_I^{j+1}W_O^{j}$ (là tích của ma trận input layer hiện tại và ma trận output layer phía trước).

Tương tự, ở component tiếp theo, ta cũng có:

Tổng quát lên, một component $i$ read information từ component $j$ (với $j < i$) thông qua virtual weight $W_I^iW_O^j$. Virutal weight chính là thứ mà connect được output của component trước đó với input của component phía sau.

4. Không gian con và “thông lượng” (bandwitdh) của Residual Stream

Residual stream là một không gian nhiều chiều (thông thường, một mô hình nhỏ, ví dụ như bài báo gốc thì $d = 512$, hay những model lớn hơn thì $d$ có thể hàng ngàn như ở Gemma-2-2B, $d=2304$). Điều này nghĩa là các layers gửi information (send information) của mình đến các layer khác bằng việc lưu trữ information trong một không gian con khác (ở đây residual stream là một không gian có $d$ chiều, trong đó sẽ có nhiều không gian khác nhau và các layers sẽ store information của mình trên những không gian con ấy).

Ở attention layer (thông thường là multi-head attetion), khi đó mỗi attention head sẽ có chiều là $d / h$ với $h$ là số lượng attention head (ở paper gốc, attention head có chiều là $64$), có thể thấy chiều của attention rất nhỏ so với residual stream. Do đó khi các attention head write (thông qua concatenation và linear projection) thì các information sẽ nằm trên một phần (hay một không gian con) của không gian vector residual stream, và các không gian con của các head ấy có thể disjoint với nhau ⁵.

Khi mà information được thêm (added) vào residual stream, nó sẽ nằm đó (persist) cho đến khi một layer khác thay đổi nó (bằng cách overwrite ví dụ như delete, add more, …). Từ góc nhìn này, ta có thể xem residual stream chính là memory hay bandwitdth, vì vậy (về mặt lý thuyết) việc tăng $d_{model}$ lên thì ta có thể lưu trữ được nhiều information từ các component hơn.

Một điều nữa là token embedding hay token unembedding thông thường chỉ tương tác (interact) với một phần nhỏ các chiều trong residual stream ⁶, do đó còn rất nhiều các chiều trống trong residual stream (nói cách khác là những phân vùng trống trong bộ nhớ) cho các layer còn lại lưu trữ thông tin của mình.

Ngoài ra, residual stream in very high demand (bởi vì đây là nơi duy nhất mà các layer có thể communicate với nhau, ngoài ra các layer liên tục read and write vào residual stream). Tuy nhiên, số chiều tính toán (computational dimensions) thường lớn hơn rất nhiều so với số chiều của residual stream (nghĩa là thông tin mà các component có thường rất lớn so với khả năng lưu trữ của residual stream memory).

• Ví dụ, residual stream chính là một bottleneck activation, các MLP ở các layer khác nhau (thông thường có activation nhiều chiều hơn residual stream) muốn comunicate với nhau phải thông qua residual stream (do đó residual stream nằm giữa 2 layer MLP có (rất) nhiều chiều hơn so với mình). Ngoài ra, residual stream chính là con đường duy nhất để một layer MLP bất kỳ communicate với các layer MLP phía sau, hơn nữa, residual stream còn phải lưu trữ các information đến từ các lớp MLP khác trên đường dẫn đến extreme bottleneck.
• Tương tự, một value vector (trong $Q, K, V$ của attention head) cũng là một bottleneck activation.
  • Đầu tiên, value vector sẽ có chiều là $d_{model} / h$ với $h$ là số attention head do đó có chiều nhỏ hơn rất nhiều so với residual stream.
  • Đặt $x_s$ là residual stream tại token ở position $s$, do đó value vector của $s$ sẽ là $v_s = x_s W_V$. Khi đó, value vector $v_s$ được dùng để update residual stream ở position $t$, tức là:
  $$x_{t} = x_{s} + (\text{attention\_score} \times v_{s} )W_{O}$$
  • Có thể thấy, information từ residual stream $x_s$ sẽ compressed (nén vào) vào $v_s$ sau đó được chuyển sang residual stream ở token position mới là $t$, vì vậy giữa 2 residual stream, value vector chính là bottleneck activation, ngoài ra có thể thấy value vector là cách duy nhất để move information từ token này sang token khác.

Bởi vì tính chất high demand của residual stream bandwidth, ta có thể thấy các một vài MLP neurons hoặc các attention head thực hiện chức năng memory management ví dụ như clearing residual stream dimensions by other layer bằng cách đọc information và write negative của information đó (giống như việc có một app thực hiện dọn bộ nhớ, dọn ở đây chính là ghi vào một thông tin “negative” của thông tin cần xoá, vì vậy triệt tiêu nhau).

5. Các Attention Head hoạt động độc lập với nhau

Có thể thấy, attention layers của Transformers bao gồm nhiều attention heads $h \in H$ hoạt động độc lập và song song với nhau và sau đó, mỗi attention head sẽ thêm output của mình vào residual stream.

Nhắc lại Attention:

Giả sử, $r^{h_{1}}$ là output của attention head thứ $1$, tương tự ta có $r^{h_n}$ là output của attention head thứ $n$ (có chiều là $[1 \times d_v]$ và là vector dòng). Theo paper gốc, khi hoàn thành output, ta concat các attention head lại với nhau và nhân với output matrix $W_O$. Đặt $W_O^i$ là component (ma trận con) thứ $i$ của $W_O$ có chiều là $[d_v \times d_{model}]$, tức là sẽ gồm $n$ component của $W_O$ và các component được vstack với nhau, ta có:

Có thể thấy việc tính toán output cuối cùng tương tự như việc tính từng attention head một cách độc lập với nhau sau đó nhân với từng output matrix tương ứng (là các component của output matrix tổng), cuối cùng cộng lại rồi được kết quả cuối và lấy kết quả cuối cùng ấy đưa vào residual stream.

6. Attentions Heads chính là nơi để thông tin “di chuyển”

Câu hỏi: Nếu các Attention Head hoạt động độc lập với nhau, chúng thực sự làm gì ?

Hoạt động cơ bản nhất của attention heads đó chính là moving information. Các attention head read information từ residual stream của một token này và sau đó write information ấy vào residual stream của token khác.

The main observation to take away from this section is that which tokens to move information from is completely separable from what information is “read” to be moved and how it is “written” to the destination (note: mình không hiểu câu này lắm).

Để nhìn rõ hơn, ta nhìn kĩ vào cách tính toán của Attention Head:

• Đặt $x_i$ là residual stream value tại token $i$ (ở đây mình sẽ xem các vector này vector cột giống như paper thay vì vector dòng), khi đó value vector sẽ là $v_i = W_Vx_i$.
• Khi đó giá trị head (hay result vector trước khi nhân với output matrix) là $r_i = \sum_{j} A_{i,j}v_{j}$ (trong đó $A_{i, j}$ là giá trị softmax của score giữa hai token $i$ và $j$).
• Cuối cùng, tính giá trị output vector $h(x)_i = W_{O}r_{i}$ (ở đây ta giả sử chỉ có 1 attention head duy nhất, nếu là $n$ attention head thì $h(x)_i = \sum_{k=1}^n W_{O}^k r^k_{i}$).

• Đặt $x$ là ma trận gồm các embedding vector, mỗi embedding được xếp thành cột, mục đích của ta là nhân $W_V$ cho từng cột của $X$, vì vậy, dùng tensor product (ta hiểu ở đây là Kronecker product), ta có:

• Tương tự, sau đó đưa qua $A$, ta có:

• Và cuối cùng,

• Viết lại gọn hơn, ta có: