Chào các bạn, trong bài blog này mình sẽ trình bày cách quá trình và cách cài đặt hệ thống dịch máy mạng nơ ron (Neural Machine Translation hay NMT) từ đầu bằng PyTorch. Từ việc chọn data, xử lý data cho đến việc huấn luyện, đánh giá và tương tác với mô hình.

Tại sao mình viết blog này?

Mình hay tiếp cận vấn đề theo kiểu Top-Down, tức là với bài toán NMT, mình sẽ sử dụng các framework như HuggingFace và Fairseq để hình dung được bức tranh lớn của quá trình. HuggingFace cho phép sử dụng các mô hình huấn luyện sẵn và Fairseq tối ưu quá trình cài đặt hệ thống, cho phép mình có thể tạo một hệ thống NMT xịn xò mà không cần hiểu nhiều về quá trình xử lý dữ liệu, cách huấn luyện và đánh giá mô hình. Sau khi biết được hệ thống NMT gồm những quá trình gì, mình tiếp tục tìm hiểu từng thành phần và thử cài đặt lại chúng. Việc này giúp mình hiểu sâu hơn về vai trò của các module và nó ảnh hưởng thế nào đến chất lượng của mô hình, cũng như cho phép mình cài đặt lại chúng theo cách của bản thân. Tuy nhiên với mình, đây là bước khó chịu nhất vì:

• Quá nhiều thông tin để tiếp thu: Việc này hay xảy ra khi mình đọc các bài blog trên mạng. Một bài blog, hay một cuốn sách có thể cung cấp quá nhiều kiến thức về module bạn đang tìm hiểu, trong khi bạn chỉ muốn biết nó hoạt động như thế nào trong mô hình NMT thôi. Đặt biệt là khi bạn đọc source code của framework. Code của framework đã được tối ưu để chương trình chạy nhanh hơn, đánh đổi với việc đọc hiểu dễ dàng (như Beam Search của Fairseq).
• Kết nối các thành phần lại với nhau: Khi bạn đã có thể cài đặt từng module, bước tiếp theo là bạn cần phải ghép chúng lại với nhau. Nhưng khi bạn đọc từ nhiều nguồn khác nhau, và module này không được cài đặt để hỗ trợ module tiếp theo như output của ViTokenizer.tokenize thuộc PyVi không trả về kiểu list như mong đợi lớp Field thuộc TorchText). Thế là bạn phải bỏ ra thời gian để đọc hiểu code của cả hai.
• Debug: Pipeline của bạn đã hoàn chỉnh rồi, chạy được rồi, nhưng kết quả lại tệ hơn mong đợi 😢. Nếu kết quả thấp một cách vô lý, rất có thể bạn xử lý dữ liệu bị sai, mô hình không phù hợp hoặc bộ đo cài đặt không chính xác, thế là phải tốn một khoảng thời gian nữa để debug....

Tất cả các vấn đề trên mình đều gặp phải trong lúc cài đặt hệ thống NMT bằng PyTorch mà không dùng các thư viện mạnh mẽ như HuggingFace và Fairseq. Trong bài này mình sẽ cung cấp các bạn:

• ✅ Quá trình cài đặt môt hệ thống NMT từ giai đoạn đọc dữ liệu cho đến giai đoạn sử dụng mô hình để dịch câu của bạn.
• ✅ Các thuật ngữ, phương pháp, module được sử dụng trong từng quá trình. Mình trình bày những thứ được sử dụng phổ biến trong các hệ thống NMT mà mình từng sử dụng.
• ✅ Mã nguồn cho tất tần tật các quá trình trên, có cả link Google Colab demo. Cung cấp bạn một project mẫu để bạn có thể tự custom theo ý mình.
• ✅ Một website demo nho nhỏ được cài đặt bằng Streamlit cho bạn nào muốn tương tác trực quan.

Và để cho rõ ràng, mình sẽ không:

• ❌ Hướng dẫn cách cài đặt Tokenizer, mặc dù mình nói là cài đặt lại từ đầu, nhưng vẫn có một vài phần nhỏ mình sử dụng các thư viện ngoài PyTorch. Tokenizer mình xin hẹn trình bày ở bài viết tiếp theo.
• ❌ Hướng dẫn chi tiết về mô hình NMT hay Optimizer, hàm Loss. Mục tiêu của mình là hướng dẫn quá trình cài đặt hệ thống NMT hơn là chỉ trình bày về mô hình NMT. Nên trong bài này mình chỉ trình bày ngắn gọn về cách hoạt động của chúng.
• ❌ Cách cài đặt tối ưu cho từng module. Mục tiêu mình là hiểu sao code vậy, hơn là viết code chạy thật nhanh hay sử dụng GPU ít nhất có thể.
• ❌ Cách cài đặt cụ thể cho từng module. Vì code lại từ đầu nên source code của mình khá dài, nên trong bài này mình sẽ không hướng dẫn từng đoạn code mà chủ yếu là lý thuyết và một số vấn đề gặp phải trong lúc cài đặt.

Bạn có thể đọc toàn bộ source code tại repo của mình: https://github.com/egliette/EnVi_NMT

1. Bức tranh lớn về hệ thống NMT

Như các bài toán học sâu khác, đầu tiên mình cần thu thập dữ liệu, sau đó xử lý chúng để đưa về dạng số mà mô hình có thể nhận làm đầu vào, tiếp theo sẽ là huấn luyện cho đến khi mô hình hội tụ và cuối cùng là kiểm tra trên mô hình vừa huấn luyện xong. Và tùy vào bài toán mà nội dung của các công đoạn trên sẽ khác nhau. Trong các mục tiếp theo, mình sẽ trình bài chi tiết mỗi công đoạn trên tương ứng với tác vụ NMT.

2. Thu thập dữ liệu

NMT là bài toán dịch từ source language (ngôn ngữ gốc hay ngôn ngữ nguồn) sang target language (ngôn ngữ đích). Nên dữ liệu cho bài toán này thường được lưu dưới dạng các bộ ngữ liệu song ngữ (parallel corpus), bao gồm các cặp câu của cả hai ngôn ngữ như ("Hello, how are you?", "Xin chào, bạn khỏe không?""). Để chọn bộ dữ liệu để huấn luyện, mình thường hay trả lời các câu hỏi sau:

• Dịch từ ngôn ngữ nào sang ngôn ngữ nào?: Với mỗi ngôn ngữ sẽ có các cách xử lý dữ liệu khác nhau. Ở đây để cho dễ đánh giá, mình chọn dịch từ tiếng Anh sang tiếng Việt.
• Chọn tập ngữ liệu nào?: Khi biết được ngôn ngữ gốc và ngôn ngữ đích, ta cần phải chọn tập dữ liệu song ngữ tương ứng. Trong lúc tìm các tập dữ liệu trên mạng, bên cạnh kích thước dữ liệu, ta cũng cần xem số cặp câu của bộ ngữ liệu. Như các mạng học sâu khác, bạn có thể bị overfit nếu mô hình quá phức tạp trong khi tập dữ liệu quá nhỏ. Vì mình sắp huấn luyện một mô hình Transformer phiên bản nhỏ (tầm 25 triệu tham số), nên mình chọn tập IWSLT'15 Anh-Việt gồm khoảng 133,000 cặp câu, vừa đủ để đạt kết quả tạm ổn.

  Một lưu ý nhỏ là tập ngữ liệu dịch từ tiếng Anh sang tiếng Việt không có nghĩa là nó sẽ hoạt động tốt trên tác vụ dịch từ tiếng Việt sang tiếng Anh. Trong quá trình huấn luyện, mô hình sẽ đọc hiểu câu gốc, sau đó mới dựa trên đó để dịch sang câu đích. Do đó mô hình có thể đọc hiểu tốt tiếng Anh nhưng lại khá tệ trên tiếng Việt, dẫn đến khác biệt về kết quả giữa hai chiều dịch.

• Định dạng dữ liệu ra sao?: Thường thì khi bạn tìm được tập dữ liệu thì bạn đã trả lời được câu này rồi. Các dữ liệu văn bản có thể được lưu dưới dạng file text, csv hay json. Nhưng đối với dịch máy, mình thích định dạng text hơn, lúc này các câu sẽ được ngăn cách bởi dấu xuống hàng "\n", đơn giản và dễ đọc ghi. Và dữ liệu thường sẽ có dạng train.en, train.vi cho tập huấn luyện của mỗi ngôn ngữ, tương tự với tập valid và test.

Một điểm mạnh của tập dữ liệu văn bản là nó khá là nhẹ để lưu trữ, mình thậm chí đặt dữ liệu trên repo github mà không sợ bị vượt quá dung lượng cho phép 😊.

3. Xử lý dữ liệu

Khi nhắc đến các công trình nghiên cứu đến NMT, bạn thường chỉ nghe thấy các tên gọi rất kêu của mô hình như Transformer, BERT-NMT, GPT-2, T5, .... Nhưng để các mô hình NMT học tốt nhất có thể, một quá trình xử lý dữ liệu hiệu quả là không thể thiếu. Trước khi đi vào tìm hiểu các bước, mình sẽ trình bày về một vài thuật ngữ mà bạn sẽ gặp xuyên suốt cả bài hoặc thậm chí là các bài toán xử lý ngôn ngữ tự nhiên (Natural Language Processing hay NLP) khác.

• Source sentence (hay câu gốc, câu nguồn): dùng để chỉ các câu thuộc ngôn ngữ gốc, là câu dùng để dịch sang ngôn ngữ đích.
• Target sentence (hay reference sentence, câu đích): dùng để chỉ câu thuộc ngôn ngữ đích, dùng để so sánh với câu dịch của mô hình.
• Predicted sentence (hay model candidate sentence, câu dịch, câu dự đoán): dùng để chỉ câu được dịch bởi model từ câu gốc, câu này cũng thuộc ngôn ngữ đích.
• Tokenizer (hay word segmenter, bộ tách từ: dùng để tách một câu thành các token, như "Hello, how are you?" thành ["Hello", ",", "how", "are", "you", "?"].

  Vì kết quả trả về của Tokenizer không phải lúc nào cũng là một từ như các ký tự ",", "?" hoặc các cụm kí tự như "He", "llo" (đối với Subword Tokenizer). Nên từ mục này về sau mình sẽ sử dụng từ Token thay cho "từ".

• Vocabulary (hay bộ từ vựng): dùng để chuyển token thành số (index) và ngược lại. Các token được tách bởi Tokenizer sẽ được lưu bởi Vocabulary vào một dict(str, int). Một dict(int, str) cũng có thể được dùng để chuyển index thành token.
• BOS (hay SOS), EOS, UNK, PAD: lần lượt là các token bắt đầu câu (Begin/Start Of String), token kết thúc câu (End Of String), token không xác định (Unknown) và token đệm (Padding). Các token đặc biệt này thường sẽ được khởi tạo đầu tiên trong Vocabulary. Nhằm giúp mô hình biết vị trí nào là bắt đầu câu, vị trí nào câu đã kết thúc, token nào là chưa xác định hoặc vô nghĩa như UNK và PAD.

Few! Khá nhiều thứ để quan tâm trong phần xử lý dữ liệu đối với một người mới tìm hiểu NMT. Nhưng yên tâm, sau khi nắm được các thuật ngữ trên thì việc đọc hiểu các bước sau sẽ dễ dàng hơn nhiều 😄:

Bước 1: Làm sạch dữ liệu: một số dữ liệu văn bản được crawl thường sẽ có thêm các thẻ html, và chúng vô nghĩa trong tác vụ dịch máy nên cần phải lọc bớt các thẻ này. May mắn thay, tập IWSLT'15 Anh-Việt đã làm hộ mình cái công đoạn này.

Bước 2: Tokenize dữ liệu: thông thường sẽ có 2 tokenizer cho 2 ngôn ngữ gốc và đích (mặc dù chúng ta cũng có thể dùng chỉ 1 tokenizer). Để đơn giản, mình chỉ trình bày word tokenizer và multi-word tokenizer. Ôn lại bài cũ một xíu, vì Tiếng Việt có bao gồm từ ghép được tạo bởi 2 tiếng trở lên (như "xây_dựng", "học_sinh") cho nên Tiếng Việt sẽ có 2 cách để tokenize, một là tách từ ghép "học_sinh" thành 2 token riêng biệt bằng word tokenizer, hai là giữ nguyên "học_sinh" như 1 token bằng multi-word tokenizer. Ở đây, mình sẽ chọn word tokenizer của Spacy để tokenize dữ liệu. Thoạt nhìn có vẻ vô lý nhưng nó lại đạt hiệu quả cao hơn so với multi-word trong tác vụ NMT 😵, lí do mình sẽ giải thích trong mục kiểm tra mô hình nhé 😄. Và để giảm bớt lượng token để thêm vào Vocabulary, mình lowercase các câu trước khi tokenize.

Trừ một số loại rule-based tokenizer, các tokenizer khác đặc biệt là subword tokenizer như BPE, WordPiece hay SentencePiece, bạn cần phải huấn luyện cả tokenizer trước khi dùng nó để tokenize câu. Phần này mình hẹn sẽ trình bày trong phần khác nha 😅.

Bước 3: Lọc bớt câu bằng tokenizer: một mô hình NMT chỉ nhận vào list các token với độ dài cho trước. Do đó, chúng ta cần phải bỏ đi những câu chứa quá nhiều token hoặc tăng độ dài tối đa của mô hình (đồng nghĩa với mô hình sẽ to ra). Lưu ý là chúng ta đang xử lý dữ liệu song ngữ, nên chúng ta sẽ lượt bỏ một lúc cả hai câu gốc và đích nếu một trong hai chúng quá dài hoặc quá ngắn (độ dài bằng 0).

Trong IWSLT'15 Anh-Việt vì lấy từ các bài diễn thuyết TedxTalk, nên sẽ có một số câu phụ đề âm thanh tiếng Anh không có chữ nào, nhưng tiếng Việt thì lại có như "vỗ tay" 😵. Rõ ràng là chả có giá trị gì để huấn luyện trên mấy cặp câu như thế này.

Bước 4: Thêm các token vào Vocabulary (build Vocabulary): sau khi tokenize và loại bỏ những câu không hợp lệ, chúng ta sẽ thêm các token vào Vocabulary. Và cả bước này, bạn cũng có nhiều sự lựa chọn khác nhau để tạo một Vocabulary 😢:

• Đầu tiên là số lượng token của Vocabulary, nếu số lượng token quá nhiều có thể làm mô hình bạn phình to lên (do nó số lượng token chính là số lượng vector embedding của mô hình). Nếu tổng số token của dữ liệu vượt quá số lượng cho phép, ta sẽ cắt bớt một lượng token và thay thế các token bị cắt trong câu bằng token UNK. Để giảm lượng UNK trong câu, mình chỉ loại bỏ các token ít xuất hiện nhất. Trong source code, mình lấy hết các token (tổng cộng là 42,706 token cho tiếng Anh và 21,037 token cho tiếng Việt, không tính 4 token đặc biệt 🥶).
• Thứ hai là lọc bớt các token chỉ xuất hiện 1-2 lần như tên người hoặc tên địa danh, vì không xuất hiện nhiều nên mô hình thật sự không học được gì nhiều từ những token này nên có thể lọc bớt để giảm kích thước mô hình. Mô hình mình sắp dùng khá nhỏ nên mình sẽ dùng hết các token có trong tập huấn luyện.
• Cuối cùng là chọn tập dữ liệu để thêm token vào Vocabulary. Dữ liệu trong một project Deep Learning thường chia làm 3 tập nhỏ: train, valid và test. Và chắc chắn một điều là tập train sẽ luôn được dùng để build Vocabulary. Còn tập test và valid thì bạn có thể dùng hoặc không. Nếu bạn không dùng hai tập này để build Vocabulary, bạn có thể giảm kích thước model và giữ nguyên hiệu suất, bởi vì những token không xuất hiện trong tập train sẽ không được học bởi model. Và ngược lại nếu bạn dùng chúng để build Vocabulary, lượng token sẽ tăng lên và mô hình cũng to hơn, nhưng bù lại lúc bạn detokenize (ghép từ - ngược lại với tokenize) sẽ ít xuất hiện các token UNK.

Bước 5: Lưu Tokenizer và Vocabulary: cuối cùng cũng xong công đoạn xử lý dữ liệu 😅. Chúng ta đã thu được: 2 Tokenizer và 2 Vocabulary tương ứng với 2 ngôn ngữ gốc và đích. Và để không tốn thời gian xây dựng lại các đối tượng này, mình sẽ lưu chúng sau khi xử lý dữ liệu xong. Vocabulary thường được xem như một thuộc tính của lớp Tokenizer, nên chỉ cần lưu Tokenizer là đủ. Bạn có thể lưu bằng torch.save và load bằng torch.load nguyên cả Tokenizer, nhưng cách này có thể gặp lỗi khi thay đổi tên file hoặc folder liên quan. Cách giải quyết là bạn chỉ cần lưu dict của Vocabulary vào một file .txt, khi cần dùng chỉ cần load lại file này vào Tokenizer thôi.

Vì các vector embedding trong model được sắp xếp theo thứ tự tương ứng với thứ tự các token trong Vocabulary, nên việc sử dụng Tokenizer có Vocabulary khác sẽ ảnh hưởng đáng kể đến hiệu suất mô hình, thậm chí là không thể nạp checkpoint vào mô hình vì sai khác kích thước.

4. Tạo dữ liệu huấn luyện

Một deep leanring model thường sẽ được nhận vào dữ liệu số và trả về dữ liệu số. Với PyTorch model, dữ liệu số này cần được lưu dưới dạng Tensor. Mình cũng sẽ làm rõ một số lớp sau trước khi đi vào các bước:

• Dataset: như mọi project PyTorch, Dataset như một list với mỗi phần tử của list là một dữ liệu mà mô hình có thể huấn luyện. Ở đây, mình sẽ cài đặt Dataset sao cho mỗi phần tử có dạng dict("src": list(int), "tgt": list(int)), lúc này mỗi phần tử sẽ mang một cặp danh sách các index của câu gốc (src) và câu đích (tgt).
• DataLoader: dùng để chia Dataset thành các batch nhỏ có số lượng câu bằng nhau. Mô hình sẽ học qua tất cả các câu của Dataset trong 1 epoch. Và trong mỗi epoch, mô hình sẽ học qua nhiều batch khác nhau.

Trong quá trình này bạn sẽ thực hiện các bước sau:

Bước 1: Chuyển token sang số: Từ các token thu được ở công đoạn trên, ta sẽ sử dụng Vocabulary để chuyển các token sang index tương ứng. Với những token không có trong Vocabulary, ta chuyển thành index của token UNK.

Bước 2: Tạo Dataset của các cặp câu: như mô tả ở trên, mỗi cặp câu sẽ được lưu vào một dict chứa hai khóa là src và tgt và giá trị tương ứng của chúng sẽ là list các index của token của câu. Do đó Dataset sẽ trông như một list(dict). Mình sử dụng lớp Dataset của PyTorch để tiện trong lúc tạo DataLoader.

Một số mã nguồn sẽ để Dataset trả về list các token thay vì index của chúng để dễ debug bằng cách in từng phần tử ra. Việc chuyển sang index sẽ được thực hiện trong lúc chia batch. Mình chọn chuyển token thành index trong lúc tạo Dataset để tiết kiệm thời gian chia batch lúc huấn luyện.

Bước 3: Tạo DataLoader từ Dataset: với batch size cho trước, DataLoader sẽ chọn batch size phần tử từ Dataset để huấn luyện model cho mỗi bước. Thông thường batch size càng lớn thì model học càng tốt, nhưng batch size quá lớn sẽ làm bạn bị tràn GPU và không huấn luyện được nữa. Tại bước này, bạn gom các dict lại với nhau thành một dict của batch, vẫn gồm hai khóa src và tgt nhưng trị của nó sẽ là Tensor chứa list index của từng câu. Nhưng bạn chỉ có thể tạo Tensor khi các list này có độ dài bằng nhau. Khi đó, ta cần tìm độ dài lớn nhất của các câu trong batch, sau đó đệm vào cuối các câu còn lại bằng index của token PAD cho đến khi tất cả các câu có độ dài bằng độ dài lớn nhất. Khi thu được các list index có cùng độ dài, ta chuyển nó sang dạng Tensor. Tất nhiên khi bắt đầu mỗi epoch, ta cần xáo tập train trước khi chia batch, tránh việc trình tự các câu ảnh hưởng đến khả năng học của mô hình.

Vì token PAD không mang nhiều ý nghĩa nên có quá nhiều PAD trong một batch có thể ảnh hưởng đến kết quả của mô hình. Cho nên nếu bạn sắp xếp các câu theo chiều tăng dần sẽ giảm đáng kể số PAD trong một batch. Mình đã thử sắp xếp Dataset của tập test và valid tăng dần theo chiều dài, không xáo khi bắt đầu mỗi epoch, và sắp xếp giảm dần trong mỗi batch thì thu được loss của mô hình giảm và BLEU tăng đáng kể so với việc giữ nguyên 😋. Mình cũng áp dụng chiến lược trên trong mã nguồn của mình. Nhưng lưu ý rằng nếu bạn đang tham dự cuộc thi hay cố vượt qua một benchmark cho trước, hãy hỏi ban tổ chức rằng việc sắp xếp tập test và valid như thế có hợp lệ hay không nhé ⚠️.

Bước 4: Lưu các DataLoader (không bắt buộc): nếu bạn muốn tiết kiệm thời gian công đoạn này, bạn có thể lưu một lần rồi sử dụng nhiều lần. Đặc biệt khi bạn muốn chia nhỏ source code thành từng file, thay vì nhồi thật nhiều code vào một file. Riêng công đoạn này xử lý khá nhanh nên mình sẽ lưu cả DataLoader bằng torch.save, nếu load có bị lỗi thì cũng không tốn nhiều thời gian tạo lại.

5. Huấn luyện mô hình

Sau một hồi vật lộn với mớ dữ liệu, ta cuối cũng có thể huấn luyện mô hình NMT xịn xò. Như đã nói, mình sẽ không trình bày chi tiết kiến trúc mô hình nhưng chỉ nói tóm lượt thôi. Mình sẽ tái sử dụng mô hình Transformer của Ben Trevett và họ cũng đã giải thích chi tiết chức năng của từng thành phần trong mô hình trong bài viết. Mình cũng dùng optimizer là Adam và hàm loss là Cross Entropy như bài viết gốc.

5.1. Mô hình huấn luyện như thế nào

Một mô hình NMT thông thường sẽ gồm 2 phần: encoder và decoder. Bạn có thể tưởng tượng encoder nhận câu gốc làm đầu vào, đọc hiểu câu gốc và nén những hiểu biết đó thành một Tensor trung gian (hay được gọi là Encoder Output). Tensor trung gian này sẽ cùng với câu đích được dùng làm dữ liệu đầu vào cho decoder. Khi câu trước, có thể bạn sẽ hỏi tại sao câu đích lại được dùng để làm đầu vào? Chẳng phải chúng chỉ được dùng để tính loss thôi sao?. Để trả lời cho câu này chúng ta cần phải hình dung được decoder sẽ làm gì để sinh ra câu dự đoán.

Decoder là một loại mô hình tự hồi quy (autoregressive model), tức decoder sẽ nhận input là một chuỗi dữ liệu có cả output của chính model trước đó và cho ra output là dữ liệu tiếp theo của chuỗi. Trong trường hợp NMT, input của decoder sẽ là encoder output và một chuỗi index của token, đầu tiên chỉ bao gồm token BOS. Decoder dựa vào đó để dự đoán token tiếp theo như "Xin". Sau đó thêm token "Xin" vào chuỗi index của input trước đó, Decoder lặp lại quá trình này liên tục cho đến khi dự đoán token EOS hoặc câu dự đoán dài quá quy định.

Vấn đề là model trong giai đoạn đầu khá là ...ngu, nên chờ nó dự đoán ra EOS hoặc dài quá quy định thì rất tốn thời gian. Ngoài ra các output của decoder khá tệ, dẫn đến input tiếp theo tệ khiến train model không hiệu quả. Cho nên ta sẽ dùng target sentence để giải quyết vấn đề này (như hình vẽ). Bởi vì chúng ta muốn decoder dự đoán đúng ký tự tiếp theo dựa trên chuỗi cho trước nên chúng ta không nhất thiết phải lấy output làm input. Do đó, ta lấy target sentence rồi thêm từng từ sau mỗi bước. Và ta buộc mô hình dừng khi đã thêm hết token từ target sentence. Phương pháp này còn được gọi là Teacher Forcing: cho phép model sử dụng nhãn của các bước trước đó để dự đoán bước hiện tại. Quá trình tạo input từ target sentence có thể được tính toán song song bằng cách dùng ma trận tam giác trái.

Tuy nhiên trong giai đoạn kiểm tra trên tập test, bạn không nên dùng Teacher Forcing mà để decoder tự sử dụng output của nó làm input. Việc này sẽ giúp đánh giá khả năng của mô hình tốt hơn. Đặc biệt là khi dịch câu của người dùng, bạn sẽ không có target sentence để áp dụng Teacher Forcing.

5.2. Lưu checkpoint của mô hình như thế nào

Mình cần lưu checkpoint của model qua các epoch để có thể tiếp tục huấn luyện (trường hợp sử dụng GPU của Google Colab quá thời gian cho phép và bạn phải chờ nó hồi lại) hoặc sử dụng để kiểm tra (test và inference). Tương tự như lưu Tokenizer, mình không ủng hộ việc dùng torch.save và torch.load để lưu toàn bộ model, thay vào đó chỉ sử dụng chúng để lưu một dict chứa các giá trị cần thiết bao gồm:

• epoch: để bạn biết checkpoint này là của epoch thứ mấy để tiếp tục train.
• loss: lưu loss của train hoặc valid tại epoch hiện tại, tại đây mình sẽ lưu valid loss.
• model_state_dict: lưu model.state_dict() - là một dict với khóa là tên các layer và trị là các Tensor parameter của các layer tương ứng. Để load lại trạng thái của model, ta chỉ cần truyền dict này vào model.load_state_dict() thôi.
• optimizer_state_dict: lưu optimizer.state_dict() - cũng là một dict chứa các giá trị của optimizer. Vì optimizer đang dùng là Adam sẽ tính learning rate tương ứng cho từng parameter, nên các learning rate này sẽ thay đổi theo từng epoch. Khi load lại, ta cũng truyền dict này vào optimizer.load_state_dict().

Nếu bạn lưu chỉ để kiểm tra thôi thì chỉ cần lưu model_state_dict. Trong lúc huấn luyện, mình hay lưu 2 file checkpoint. Một checkpoint_best để chứa checkpoint có loss thấp nhất, checkpoint này sẽ được dùng để kiểm tra. Và một checkpoint_last để chứa checkpoint tại epoch sau cùng, checkpoint này sẽ được dùng để tiếp tục huấn luyện.

6. Kiểm tra mô hình

Đến công đoạn cuối rồi 👏, nhưng không có nghĩa là công đoạn dễ nhất hay ngắn nhất 😂. Bình thường thì chúng ta sẽ xem loss và metric của checkpoint trên tập test. Mình sẽ phân biệt loss và metric như thế này:

• Loss (hay giá trị mất mát): là giá trị để model đọc hiểu, nên càng thấp càng tốt. Mục tiêu của model sẽ là làm giảm giá trị loss bằng cách thay đổi tham số chính nó. Loss còn được dùng để dự báo model có bị overfit hay underfit hay không bằng cách xem loss của tập train và tập valid thay đổi như thế nào qua từng epoch.
• Metric (hay bộ đo): là giá trị để con người đọc hiểu, nên càng cao càng tốt (giống như điểm học tập của bạn vậy 💯). Vì mục đích là đọc hiểu, nên mỗi metric sẽ có một ý nghĩa riêng dựa trên đánh giá của con người. ⚠️ Vì mục tiêu của model là giảm loss, chứ không phải là tăng metric, nên việc model đạt loss thấp không đồng nghĩa với model sẽ có metric cao.

Với loss, chúng ta sẽ tính bằng hàm Cross Entropy Loss như lúc huấn luyện. Riêng metric mình sẽ sử dụng BLEU - một metric phổ biến cho bài toán NMT.

6.1. BLEU - BiLingual Evaluation Understudy

Trước khi giải thích về BLEU, mình sẽ nói về n-gram trước. n-gram có thể hiểu là cụm gồm n token phân biệt đứng kề nhau. Ví dụ trong câu ["một", "cộng", "một", "bằng", "hai"] sẽ có các n-gram là

• 1-gram (hay unigram): gồm 4 unigram sau ["một", "cộng", "bằng", "hai"].
• 2-gram (hay bigram): gồm 4 bigram sau ["một cộng", "cộng một", "một bằng", "bằng hai"].
• 3-gram (hay trigram): gồm 3 trigram sau ["một cộng một", "cộng một bằng", "một bằng hai"]
• 4-gram: gồm 2 4-gram sau ["một cộng một bằng", "cộng một bằng hai"].

Khi đó ý tưởng chính của BLEU là đếm số n-gram trùng lặp giữa target sentence và predicted sentece (với n thường là cả 1, 2, 3 và 4).

Cơ bản BLEU được tính như sau:

Nhìn hơi phức tạp nhỉ, nhưng sẽ đơn giản hơn khi ta chia BLEU thành hai phần như sau:

• N-gram Overlap (hay Geometric Average Precision, Geometric Mean): tính tỉ lệ n-gram trùng khớp giữa predicted sentence và target sentence thông qua precision-i (với i mang giá trị từ 1 đến n).
• Brevity Penalty: dùng để phạt khi predicted sentence quá ngắn. Vì precision-i được tính bằng cách lấy tổng số i-gram trùng khớp chia cho tổng số i-gram của predicted sentence. Dẫn đến trường hợp predicted sentence rất ngắn nhưng lại chứa các n-gram trong target sentence, sẽ cho ra N-gram Overlap mang giá trị cao. Ví dụ như target sentence: "Xin chào, bạn như thế nào, tôi vẫn khỏe" và predicted sentence: "Xin chào, bạn như" sẽ cho ra precision-i là 100%. Ý tưởng của Brevity Penalty là khi số token của predicted sentence ít hơn target sentence, giảm BLEU theo hàm số mũ.

Trong lúc mình cài đặt BLEU và so sánh với hàm tính BLEU của TorchText và NLTK thì có một số nhận xét cũng như lưu ý như sau:

• BLEU bằng 0 cho dù có nhiều n-gram trùng nhau: bởi vì N-gram Overlap tính bằng phép nhân giữa các precision-i, nên chỉ cần một precision-i bằng 0 cũng làm BLEU bằng 0. Nên với những câu dịch không có 4-gram nào trùng khớp thì BLEU = 0 cho dù các n-gram còn lại trùng khớp rất nhiều. Hoặc target sentence và predicted sentence có ít hơn 4 token làm precision-4 bằng 0.

  Để giải quyết vấn đề này, đã có một số phương pháp làm mượt (smoothing) để ngăn không cho N-gram Overlap bằng 0 khi có ít nhất một n-gram trùng khớp như BLEUS. Hàm tính BLEU của thư viện NLTK cũng có cài đặt phương pháp làm mượt. Mình tạm thời không cài đặt phương pháp này trong mã nguồn.

• BLEU của nhiều câu (hay corpus BLEU) không tính bằng trung bình cộng của BLEU từng câu (hay sentence BLEU): không giống như loss của epoch được tính bằng trung bình cộng của loss từng batch. Để tính BLEU trên nhiều câu, biểu thức phân số của N-gram Overlap và Brevity Penalty sẽ có tử là tổng các tử của từng câu, sau đó đem chia cho mẫu cũng được tính bằng tổng các mẫu. N-gram Overlap và Brevity Penalty thu được nhân lại với nhau như bình thường, ta thu được corpus BLEU. Cách tính này gọi là micro-average precision, bạn có thể tham khảo tại đây.

• BLEU không phản ánh mặt ngữ nghĩa của câu: vì cơ bản nó chỉ làm đếm số token trùng khớp, nên "tôi" và "tớ" cơ bản là hai token khác nhau. Vì thế việc này dẫn đến word tokenizer sẽ cho ra kết quả BLEU tốt hơn so với multi-word tokenizer như mình đã nói trong mục . Xét hai từ "y tế" và "y khoa", với multi-word tokenizer, ta sẽ thu được 1 token cho mỗi từ: "y_tế", "y_khoa", dẫn đến không có token nào trùng khớp trong khi cả hai từ này có nghĩa gần nhau. Trong khi với word tokenizer, ta thu được 2 token cho mỗi từ: ["y", "tế"], ["y", "khoa"], trùng nhau 1 token "y". Tương tự vậy cho các từ ghép tiếng Việt khác, dần dần word tokenizer sẽ có số n-gram trùng khớp nhiều hơn multi-word tokenizer, dẫn đến BLEU cao hơn.

  Một phiên bản đầy đủ của BLEU là sẽ nhận nhiều target sentence cho một predicted sentence, cho phép sử dụng nhiều câu target sentence đồng nghĩa để đánh giá predicted sentence. Tuy nhiên việc chuẩn bị nhiều target sentence cho mỗi predicted sentence sẽ tốn rất nhiều công sức nên thường chỉ được cung cấp một target sentence tương ứng cho một source sentence. Hiện tại mình chỉ cài đặt BLEU nhận vào một target sentence cho mỗi predicted sentence

• Loss giảm không đồng nghĩa BLEU tăng: mục đích của việc huấn luyện model là giảm loss, khiến mô hình hội tụ, chứ không phải là tìm cách tăng BLEU. Sẽ không có gì lạ nếu mô hình của bạn đạt loss rất thấp trên cả hai tập train và valid nhưng lại có BLEU rất thấp trên tập valid. Do đó tốt nhất vẫn nên kiểm tra BLEU của mô hình sau một vài epoch, để đảm bảo mô hình hoạt động tốt hay loss giảm, BLEU tăng. Từ đó phát sinh ra vấn đề sau.

• BLEU tính rất lâu trên các model mới train: để tính BLEU, ta cần predicted sentence của model ứng với target sentence. Như , mô hình mới train trong 1, 2 epoch thường sẽ tốn rất nhiều thời gian để sinh predicted sentence. Trong khi ta cần kiểm tra mô hình có hoạt động tốt hay không, đôi khi ta cần tính BLEU cho từng epoch. Để giảm thời gian, mình áp dụng Teacher Forcing để sinh predicted sentence, sau đó mới tính BLEU. Như đã nói, Teacher Forcing không nên áp dụng khi tính BLEU trên tập test để có thể đánh giá đúng khả năng của mô hình.

Ngoài ra còn có một phương pháp có thể giúp tăng BLEU của model, dựa trên một heuristic tên là Beam Search Decoding.

6.2. Beam Search Decoding

Như thường lệ, mình sẽ trình bày một vài nội dung cần thiết trước khi giải thích Beam Search Decoding. Như các mình đã nói ở mục , decoder sẽ nhận chuỗi index để dự đoán token đầu ra. Token đầu ra này mặc định sẽ được chọn là Token có xác suất cao nhất trong Vocabulary. Mục đích cuối cùng là tìm chuỗi các token sao cho tích các xác suất cao nhất. Phương pháp tìm tích các xác suất cao nhất bằng cách chọn token có xác suất cao nhất trong từng bước này gọi là Greedy Search.

Vì là một heuristic, Greedy Search sẽ không 100% cho ra tích xác suất cao nhất có thể. Vẫn có trường hợp:

• Tại bước i, token x mang giá trị cao nhất là 0.5, kế đến là token y mang giá trị cao nhì là 0.4.
• Tại bước i+1, với token x, model cho ra token x_1 mang xác suất cao nhất là 0.4, tích xác suất lúc này là 0.5 x 0.4 = 0.2. Trong khi với y, model cho ra y_1 mang xác suất cao nhất là 0.9, tích xác suất lúc này là 0.4 x 0.9 = 0.36. Cao hơn so với x mang xác suất cao nhất trong bước trước.

Cách duy nhất để chắc chắn tìm được chuỗi token có tích xác suất cao nhất xét hết tất cả khả năng có thể. Phương pháp này gọi là Exhaustive Search. Việc xét từng token của Vocabulary cho mỗi bước sẽ có độ phức tạp lũy thừa $O(c^n)$ với $c$ là số token của Vocabulary .

Beam Search Decoding (mình gọi tắt là Beam Search) là một cách để cân bằng giữa Greedy Search và Exhaustive Search. Ý tưởng là model sẽ giữ $n$ bản dịch tốt nhất tại mỗi bước. Với $n$ (hay beam size, beam width, number of beam) là một siêu tham số do mình tự chọn.

Một thay đổi nhỏ trước khi trình bày các bước của Beam Search là vì các xác suất này đều nằm trong khoảng từ 0 đến 1 nên tích của chúng sẽ càng ngày càng nhỏ. Do đó mình sẽ sử dụng hàm lognepe (ở đây mình viết tắt là log) là một hàm tăng đơn điệu và nhận giá trị âm trong khoảng từ 0 đến 1. Khi đó bài toán sẽ chuyển từ tìm tích xác suất lớn nhất $max(\prod p)$ thành tìm giá trị lớn nhất của tổng các log $max(\sum log(p))$ (để có được kết quả dương, bạn có thể tìm $min(-\sum log(p))$).

Với:

• $max\_len$ là độ dài tối đa cho phép của câu dịch.
• $beam\_size$ là số bản dịch tốt nhất tối đa tại mỗi bước.
• $candidate\_list$ là danh sách các chuỗi ứng viên để chọn ra chuỗi dịch hoàn tất, kèm theo mỗi chuỗi là tổng log xác suất từ bước đầu cho đến hiện tại. Ban đầu $candidate\_list$ chỉ gồm một phần tử là $([bos\_id], 0.0)$.
• $completed\_candidates$: là danh sách các chuỗi ứng viên có token cuối là EOS hoặc đạt độ dài tối đa.

Khi đó Beam Search được thực hiện như sau:

# Với i là chiều dài hiện tại của các câu ứng viên, # duyệt cho đến khi i đạt max_len
FOR i = 0..(max_len - 1) DO n_cands := LENGTH(candidate_list) new_candidates = [] # Với mỗi ứng viên, tìm top beam_size token kế tiếp # đạt log softmax cao nhất FOR j = 0..(n_cands - 1) DO pred_indexes, log = candidate_list[j] output := model(pred_indexes) logs := log_softmax(output) # Chỉ lấy top beam_size token đạt log cao nhất top_logs, indexes := top_k(logs, k=beam_size) # Thêm top beam_size token đạt log cao nhất # vào new_candidates FOR k = 0..(beam_size-1) DO new_indexes := pred_indexes + [indexes[k]] new_log := log + top_logs[k] new_candidates := (new_indexes, new_log) # Trước khi chọn ra candidate hoàn chỉnh, # từ new_candidates chọn ra top_new_candidates # bao gồm các candidate có log cao nhất sao cho có độ dài thỏa # beam_size == LENGTH(top_new_candidates) + LENGTH(completed_candidates) # đảm bảo mình đang giữ beam_size chuỗi tốt nhất. top_new_candidates := top_k(new_candidates, k=beam_size - LENGTH(completed_candidates), key=candidate_log) # Chọn ra các chuỗi ứng viên có độ dài i+1 hoặc thêm vào completed_candidates candidate_list = [] FOR candidate IN top_new_candidates DO indexes, log = candidate IF indexes[-1] == EOS_ID THEN completed_candidates += [candidate] # Nếu thu thập đủ beam_size câu hoàn chỉnh thì không cần duyệt nữa IF LENGTH(completed_candidates) == beam_size): BREAK ELSE: candidate_list += [candidate] # Nếu duyệt đủ max_len vòng, candidate_list sẽ chứa các chuỗi đạt độ dài tối đa
# các chuỗi này sẽ được thêm vào completed_candidate
completed_candidates += candidate_list

Bạn có thể quan sát ví dụ từng bước sau:

Vì ví dụ chỉ mang tính chất minh họa 😅 nên mình đã tạo ra trường hợp đẹp nhất, tức là Beam Search giúp tìm ra câu có nhiều token khớp với target sentence hơn, từ đó làm tăng điểm BLEU. Nhưng vẫn có một số vấn đề cần lưu ý sau về Beam Search:

• Beam Search trở thành Greedy Search khi beam size = 1: cái này dễ hiểu mà, mình chỉ nhấn mạnh thôi 😙.

• Predicted sentence có log cao nhất không đồng nghĩa là câu này cho ra BLEU cao nhất: nó chỉ là các token mà mô hình cho rằng hay xuất hiện với nhau nhất trong tập train với beam size = n thôi. Các câu còn lại thuộc top 2, top 3 vẫn có thể gần giống với target sentence hơn top 1.

• Predicted sentence ngắn nhất có xu hướng có log cao nhất: Vì càng nhiều token, càng nhiều giá trị âm được cộng vào, dẫn đến các predicted sentence có càng ít token sẽ có log càng lớn. Giống như ví dụ minh họa của mình ở trên, candidate 1 ít hơn candidate 2 một token, nên dễ dàng đạt log cao hơn. Để giải quyết vấn đề này, ta có thể sử dụng length normalization (hay length penalty). Ý tưởng là sẽ chia log cho một hàm phụ thuộc vào số token, thu được score mới ít chịu tác động bởi số token của câu. Mình hiện tại khá hài lòng với kết quả nên sẽ không cài đặt phương pháp length normalization, các bạn có nhu cầu tìm hiểu có thể xem qua video Refining Beam Search trên youtube.

• Sử dụng kết quả Beam Search như thế nào để tính BLEU: với beam size = n, ta sẽ thu được n predicted sentence ứng với mỗi target sentence. Mình sẽ cài đặt như Fairseq là chọn câu cho ra log cao nhất để tính BLEU với target sentence.

  Bạn có thể chọn toàn bộ hoặc top k các predicted sentence để tính BLEU. Để đo BLEU bằng các kết quả của Beam Search thay vì chỉ sử dụng top 1, bạn chỉ cần thêm vào các cặp câu (predicted 1, target), (predicted 2, target),... như cặp câu bình thường. Tuy nhiên nếu các câu này có BLEU thấp hơn top 1, BLEU chung sẽ bị giảm đáng kể, và mình cũng chưa thấy ai làm như thế này 😗.

• Tăng beam size không đồng nghĩa với BLEU của model sẽ tăng: việc này tương tự với loss và BLEU, việc tăng beam size nhằm chọn ra tổng các log mang giá trị cao nhất trong phạm vi cho phép, chứ không phải cố làm tăng BLEU của model. Tuy nhiên, theo giải thích ở trên, các token có log cao hay xuất hiện nhau với nhau trong tập train, nên khả năng câu sẽ đúng về mặt ngữ pháp và hợp lý hơn.

• Chọn beam size như thế nào: mình hay chọn bắt đầu từ 1 (Greedy Search), sau đó tăng dần theo bội của 5 như 5, 10, 15. Đến khi nào BLEU có vẻ giảm dần thì chọn xung quanh số mà BLEU bắt đầu giảm.

• Beam Search sinh ra các token bị lặp: Vấn đề này hay gặp ở các bài toán NMT, cho dù có sử dụng Beam Search hay không. Ví dụ như trong tập train có chứa các token hay lặp lại liên tục nhiều lần như "ha" trong "He laughed: ha ha ha ha". Việc này sẽ làm model cho rằng token "ha" hay theo sau token "ha", tạo thành một vòng lặp "ha ha ha ha ha..." làm ảnh hưởng nghiêm trọng đến kết quả mô hình. Một giải pháp cho phương pháp này là các kỹ thuật sampling. Ý tưởng là thay vì 100% chọn token có xác suất cao nhất, mình chỉ 80% chọn token cao nhất, các token còn lại vẫn có xác suất được chọn. Từ đó cho phép mô hình thoát khỏi vòng lặp "ha ha" trên. Hiện tại mình không cài đặt phương pháp sampling trong source code, các bạn nếu tò mò thì có thể tham khảo tại bài blog khá hay này.

Cài đặt đến bước này rồi thì cơ bản bạn cũng đã có một hệ thống NMT hoàn chỉnh rồi 🥳. Bạn chỉ cần thay đổi các thông số như kích thước Vocabulary, các kích thước model, kích thước batch, số epoch,... để model hoạt động tốt nhất. Sau đây mình sẽ cung cấp source code để tương tác với model bằng Streamlit như một lời cảm ơn các bạn đã đọc đến đây 🤗.

7. Tương tác với mô hình bằng Streamlit

Khi đã train xong xuôi rồi, bạn chỉ cần chạy dòng lệnh sau với file config.yml là file config bạn dùng để train mô hình:

streamlit run inference_streamlit.py -- --config config.yml

Tuy nhiên, nếu máy bạn không có GPU, mình cũng có cung cấp link Google Colab mẫu hướng dẫn chạy server bằng localtunnel, sử dụng model và vocab mà mình huấn luyện trước: https://colab.research.google.com/drive/1grZfoaD9SCfwOJi-QDOdGRJMlaekDY3f

Website này sẽ cho bạn dịch câu tiếng Anh sang tiếng Việt, sau đó hiển thị Attention Matrix cho head đầu tiên của mô hình, và cuối cùng là danh sách kết quả Beam Search với beam size được xác định trong file config.yml. ⚠️ Vì hiện tại mình chưa code lấy Attention Matrix cho các Beam Search candidate nên kết quả dịch và Attention Matrix mà các bạn thấy trên là kết quả của Beam Search với beam size = 1, còn các kết quả của Beam Search mình sẽ để ở cuối website.

8. Kết bài

Bài viết này dài hơn mình tưởng 😅. Trong bài này chúng ta đã:

• Trình bày chi tiết từng công đoạn từ xử lý dữ liệu đến kiểm tra model.
• Các vấn đề và cách giải quyết trong mỗi công đoạn.
• Tương tác với model thông qua Streamlit.

Cảm ơn bạn đã đọc đến tận đây, rất mong bài blog của mình giúp ích cho các bạn!

9. Tham khảo

• Repo của mình: https://github.com/egliette/EnVi_NMT
• Link Google Colab demo
• Link Google Colab Streamlit
• Tập ngữ liệu IWSLT'15 Anh-Việt
• Stanford CS224n Course
• Seq2Seq Tutorial of Ben Trevett
• BLEU paper
• BLEU Score - Google Cloud
• BLEU Score - Wikipedia
• BLEUS
• How to Implement a Beam Search Decoder for Natural Language Processing
• Transformer’s Evaluation Details: Greedy and Beam Search Translators
• Beam Search code of Fairseq
• Fairseq Generator
• Sampling Techniques
• Refining Beam Search