Tự xây dựng tầng nhúng Embedding Lookup Table từ ma trận trọng số và Index Slicing

Mô tả bài toán

Trong các mô hình ngôn ngữ lớn (LLM), tầng Embedding (nn.Embedding) thực chất là một bảng tra cứu (lookup table) ma trận trọng số W ∈ R^{V × D}, trong đó V là kích thước từ điển (vocab_size) và D là số chiều vector nhúng (embedding_dim).

Mỗi token ID nguyên i ∈ [0, V-1] được chuyển đổi thành vector hàng thứ i của ma trận W.

Yêu cầu kỹ thuật:

Xây dựng lớp EmbeddingLayer:

• __init__(self, vocab_size: int, embedding_dim: int, weights: np.ndarray | None = None):
• Khởi tạo ma trận self.weight shape (V, D). Nếu weights được truyền vào, gán trực tiếp (sao chép weights.copy()). Nếu không, khởi tạo bằng phân phối chuẩn tắc N(0, 0.02).
• forward(self, indices: np.ndarray) -> np.ndarray:
• Nhận mảng indices kiểu số nguyên 1D (L,) hoặc 2D (B, L).
• Kiểm tra tính hợp lệ: nếu có bất kỳ chỉ số nào < 0 hoặc ≥ V, ném ngoại lệ ValueError("Token index out of vocabulary bounds").
• Trả về tensor embeddings tương ứng: shape (L, D) nếu đầu vào 1D, shape (B, L, D) nếu đầu vào 2D.
• backward(self, grad_output: np.ndarray, indices: np.ndarray) -> np.ndarray:
• Nhận grad_output có cùng shape với kết quả forward.
• Tính đạo hàm đối với ma trận trọng số ∂ L∂ W có shape (V, D).
• Lưu ý then chốt: Nếu trong indices có các token trùng lặp (ví dụ token 5 xuất hiện 3 lần trong batch), gradient tương ứng phải được cộng dồn (accumulated) chứ không được ghi đè! Sử dụng np.add.at. Các hàng không được tra cứu có gradient bằng 0.

Ràng buộc

• Thời gian chạy tối đa: 6000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.