AI說書 - 從0開始 - 10

我想要一天分享一點「LLM從底層堆疊的技術」，並且每篇文章長度控制在三分鐘以內，讓大家不會壓力太大，但是又能夠每天成長一點。

• Attention Layer的複雜度是O(n^2 * d)
• Recurrent Layer的複雜度是O(d^2 * n)

這些資料出自於ChatGPT的關鍵技術 - Transformer的原始Google論文：Attention is All You Need, Vaswani et al. (2017)

並且我們在AI說書 - 從0開始 - 9，已經完成Colab Python GPU環境配置。現在目標是想要用Python來做模擬，且使用GPU當作運算資源。

針對Attention Layer的程式配置為：

import torch
import time

n = 512
d = 512

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

input_seq = torch.rand(n, d, device = device)

start_time = time.time()
_ = torch.mm(input_seq, input_seq.t())
attention_time = time.time() - start_time
print(f"Self-attention computation time: {attention_time} seconds")

對此，執行結果為：

而針對Recurrent Layer的程式配置為：

start_time = time.time()
hidden_state = torch.zeros(d, device = device)
for i in range(n):
	for j in range(d):
		for k in range(d):
			hidden_state[j] += input_seq[i, j] * hidden_state[k]
			compute_time = time.time() - start_time
			if compute_time > attention_time * 10:
				break

recurrent_time = time.time() - start_time
print(f"Recurrent layer computation time: {recurrent_time} seconds")

對此，執行結果為：