Ràng Buộc Trọng Số (Weight Tying) Giữa Input Embedding & Output LM Head

Mục tiêu kiến thức

• Trong mô hình ngôn ngữ có kích thước từ điển lớn (V = 128,000 như LLaMA-3):
• Input Embedding E ∈ R^{V × D} tốn 128,000 × 4096 × 2 ≈ 1 GB bộ nhớ.
• Output Linear Head W ∈ R^{V × D} cũng tốn thêm 1 GB bộ nhớ.
• Weight Tying (Press & Wolf, 2017 - Dùng trong GPT-2, Gemma, T5):
• Cho phép lm_head.weight = embedding.weight.
• Hai ma trận dùng chung cùng 1 vùng nhớ duy nhất.
• Khi tính logits: z = h E^T (trong đó h là vector ẩn cuối cùng).
• Khi tối ưu hóa, gradient từ cả hai đầu input và output được tích lũy chung vào một ma trận trọng số E, giúp biểu diễn từ vựng hội tụ nhanh hơn.

Yêu cầu

Xây dựng lớp TiedLanguageModel(torch.nn.Module):

class TiedLanguageModel(nn.Module):
    def __init__(self, vocab_size: int, d_model: int):
        pass

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        # Returns logits of shape (B, S, vocab_size)
        pass

• Phải đảm bảo:
• self.embedding = nn.Embedding(vocab_size, d_model).
• self.lm_head = nn.Linear(d_model, vocab_size, bias=False).
• Gán ràng buộc trọng số: self.lm_head.weight = self.embedding.weight.
• Tổng số tham số của mô hình chỉ có V × D (không bị nhân đôi).

Ràng buộc

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