Nhận dạng tiếng nói tự động¶
Nhận dạng tiếng nói tự động chuyển đổi âm thanh lời nói thành văn bản viết, kết nối khoảng cách giữa giọng nói con người và ngôn ngữ máy có thể đọc được. File này bao quát GMM-HMM, CTC loss, RNN-Transducer, các mô hình encoder-decoder dựa trên cơ chế chú ý (LAS), Whisper, và ASR đầu-cuối, từ các pipeline cổ điển đến kiến trúc nơ-ron hiện đại.
-
Nhận dạng tiếng nói tự động (automatic speech recognition — ASR) là tác vụ chuyển đổi âm thanh lời nói thành văn bản viết. Đây là một trong những bài toán lâu đời nhất trong AI (các hệ thống đầu tiên từ những năm 1950 nhận dạng được chữ số đơn) và cũng là một trong những ứng dụng thương mại phổ biến nhất (trợ lý giọng nói, dịch vụ phiên âm, phụ đề).
-
Khó khăn đến từ sự biến thiên to lớn của lời nói: người nói khác nhau, giọng địa phương, tốc độ nói, nhiễu nền, đặc tính microphone, và sự mơ hồ cơ bản khi ánh xạ một tín hiệu âm thanh liên tục ra các từ rời rạc.
-
Hãy hình dung ASR giống như một nhân viên tốc ký tại tòa án. Nhân viên nghe một luồng âm thanh liên tục, tinh thần phân đoạn nó thành các từ, giải quyết sự mơ hồ bằng ngữ cảnh ("their" vs "they're" vs "there"), và gõ ra kết quả. Một hệ thống ASR cũng làm điều tương tự, nhưng theo các giai đoạn có thể được làm rõ ràng và tối ưu một cách độc lập hoặc đồng thời.
-
Pipeline ASR cổ điển xử lý âm thanh qua một chuỗi các giai đoạn riêng biệt: âm thanh thô được biến đổi thành đặc trưng (MFCC hoặc log-mel spectrogram, từ file 01), một mô hình âm học (acoustic model) chấm điểm mức độ phù hợp của mỗi khung đặc trưng với mỗi đơn vị ngữ âm, một mô hình phát âm (lexicon) ánh xạ các đơn vị ngữ âm thành từ, một mô hình ngôn ngữ (language model) chấm điểm độ khả dĩ của chuỗi từ, và một bộ giải mã (decoder) tìm kiếm chuỗi từ tối đa hóa điểm tổng hợp. Mỗi thành phần được huấn luyện và tinh chỉnh riêng.
-
Âm vị (phonemes) là đơn vị âm thanh nhỏ nhất có thể phân biệt các từ trong một ngôn ngữ. Tiếng Anh có khoảng 39–44 âm vị (con số chính xác phụ thuộc vào thổ ngữ và bộ âm vị được sử dụng). Ví dụ, "bat" và "pat" khác nhau một âm vị (/b/ vs /p/). Hầu hết các hệ thống ASR mô hình hóa các âm vị phụ thuộc ngữ cảnh gọi là triphone: một âm vị được định nghĩa bởi các láng giềng trái và phải của nó (ví dụ, "a" trong ngữ cảnh "b_t" là một đơn vị khác với "a" trong ngữ cảnh "c_t"), vì hiện thực âm học của một âm vị bị ảnh hưởng mạnh bởi các láng giềng của nó (gọi là đồng cấu âm — coarticulation).
-
Số lượng triphone có thể có là khổng lồ (40 âm vị mũ 3 = 64,000), nên phân cụm cây quyết định (decision tree clustering) nhóm các triphone tương tự về mặt âm học thành senone (thường 2000–10,000 lớp). Mỗi senone có mô hình âm học riêng. Sự phân cụm này là một dạng của thuật toán cây quyết định từ chương 06.
-
GMM-HMM (Mô hình hỗn hợp Gaussian — Mô hình Markov ẩn) là phương pháp mô hình âm học thống trị từ thập niên 1980 đến đầu những năm 2010. HMM (từ chương 05) mô hình hóa cấu trúc thời gian của lời nói: mỗi âm vị là một HMM từ trái sang phải với 3–5 trạng thái, trong đó mỗi trạng thái biểu diễn một đoạn dưới âm vị (bắt đầu, giữa, kết thúc). Các chuyển tiếp trạng thái mô hình hóa thời lượng một cách ngầm định.
-
Tại mỗi trạng thái HMM, xác suất phát xạ (emission probability — khả năng một vector đặc trưng cụ thể xuất hiện khi ở trạng thái đó) được mô hình hóa bởi một mô hình hỗn hợp Gaussian (GMM): một tổng có trọng số của các phân bố Gaussian đa biến (từ chương 05):
-
trong đó \(\mathbf{x}\) là vector đặc trưng (ví dụ, MFCC 39 chiều), \(s\) là trạng thái HMM, \(M\) là số thành phần hỗn hợp (thường 8–64), \(w_m\) là trọng số hỗn hợp, và \(\boldsymbol{\mu}_m\), \(\boldsymbol{\Sigma}_m\) là trung bình và hiệp phương sai của mỗi thành phần Gaussian. Ma trận hiệp phương sai thường là đường chéo để tính toán hiệu quả (giả định các chiều đặc trưng độc lập, điều này xấp xỉ đúng cho MFCC nhờ sự bất tương quan hóa của DCT).
-
Huấn luyện dùng thuật toán Baum-Welch (một trường hợp đặc biệt của EM, từ chương 05) để ước lượng lặp các tham số GMM và xác suất chuyển tiếp HMM từ dữ liệu tiếng nói đã được phiên âm. Giải mã (tìm dãy trạng thái có khả năng nhất) dùng thuật toán Viterbi (quy hoạch động, từ chương 05):
-
trong đó \(\delta_t(j)\) là xác suất của đường đi tốt nhất kết thúc ở trạng thái \(j\) tại thời điểm \(t\), \(a_{ij}\) là xác suất chuyển tiếp từ trạng thái \(i\) sang \(j\), và \(b_j(\mathbf{x}_t)\) là xác suất phát xạ của đặc trưng \(\mathbf{x}_t\) ở trạng thái \(j\).
-
DNN-HMM (Hinton và đồng sự, 2012) đã thay thế mô hình phát xạ GMM bằng một mạng nơ-ron sâu (DNN, từ chương 06) dự đoán xác suất hậu nghiệm senone \(p(s | \mathbf{x})\) từ một cửa sổ các khung đặc trưng. HMM vẫn xử lý cấu trúc thời gian và trình tự, nhưng mạng nơ-ron cung cấp các điểm phát xạ có tính phân biệt cao hơn nhiều. Cách tiếp cận hybrid này giảm tỷ lệ lỗi từ (WER) xuống 20–30% so với GMM và là mô hình thống trị từ 2012–2016.
-
Giải mã WFST (Weighted Finite-State Transducer) là khung giải mã tiêu chuẩn cho ASR truyền thống. Mỗi thành phần (cấu trúc HMM H, phụ thuộc ngữ cảnh C, từ điển L, ngữ pháp/mô hình ngôn ngữ G) được biểu diễn như một bộ chuyển đổi trạng thái hữu hạn có trọng số, và chúng được kết hợp thành một đồ thị tìm kiếm duy nhất \(H \circ C \circ L \circ G\). Tìm kiếm Viterbi sau đó tìm đường đi có chi phí thấp nhất qua đồ thị tổng hợp này. WFST cho phép kết hợp mô-đun các nguồn tri thức và tìm kiếm quy hoạch động hiệu quả. Khung toán học này đến từ lý thuyết automata hữu hạn (liên quan đến các máy trạng thái trong chương 05).
-
ASR đầu-cuối (end-to-end) loại bỏ các thành phần riêng biệt (mô hình phát âm, bộ âm vị, bộ giải mã WFST) và huấn luyện một mạng nơ-ron duy nhất ánh xạ trực tiếp từ đặc trưng âm thanh ra ký tự hoặc từ con (word pieces). Thách thức then chốt là bài toán căn chỉnh (alignment problem): đầu vào (hàng trăm khung đặc trưng mỗi giây) và đầu ra (vài ký tự mỗi giây) có độ dài rất khác nhau, và sự căn chỉnh giữa chúng không được biết trong quá trình huấn luyện.
-
Connectionist Temporal Classification (CTC) (Graves và đồng sự, 2006) giải quyết bài toán căn chỉnh bằng cách đưa vào một token blank đặc biệt và cho phép mạng xuất ra bất kỳ dãy ký tự và blank nào, miễn là khi gộp các ký tự lặp lại và loại bỏ blank sẽ thu được bản phiên âm đúng. Ví dụ, bản phiên âm "cat" có thể được tạo ra bởi dãy đầu ra "--cc-aa-t--" (trong đó "-" là blank).
-
Một cách hình thức, CTC định nghĩa một ánh xạ nhiều-một \(\mathcal{B}\) từ tập tất cả các dãy đầu ra độ dài \(T\) (trên bảng chữ cái cộng blank) sang các dãy nhãn. Xác suất của một dãy nhãn \(\mathbf{y}\) là tổng trên tất cả các căn chỉnh gộp về nó:
-
Tính tổng này một cách ngây thơ đòi hỏi liệt kê số lượng căn chỉnh mũ (exponential), nhưng thuật toán forward-backward CTC tính nó hiệu quả trong \(O(T \cdot |\mathbf{y}|)\) bằng quy hoạch động, tương tự thuật toán forward-backward HMM từ chương 05.
-
CTC đưa ra một giả định độc lập có điều kiện (conditional independence assumption): đầu ra tại mỗi bước thời gian độc lập với mọi đầu ra khác khi đã biết đầu vào. Điều này có nghĩa CTC không thể mô hình hóa các phụ thuộc đầu ra (ví dụ, nó không thể học được rằng "q" hầu như luôn được theo sau bởi "u"). Một mô hình ngôn ngữ bên ngoài phải được dùng để xử lý các phụ thuộc như vậy.
-
Các tùy chọn giải mã CTC:
- Giải mã tham lam (greedy decoding): lấy token có xác suất cao nhất tại mỗi bước thời gian, sau đó gộp. Nhanh nhưng không tối ưu.
- Tìm kiếm chùm (beam search): duy trì \(k\) giả thuyết bộ phận tốt nhất tại mỗi bước, gộp các giả thuyết cùng gộp về cùng một tiền tố. Có thể kết hợp điểm mô hình ngôn ngữ.
- Tìm kiếm chùm tiền tố (prefix beam search): một biến thể của beam search xử lý đúng việc gộp blank CTC, đảm bảo các giả thuyết được so sánh sau khi gộp.
-
RNN-Transducer (RNN-T) (Graves, 2012) mở rộng CTC bằng cách thêm một mạng dự đoán (prediction network) rõ ràng—một RNN giống mô hình ngôn ngữ—điều kiện hóa mỗi đầu ra dựa trên các đầu ra trước đó, loại bỏ giả định độc lập có điều kiện. RNN-T có ba thành phần:
- Encoder: xử lý đặc trưng âm thanh để tạo ra biểu diễn ẩn \(\mathbf{h}_t^\text{enc}\) (thường là một chồng các lớp LSTM hoặc Conformer).
- Mạng dự đoán: một RNN tự hồi quy tạo ra biểu diễn ẩn \(\mathbf{h}_u^\text{pred}\) từ các nhãn đã phát ra trước đó.
- Mạng kết hợp (joint network): kết hợp đầu ra của encoder và mạng dự đoán tại mỗi vị trí (thời gian, nhãn) và tạo ra một phân bố trên token tiếp theo (bao gồm blank):
-
RNN-T có thể phát ra không hoặc nhiều nhãn mỗi bước thời gian (bằng cách phát ra token không-blank trước khi tiến đến bước thời gian tiếp theo, hoặc phát ra blank để tiến tới mà không có đầu ra). Huấn luyện dùng thuật toán forward-backward trên lưới 2D (thời gian, nhãn), với độ phức tạp \(O(T \cdot U)\) trong đó \(U\) là độ dài đầu ra. RNN-T là kiến trúc thống trị cho ASR streaming trên thiết bị (dùng trong điện thoại Pixel của Google và các sản phẩm tương tự) vì nó hỗ trợ streaming một cách tự nhiên: encoder xử lý âm thanh từ trái sang phải và mạng dự đoán sinh đầu ra tăng dần.
-
Listen, Attend and Spell (LAS) (Chan và đồng sự, 2016) là một mô hình encoder-decoder dựa trên cơ chế chú ý (kiến trúc sequence-to-sequence từ chương 06). Nó có ba thành phần:
- Listener (encoder): một LSTM hai chiều hình tháp xử lý toàn bộ dãy đầu vào và giảm tốc độ mẫu (downsample) 8 lần (bằng cách nối ghép các trạng thái ẩn liên tiếp ở mỗi tầng), tạo ra một dãy ngắn hơn các trạng thái ẩn encoder.
- Attention: tại mỗi bước decoder, tính trọng số chú ý trên tất cả các trạng thái encoder để tạo thành một vector ngữ cảnh (cùng cơ chế chú ý từ chương 07).
- Speller (decoder): một LSTM tự hồi quy sinh ra bản phiên âm đầu ra từng ký tự một, được điều kiện hóa bởi vector ngữ cảnh và các ký tự đã sinh trước đó.
-
LAS đạt kết quả mạnh mẽ nhưng yêu cầu toàn bộ câu nói phải có sẵn trước khi giải mã (vì cơ chế chú ý nhìn vào tất cả các trạng thái encoder), làm nó không phù hợp cho các ứng dụng streaming. Nó cũng gặp khó khăn với các câu nói rất dài vì cơ chế chú ý trên các dãy dài trở nên khuếch tán.
-
Conformer (Gulati và đồng sự, 2020) kết hợp khả năng nắm bắt mẫu cục bộ của phép tích chập với khả năng mô hình hóa phụ thuộc toàn cục của self-attention. Mỗi khối Conformer có bốn module theo cấu trúc sandwich:
- Module feed-forward (nửa bước): mạng feed-forward với kết nối dư, dùng một nửa trọng số dư.
- Module multi-head self-attention: self-attention Transformer chuẩn (từ chương 07) với mã hóa vị trí tương đối.
- Module tích chập (convolution): gồm tích chập điểm (pointwise convolution), cổng tuyến tính (GLU), tích chập chiều sâu 1D, chuẩn hóa batch, kích hoạt Swish, và một tích chập điểm khác. Tích chập chiều sâu nắm bắt ngữ cảnh cục bộ (giống n-gram trên dãy đặc trưng).
- Module feed-forward (nửa bước): giống hệt module 1.
-
Đầu ra là: \(\mathbf{y} = \text{LayerNorm}(\mathbf{x} + \frac{1}{2}\text{FFN}_1 + \text{MHSA} + \text{Conv} + \frac{1}{2}\text{FFN}_2)\). Cấu trúc macaron-like (FFN-Attention-Conv-FFN) với các kết nối dư nửa bước được tìm thấy thực nghiệm là vượt trội so với các thứ tự khác. Conformers đã trở thành encoder mặc định cho cả hệ thống CTC và RNN-T, vượt qua cả encoder Transformer thuần và LSTM thuần.
-
Whisper (Radford và đồng sự, 2023) là mô hình ASR dựa trên attention quy mô lớn của OpenAI. Nó dùng kiến trúc transformer encoder-decoder chuẩn (từ chương 07) được huấn luyện trên 680,000 giờ dữ liệu giám sát yếu (weakly supervised) thu thập từ internet (âm thanh ghép với các bản phiên âm xấp xỉ). Các lựa chọn thiết kế chính:
- Đầu vào: log-mel spectrogram 80 kênh (từ file 01) với cửa sổ 25 ms và bước nhảy 10 ms, được chuẩn hóa về trung bình zero và phương sai đơn vị.
- Encoder: encoder Transformer chuẩn với mã hóa vị trí sinusoidal và chuẩn hóa tầng tiền kích hoạt.
- Decoder: decoder Transformer tự hồi quy sinh ra token dùng bộ token hóa BPE mức byte (từ chương 07).
- Đa nhiệm: một mô hình duy nhất xử lý phiên âm, dịch thuật, nhận dạng ngôn ngữ, và dự đoán dấu thời gian, được điều kiện hóa bởi các token tác vụ đặc biệt trong phần mở đầu decoder.
- Quy mô của dữ liệu huấn luyện (chứ không phải sự đổi mới kiến trúc) là yếu tố chính tạo nên khả năng tổng quát hóa mạnh mẽ của Whisper trên nhiều lĩnh vực, giọng địa phương, và ngôn ngữ.
-
wav2vec 2.0 (Baevski và đồng sự, 2020) là một khung tự giám sát (self-supervised) cho biểu diễn tiếng nói. Ý tưởng cốt lõi là học các biểu diễn tiếng nói từ lượng lớn âm thanh không nhãn, sau đó tinh chỉnh với một lượng nhỏ dữ liệu có nhãn. Điều này theo cùng mô hình tự giám sát như BERT (từ chương 07) nhưng được điều chỉnh cho tín hiệu âm thanh liên tục.
-
Kiến trúc wav2vec 2.0 có ba phần:
- Feature encoder: một CNN 1D nhiều lớp xử lý các mẫu dạng sóng thô và tạo ra biểu diễn tiềm ẩn \(\mathbf{z}_t\) ở tốc độ khung 20 ms (một vector mỗi 320 mẫu ở 16 kHz).
- Module lượng tử hóa (quantisation module): rời rạc hóa các biểu diễn tiềm ẩn thành một bảng mã hữu hạn dùng lượng tử hóa tích (product quantisation — chia vector thành các nhóm và lượng tử hóa từng nhóm độc lập, chọn từ \(G\) bảng mã mỗi bảng \(V\) mục). Điều này tạo ra các mục tiêu \(\mathbf{q}_t\) cho hàm mục tiêu học đối lập (contrastive learning).
- Context network: một transformer encoder nhận các biểu diễn tiềm ẩn đã được che (partially masked) và tạo ra các biểu diễn ngữ cảnh hóa \(\mathbf{c}_t\).
- Trong quá trình tiền huấn luyện, các đoạn ngẫu nhiên của biểu diễn tiềm ẩn bị che (masked) (thay thế bằng một embedding mask được học), và mô hình phải xác định biểu diễn đã lượng tử hóa thực sự của vị trí bị che từ một tập các yếu tố gây nhiễu (distractors — các mẫu âm được lấy từ các vị trí khác trong cùng câu nói). Hàm mất mát đối lập là:
-
trong đó \(\text{sim}\) là độ tương tự cosine, \(\kappa\) là tham số nhiệt độ, và \(Q_t\) bao gồm mục tiêu đã lượng tử hóa thực sự cộng với các yếu tố gây nhiễu. Một hàm mất mát đa dạng (diversity loss) bổ sung khuyến khích việc sử dụng đồng đều tất cả các mục trong bảng mã. Hàm mất mát này về cơ bản là hàm mất mát đối lập InfoNCE, cùng họ với các hàm mục tiêu đối lập dùng trong học tự giám sát thị giác.
-
Sau tiền huấn luyện, một phép chiếu tuyến tính và đầu CTC được thêm lên trên, và mô hình được tinh chỉnh trên dữ liệu có nhãn. wav2vec 2.0 đạt kết quả gần như state-of-the-art với chỉ 10 phút dữ liệu có nhãn (dùng 53,000 giờ âm thanh không nhãn cho tiền huấn luyện), chứng minh sức mạnh của học tự giám sát cho nhận dạng tiếng nói ít tài nguyên.
-
HuBERT (Hsu và đồng sự, 2021) là một phương pháp tự giám sát khác thay thế hàm mục tiêu đối lập bằng một hàm mục tiêu dự đoán có che (masked prediction objective — dự đoán các gán nhãn cụm rời rạc của các khung bị che). Các mục tiêu được tạo ra bởi một bước phân cụm offline (k-means trên MFCC ở vòng lặp đầu tiên, sau đó k-means trên các đặc trưng HuBERT ở các vòng lặp tiếp theo). HuBERT đơn giản hóa pipeline huấn luyện so với wav2vec 2.0 (không cần module lượng tử hóa hay lấy mẫu đối lập) và đạt kết quả tương đương hoặc tốt hơn.
-
Fast Conformer (Rekesh và đồng sự, 2023, NVIDIA NeMo) thay thế self-attention bình phương (quadratic) trong Conformer chuẩn bằng một cơ chế chú ý hạ tốc (down-sampled attention): dãy đầu vào được nén (thường 8× qua tích chập có stride) trước khi tính attention, sau đó được mở rộng lại. Điều này giảm chi phí attention từ \(O(T^2)\) xuống \(O(T^2/64)\) trong khi vẫn giữ ngữ cảnh toàn cục, cho phép huấn luyện trên các câu nói rất dài (lên đến vài phút) mà không gặp vấn đề bộ nhớ. Fast Conformer là encoder mặc định trong bộ công cụ NeMo của NVIDIA và là xương sống cho các mô hình chất lượng sản xuất của họ.
-
Parakeet (NVIDIA, 2024) là một họ các mô hình ASR tiếng Anh độ chính xác cao xây dựng trên encoder Fast Conformer với bộ giải mã CTC và RNN-T, được huấn luyện trên 64,000 giờ tiếng Anh. Các mô hình Parakeet (0.6B và 1.1B tham số) đạt tỷ lệ lỗi từ thấp nhất trên các benchmark chuẩn tại thời điểm phát hành, vượt qua Whisper large-v3 trên hầu hết các bộ kiểm tra tiếng Anh. Các thành phần chính là kiến trúc Fast Conformer hiệu quả, tăng cường dữ liệu mạnh mẽ (SpecAugment, nhiễu tốc độ, pha trộn nhiễu), và dữ liệu huấn luyện có giám sát quy mô lớn — cho thấy rằng kỹ thuật cẩn thận các thành phần đã biết vẫn có thể đẩy mạnh state of the art.
-
Canary (NVIDIA, 2024) mở rộng khung NeMo cho ASR đa ngữ và đa nhiệm. Nó dùng encoder Fast Conformer với decoder dựa trên attention (thay vì CTC hay RNN-T) và xử lý cả phiên âm lẫn dịch thuật qua nhiều ngôn ngữ trong một mô hình duy nhất (tương tự thiết kế đa nhiệm của Whisper nhưng với backbone Fast Conformer hiệu quả hơn). Các mô hình Canary hỗ trợ tiếng Anh, Đức, Tây Ban Nha, và Pháp với độ chính xác cạnh tranh.
-
Moonshine (Useful Sensors, 2024) là một họ các mô hình ASR được tối ưu hóa đặc biệt cho triển khai trên thiết bị và biên (on-device and edge). Encoder dùng một kiến trúc hybrid thay thế các lớp transformer/conformer đầu bằng một CNN nhỏ theo sau bởi một vài lớp transformer, giảm mạnh kích thước mô hình (mô hình base dưới 30M tham số). Moonshine nhắm đến streaming thời gian thực trên CPU và các thiết bị công suất thấp nơi Whisper quá lớn và chậm, đánh đổi một ít độ chính xác để lấy độ trễ và bộ nhớ thấp hơn 5–10×.
-
Distil-Whisper (Gandhi và đồng sự, 2023) áp dụng chưng cất tri thức (knowledge distillation, chương 06) để nén Whisper thành một mô hình nhỏ hơn, nhanh hơn. Mô hình học sinh chỉ dùng 2 lớp decoder (so với 32 của Whisper) trong khi giữ nguyên encoder, và được huấn luyện để khớp với phân bố đầu ra của Whisper. Distil-Whisper đạt WER trong vòng 1% so với thầy trong khi nhanh hơn 6×, làm nó thực tế cho các ứng dụng thời gian thực nơi mô hình Whisper đầy đủ quá chậm.
-
Universal Speech Model (USM) (Zhang và đồng sự, 2023, Google) mở rộng tiền huấn luyện tự giám sát lên 12 triệu giờ âm thanh không nhãn trên 300+ ngôn ngữ, tiếp theo là tinh chỉnh có giám sát. USM cho thấy mô hình tự giám sát wav2vec 2.0 mở rộng được đến các chế độ dữ liệu thực sự lớn, đạt hiệu suất mạnh mẽ trên các ngôn ngữ ít tài nguyên với lượng dữ liệu có nhãn rất hạn chế.
-
Massively Multilingual Speech (MMS) (Pratap và đồng sự, 2023, Meta) mở rộng tiền huấn luyện wav2vec 2.0 lên hơn 1,100 ngôn ngữ dùng các bản ghi tôn giáo và các nguồn âm thanh đa ngữ khác. MMS bao phủ nhiều ngôn ngữ hơn bất kỳ hệ thống ASR nào trước đây, mở ra khả năng nhận dạng tiếng nói cho nhiều ngôn ngữ ít tài liệu lần đầu tiên.
-
Bức tranh toàn cảnh của ASR hiện đại đang hội tụ về một vài mẫu hình thống trị: (1) Encoder họ Conformer với CTC hoặc RNN-T cho streaming, (2) Transformer encoder-decoder cho offline/đa nhiệm, (3) Tiền huấn luyện tự giám sát cho các thiết lập ít tài nguyên, và (4) Quy mô — nhiều dữ liệu hơn và mô hình lớn hơn liên tục cải thiện độ chính xác. Sự lựa chọn giữa các phương án này phụ thuộc vào các ràng buộc triển khai: ngân sách độ trễ, tài nguyên tính toán, số lượng ngôn ngữ, và liệu ứng dụng là streaming hay batch.
-
Tích hợp mô hình ngôn ngữ cải thiện ASR bằng cách kết hợp tri thức ngôn ngữ học vượt xa những gì mô hình âm học nắm bắt được. Ý tưởng cơ bản là kết hợp điểm mô hình âm học \(p(\mathbf{x} | \mathbf{y})\) (mức độ phù hợp giữa âm thanh và bản phiên âm) với điểm mô hình ngôn ngữ \(p(\mathbf{y})\) (mức độ khả dĩ của bản phiên âm như một câu) trong quá trình giải mã.
-
Shallow fusion kết hợp các điểm tại thời gian tìm kiếm chùm:
-
trong đó \(\lambda\) là trọng số có thể tinh chỉnh và \(p_\text{LM}\) là một mô hình ngôn ngữ bên ngoài (thường là n-gram hoặc mạng nơ-ron LM từ chương 07). Phương pháp này đơn giản và hiệu quả, nhưng yêu cầu LM phải hoạt động trên cùng từ vựng token với mô hình ASR.
-
Deep fusion (Gulcehre và đồng sự, 2015) tích hợp mô hình ngôn ngữ bên trong mạng decoder: trạng thái ẩn của LM được nối với trạng thái ẩn của decoder và được đưa qua một cơ chế cổng trước phép chiếu đầu ra. Toàn bộ hệ thống (bao gồm LM đã tiền huấn luyện) được tinh chỉnh đồng thời. Điều này cho phép tích hợp sâu hơn nhưng phức tạp hơn để huấn luyện.
-
Cold fusion (Sriram và đồng sự, 2018) tương tự deep fusion nhưng huấn luyện decoder ASR từ đầu với mô hình ngôn ngữ được tích hợp, thay vì tinh chỉnh một decoder đã huấn luyện. Điều này buộc mô hình âm học phải học thông tin bổ sung chứ không chỉ sao chép những gì LM đã biết.
-
Rescoring (N-best rescoring) là phương pháp hai lượt: đầu tiên tạo ra \(N\) bản phiên âm ứng viên dùng beam search, sau đó xếp hạng lại chúng bằng một mô hình ngôn ngữ mạnh hơn (ví dụ, một transformer LM lớn). Cách này đơn giản để cài đặt và cho phép dùng các LM rất lớn vốn quá chậm cho giải mã lượt đầu.
-
Internal language model estimation (ILME) giải quyết một vấn đề tinh tế: các mô hình đầu-cuối ngầm học một LM nội bộ từ các bản phiên âm huấn luyện, điều này có thể xung đột với LM bên ngoài trong shallow fusion (về cơ bản là đếm hai lần tiên nghiệm ngôn ngữ). ILME ước lượng LM nội bộ và trừ điểm của nó trong quá trình fusion:
-
ASR streaming so với offline là một lựa chọn kiến trúc cơ bản. Offline (hay batch) ASR xử lý toàn bộ câu nói trước khi tạo ra bất kỳ đầu ra nào. Streaming ASR tạo ra đầu ra tăng dần khi âm thanh tới, với độ trễ có giới hạn.
-
Streaming là thiết yếu cho các ứng dụng thời gian thực: phụ đề trực tiếp, trợ lý giọng nói (người dùng mong đợi phản hồi trước khi họ nói xong), phiên âm cuộc gọi điện thoại. Thách thức là một số ngữ cảnh tương lai có ích cho nhận dạng (biết từ tiếp theo là "York" giúp phân biệt "New"), nhưng hệ thống streaming không thể chờ đợi ngữ cảnh tương lai dài tùy ý.
-
Các bộ mã hóa một chiều (unidirectional encoders — LSTM trái-sang-phải, tích chập nhân quả, transformers nhân quả) hỗ trợ streaming một cách tự nhiên vì mỗi đầu ra chỉ phụ thuộc vào đầu vào quá khứ và hiện tại. Các bộ mã hóa hai chiều (bidirectional encoders — nhìn vào ngữ cảnh tương lai) không hỗ trợ streaming trực tiếp.
-
Chunked attention (còn gọi là blockwise hay segmental attention) chia đầu vào thành các đoạn (chunk) có độ dài cố định và áp dụng self-attention chỉ trong mỗi đoạn (và tùy chọn với một vài đoạn trước đó). Điều này giới hạn độ trễ ở kích thước đoạn cộng thời gian xử lý, trong khi vẫn cho phép một ít ngữ cảnh hai chiều cục bộ trong mỗi đoạn. Sự đánh đổi là độ chính xác giảm khi kích thước đoạn nhỏ đi.
-
Lookahead cho phép một bộ mã hóa streaming nhìn trước một số lượng nhỏ khung tương lai (ví dụ, 300–900 ms) trước khi tạo ra đầu ra cho khung hiện tại. Điều này được cài đặt bằng cách thêm một ngữ cảnh phải nhỏ vào tính toán một chiều. Cửa sổ lookahead thêm độ trễ nhưng cải thiện đáng kể độ chính xác.
-
Độ trễ trong ASR streaming có vài thành phần:
- Độ trễ thuật toán (algorithmic latency): độ trễ từ khi âm thanh tới đến khi mô hình có thể xử lý nó (xác định bởi kích thước đoạn, lookahead, và trích xuất đặc trưng).
- Độ trễ tính toán (computational latency): thời gian chạy lượt forward của mô hình.
- Độ trễ endpoint (endpointer latency): độ trễ trong việc phát hiện người dùng đã nói xong.
- Độ trễ token đầu tiên (first-token latency): tốc độ xuất hiện từ đầu tiên. Độ trễ hoàn thiện (finalization latency): tốc độ đầu ra cuối được xác nhận (hệ thống streaming thường tạo ra đầu ra tạm thời có thể được sửa khi thêm âm thanh tới).
-
Các metric đánh giá cho ASR:
-
Tỷ lệ lỗi từ (Word Error Rate — WER) là metric chính. Nó được tính bằng cách căn chỉnh giả thuyết (đầu ra hệ thống) với tham chiếu (bản phiên âm gốc) dùng khoảng cách biên tập (edit distance — số thay thế, chèn, xóa tối thiểu để biến đổi cái này thành cái kia), sau đó:
-
trong đó \(S\) là số thay thế (substitutions), \(D\) là số xóa (deletions), \(I\) là số chèn (insertions), và \(N\) là tổng số từ trong tham chiếu. WER có thể vượt quá 100% nếu có nhiều chèn. WER 5% được coi là xấp xỉ mức người cho tiếng nói đọc rõ ràng; tiếng nói trò chuyện hoặc nhiễu khó hơn nhiều (10–20%+).
-
Tỷ lệ lỗi ký tự (Character Error Rate — CER) là công thức tương tự áp dụng ở mức ký tự thay vì mức từ. CER có nhiều thông tin hơn cho các ngôn ngữ không có ranh giới từ rõ ràng (tiếng Trung, tiếng Nhật) và để đánh giá các sai sót gần đúng ("cat" vs "bat" là 100% WER nhưng 33% CER).
-
Word Information Lost (WIL) và Word Information Preserved (WIP) là các phương án thay thế dựa trên lý thuyết thông tin tính toán mối tương quan giữa tham chiếu và giả thuyết chính xác hơn WER, nhưng chúng ít được báo cáo hơn.
-
Hệ số thời gian thực (Real-Time Factor — RTF) đo hiệu quả tính toán: tỷ lệ giữa thời gian xử lý và thời lượng âm thanh. RTF < 1 nghĩa là hệ thống chạy nhanh hơn thời gian thực; RTF > 1 nghĩa là nó không thể theo kịp âm thanh trực tiếp. Hệ thống streaming phải duy trì RTF < 1.
-
Tăng cường dữ liệu (data augmentation) là yếu tố then chốt cho ASR mạnh mẽ. Các kỹ thuật phổ biến:
- Nhiễu tốc độ (speed perturbation): lấy mẫu lại âm thanh ở tốc độ 0.9× và 1.1× (thay đổi cao độ và thời lượng).
- SpecAugment (Park và đồng sự, 2019): che (mask) các dải tần số và bước thời gian ngẫu nhiên trong spectrogram. Đây là phiên bản âm thanh của dropout và là một trong những kỹ thuật chính quy hóa hiệu quả nhất cho ASR. Nó không yêu cầu thêm dữ liệu.
- Tăng cường nhiễu (noise augmentation): trộn tiếng nói sạch với nhiễu ghi âm tại các tỷ lệ tín hiệu-trên-nhiễu khác nhau.
- Mô phỏng đáp ứng xung phòng (room impulse response simulation): tích chập tiếng nói sạch với âm học phòng mô phỏng để mô phỏng môi trường có tiếng vang.
-
Token hóa (tokenisation) cho ASR xác định từ vựng đầu ra của mô hình. Các tùy chọn bao gồm:
- Ký tự (characters): đơn giản, từ vựng nhỏ (~30 cho tiếng Anh), nhưng dãy đầu ra dài và không có mô hình hóa ngôn ngữ ngầm.
- Từ con / BPE (từ chương 07): các đơn vị dưới-từ cân bằng kích thước từ vựng và độ dài dãy. Tiêu chuẩn cho các hệ thống hiện đại (Whisper dùng BPE mức byte với ~50,000 token).
- Từ (words): từ vựng lớn (50,000+), dãy đầu ra ngắn, nhưng không thể xử lý các từ ngoài từ vựng.
- Âm vị (phonemes): dựa trên ngôn ngữ học, nhỏ gọn, nhưng yêu cầu một từ điển phát âm.
-
Sự tiến hóa của ASR có thể tóm tắt như một sự phát triển từ các hệ thống mô-đun được kỹ thuật nặng (GMM-HMM + WFST decoding, 1990–2010s) đến hệ thống hybrid (DNN-HMM, 2012–2016) đến hệ thống đầu-cuối hấp thụ ngày càng nhiều pipeline vào một mạng nơ-ron duy nhất (CTC, RNN-T, LAS, 2016–2020) đến các mô hình tiền huấn luyện quy mô lớn tận dụng lượng lớn dữ liệu không nhãn hoặc nhãn yếu (wav2vec 2.0, Whisper, 2020–nay). Mỗi bước chuyển tiếp đơn giản hóa kỹ thuật đồng thời cải thiện độ chính xác, theo xu hướng rộng hơn trong học máy: học biểu diễn từ dữ liệu thay vì thiết kế tay (cùng câu chuyện đã kể trong chương 06 cho đặc trưng ảnh được thay thế bởi CNN, và trong chương 07 cho đặc trưng NLP được thay thế bởi transformer).
Bài tập lập trình (dùng CoLab hoặc notebook)¶
-
Cài đặt CTC loss từ đầu trong JAX. Tạo một ví dụ đồ chơi với một dãy logits ngắn và một nhãn đích, tính thuật toán forward CTC để có tổng xác suất, và tính negative log-likelihood loss.
import jax import jax.numpy as jnp import matplotlib.pyplot as plt def ctc_forward(log_probs, targets): """ CTC forward algorithm (log-domain for numerical stability). log_probs: (T, V) log probabilities over vocabulary (index 0 = blank) targets: (U,) target label indices (no blanks) Returns: log probability of the target sequence under CTC. """ T, V = log_probs.shape U = len(targets) # Build the extended label sequence with blanks: [blank, y1, blank, y2, ..., yU, blank] S = 2 * U + 1 labels = jnp.zeros(S, dtype=jnp.int32) # all blanks for i in range(U): labels = labels.at[2 * i + 1].set(targets[i]) # Initialise alpha (log domain) NEG_INF = -1e30 alpha = jnp.full((T, S), NEG_INF) alpha = alpha.at[0, 0].set(log_probs[0, labels[0]]) # start with blank alpha = alpha.at[0, 1].set(log_probs[0, labels[1]]) # or first label # Fill forward for t in range(1, T): for s in range(S): # Same state a = alpha[t - 1, s] # From previous state if s > 0: a = jnp.logaddexp(a, alpha[t - 1, s - 1]) # Skip blank (if current and two-back labels are different) if s > 1 and labels[s] != 0 and labels[s] != labels[s - 2]: a = jnp.logaddexp(a, alpha[t - 1, s - 2]) alpha = alpha.at[t, s].set(a + log_probs[t, labels[s]]) # Total log probability: sum of last two states at final time step log_prob = jnp.logaddexp(alpha[T - 1, S - 1], alpha[T - 1, S - 2]) return log_prob, alpha # --- Toy example --- T = 12 # input length (time steps) V = 5 # vocab size (0=blank, 1='c', 2='a', 3='t', 4='x') targets = jnp.array([1, 2, 3]) # "c", "a", "t" # Create random logits and convert to log-probabilities key = jax.random.PRNGKey(42) logits = jax.random.normal(key, (T, V)) log_probs = jax.nn.log_softmax(logits, axis=-1) log_prob, alpha = ctc_forward(log_probs, targets) ctc_loss = -log_prob print(f"Target sequence: {targets.tolist()} ('c', 'a', 't')") print(f"Input length T={T}, Vocab size V={V}") print(f"CTC log-probability: {log_prob:.4f}") print(f"CTC loss (neg log-prob): {ctc_loss:.4f}") # Visualise the forward variable (alpha) lattice fig, ax = plt.subplots(figsize=(12, 5)) # Convert from log to linear for visualisation alpha_linear = jnp.exp(alpha - jnp.max(alpha)) # normalise for visibility im = ax.imshow(alpha_linear.T, aspect='auto', origin='lower', cmap='viridis') ax.set_xlabel('Time step (t)') ax.set_ylabel('Extended label index (s)') label_names = ['_', 'c', '_', 'a', '_', 't', '_'] # _ = blank ax.set_yticks(range(len(label_names))) ax.set_yticklabels(label_names) ax.set_title(f'CTC Forward Variable (alpha lattice) | Loss = {ctc_loss:.2f}') plt.colorbar(im, ax=ax, label='Normalised probability') plt.tight_layout(); plt.show() -
Xây dựng một mô hình ASR encoder-decoder dựa trên attention đơn giản (kiến trúc LAS thu nhỏ) trong JAX. Dùng encoder tích chập 1D và decoder một lớp với dot-product attention. Chạy trên dữ liệu tổng hợp và trực quan hóa các trọng số attention.
import jax import jax.numpy as jnp import matplotlib.pyplot as plt # --- Minimal attention-based encoder-decoder for ASR --- def init_params(key, input_dim, hidden_dim, vocab_size): """Initialise parameters for a tiny LAS-like model.""" keys = jax.random.split(key, 8) scale = 0.1 params = { # Encoder: simple linear projection (simulating conv output) 'enc_w': jax.random.normal(keys[0], (input_dim, hidden_dim)) * scale, 'enc_b': jnp.zeros(hidden_dim), # Attention: query, key, value projections 'attn_q': jax.random.normal(keys[1], (hidden_dim, hidden_dim)) * scale, 'attn_k': jax.random.normal(keys[2], (hidden_dim, hidden_dim)) * scale, 'attn_v': jax.random.normal(keys[3], (hidden_dim, hidden_dim)) * scale, # Decoder RNN (simple Elman RNN for illustration) 'dec_wh': jax.random.normal(keys[4], (hidden_dim, hidden_dim)) * scale, 'dec_wx': jax.random.normal(keys[5], (vocab_size, hidden_dim)) * scale, 'dec_wc': jax.random.normal(keys[6], (hidden_dim, hidden_dim)) * scale, 'dec_b': jnp.zeros(hidden_dim), # Output projection 'out_w': jax.random.normal(keys[7], (hidden_dim, vocab_size)) * scale, 'out_b': jnp.zeros(vocab_size), } return params def encode(params, x): """Encoder: linear projection (placeholder for conv/LSTM stack).""" return jnp.tanh(x @ params['enc_w'] + params['enc_b']) def attend(params, query, enc_out): """Dot-product attention over encoder outputs.""" q = query @ params['attn_q'] # (hidden,) k = enc_out @ params['attn_k'] # (T_enc, hidden) v = enc_out @ params['attn_v'] # (T_enc, hidden) d_k = q.shape[-1] scores = (k @ q) / jnp.sqrt(d_k) # (T_enc,) weights = jax.nn.softmax(scores) # (T_enc,) context = weights @ v # (hidden,) return context, weights def decode_step(params, h_prev, y_prev_onehot, enc_out): """Single decoder step: RNN + attention.""" # Embed previous token y_emb = y_prev_onehot @ params['dec_wx'] # (hidden,) # Attend to encoder context, attn_w = attend(params, h_prev, enc_out) # RNN update h = jnp.tanh(h_prev @ params['dec_wh'] + y_emb + context @ params['dec_wc'] + params['dec_b']) # Output logits logits = h @ params['out_w'] + params['out_b'] return h, logits, attn_w # --- Setup --- key = jax.random.PRNGKey(0) input_dim = 40 # e.g., 40 mel bands hidden_dim = 64 vocab_size = 10 # small vocab for demo T_enc = 30 # encoder time steps T_dec = 8 # decoder steps params = init_params(key, input_dim, hidden_dim, vocab_size) # Synthetic input: random mel-like features key, subkey = jax.random.split(key) x = jax.random.normal(subkey, (T_enc, input_dim)) # Encode enc_out = encode(params, x) # Decode (teacher forcing with random targets) key, subkey = jax.random.split(key) targets = jax.random.randint(subkey, (T_dec,), 0, vocab_size) h = jnp.zeros(hidden_dim) all_logits = [] all_attn = [] for t in range(T_dec): y_prev = jax.nn.one_hot(targets[t] if t > 0 else 0, vocab_size) h, logits, attn_w = decode_step(params, h, y_prev, enc_out) all_logits.append(logits) all_attn.append(attn_w) all_attn = jnp.stack(all_attn) # (T_dec, T_enc) all_logits = jnp.stack(all_logits) # (T_dec, vocab_size) # --- Visualise attention weights --- fig, axes = plt.subplots(1, 2, figsize=(14, 5)) im = axes[0].imshow(all_attn, aspect='auto', cmap='Blues', origin='lower') axes[0].set_xlabel('Encoder time step') axes[0].set_ylabel('Decoder step') axes[0].set_title('Attention Weights (decoder -> encoder)') plt.colorbar(im, ax=axes[0]) # Show predicted token distribution for each decoder step im2 = axes[1].imshow(jax.nn.softmax(all_logits, axis=-1), aspect='auto', cmap='Oranges', origin='lower') axes[1].set_xlabel('Vocabulary index') axes[1].set_ylabel('Decoder step') axes[1].set_title('Output Token Probabilities') plt.colorbar(im2, ax=axes[1]) plt.suptitle('Minimal Attention-based ASR Model (untrained)') plt.tight_layout(); plt.show() -
Tính Word Error Rate (WER) từ đầu dùng quy hoạch động (khoảng cách biên tập), và đánh giá nhiều giả thuyết so với một tham chiếu. Trực quan hóa ma trận khoảng cách biên tập.
import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np def compute_wer(reference, hypothesis): """ Compute WER using dynamic programming (Levenshtein distance on words). Returns WER, number of substitutions, deletions, insertions, and the DP matrix. """ ref_words = reference.split() hyp_words = hypothesis.split() N = len(ref_words) M = len(hyp_words) # DP matrix: d[i][j] = edit distance between ref[:i] and hyp[:j] d = np.zeros((N + 1, M + 1), dtype=np.int32) # Backtrack matrix to count S, D, I ops = np.zeros((N + 1, M + 1, 3), dtype=np.int32) # [sub, del, ins] for i in range(N + 1): d[i][0] = i # all deletions for j in range(M + 1): d[0][j] = j # all insertions for i in range(1, N + 1): for j in range(1, M + 1): if ref_words[i - 1] == hyp_words[j - 1]: sub_cost = d[i - 1][j - 1] # match, no edit else: sub_cost = d[i - 1][j - 1] + 1 # substitution del_cost = d[i - 1][j] + 1 # deletion ins_cost = d[i][j - 1] + 1 # insertion d[i][j] = min(sub_cost, del_cost, ins_cost) # Backtrack to count operations i, j = N, M S, D, I = 0, 0, 0 while i > 0 or j > 0: if i > 0 and j > 0 and d[i][j] == d[i-1][j-1] and ref_words[i-1] == hyp_words[j-1]: i -= 1; j -= 1 # correct elif i > 0 and j > 0 and d[i][j] == d[i-1][j-1] + 1: S += 1; i -= 1; j -= 1 # substitution elif i > 0 and d[i][j] == d[i-1][j] + 1: D += 1; i -= 1 # deletion elif j > 0 and d[i][j] == d[i][j-1] + 1: I += 1; j -= 1 # insertion else: break wer = (S + D + I) / N if N > 0 else 0.0 return wer, S, D, I, d # --- Test cases --- reference = "the cat sat on the mat" hypotheses = [ "the cat sat on the mat", # perfect "the cat sit on the mat", # 1 substitution "the cat on the mat", # 1 deletion "the big cat sat on the mat", # 1 insertion "a dog sat in a rug", # multiple errors ] print(f"Reference: '{reference}'\n") print(f"{'Hypothesis':<40s} {'WER':>6s} {'S':>3s} {'D':>3s} {'I':>3s}") print("-" * 60) results = [] for hyp in hypotheses: wer, S, D, I, dp = compute_wer(reference, hyp) results.append((hyp, wer, S, D, I, dp)) print(f"'{hyp}':<40s} {wer:>6.1%} {S:>3d} {D:>3d} {I:>3d}") # Visualise the DP matrix for the worst case worst = results[-1] hyp_words = worst[0].split() ref_words = reference.split() dp_matrix = worst[5] fig, axes = plt.subplots(1, 2, figsize=(14, 5)) # DP matrix im = axes[0].imshow(dp_matrix, cmap='YlOrRd', origin='upper') axes[0].set_xticks(range(len(hyp_words) + 1)) axes[0].set_xticklabels([''] + hyp_words, rotation=45, ha='right', fontsize=9) axes[0].set_yticks(range(len(ref_words) + 1)) axes[0].set_yticklabels([''] + ref_words, fontsize=9) axes[0].set_xlabel('Hypothesis words') axes[0].set_ylabel('Reference words') axes[0].set_title(f'Edit Distance Matrix\nWER = {worst[1]:.1%}') for i in range(dp_matrix.shape[0]): for j in range(dp_matrix.shape[1]): axes[0].text(j, i, str(dp_matrix[i, j]), ha='center', va='center', fontsize=8) plt.colorbar(im, ax=axes[0]) # WER comparison bar chart names = [f'Hyp {i+1}' for i in range(len(results))] wers = [r[1] * 100 for r in results] colors = ['#27ae60' if w == 0 else '#f39c12' if w < 30 else '#e74c3c' for w in wers] axes[1].barh(names, wers, color=colors) axes[1].set_xlabel('WER (%)') axes[1].set_title('Word Error Rate Comparison') for i, (w, r) in enumerate(zip(wers, results)): axes[1].text(w + 1, i, f'{w:.0f}% (S={r[2]}, D={r[3]}, I={r[4]})', va='center', fontsize=9) axes[1].set_xlim(0, max(wers) * 1.4) plt.tight_layout(); plt.show() -
Cài đặt SpecAugment (che tần số và che thời gian) trên một log-mel spectrogram và trực quan hóa các phiên bản gốc so với được tăng cường. Sinh spectrogram từ một tín hiệu tổng hợp.
import jax import jax.numpy as jnp import matplotlib.pyplot as plt # --- Generate synthetic log-mel spectrogram --- key = jax.random.PRNGKey(42) fs = 16000 duration = 2.0 t = jnp.arange(0, duration, 1.0 / fs) # Simulate speech: chirp signal with harmonics f0 = 120.0 x = sum(jnp.sin(2 * jnp.pi * f0 * k * t * (1 + 0.1 * t)) / k for k in range(1, 10)) key, subkey = jax.random.split(key) x = x + 0.05 * jax.random.normal(subkey, t.shape) # Compute log-mel spectrogram (simplified) frame_len = 400 # 25 ms hop_len = 160 # 10 ms n_fft = 512 n_mels = 80 n_frames = (len(x) - frame_len) // hop_len + 1 hamming = 0.54 - 0.46 * jnp.cos(2 * jnp.pi * jnp.arange(frame_len) / (frame_len - 1)) frames = jnp.stack([x[i * hop_len : i * hop_len + frame_len] for i in range(n_frames)]) windowed = frames * hamming spectra = jnp.abs(jnp.fft.rfft(windowed, n=n_fft)) ** 2 # Simple mel filterbank def hz_to_mel(f): return 2595 * jnp.log10(1 + f / 700) def mel_to_hz(m): return 700 * (10 ** (m / 2595) - 1) mel_points = jnp.linspace(hz_to_mel(0), hz_to_mel(fs / 2), n_mels + 2) hz_pts = mel_to_hz(mel_points) bins = jnp.floor((n_fft + 1) * hz_pts / fs).astype(jnp.int32) n_freqs = n_fft // 2 + 1 fb = jnp.zeros((n_mels, n_freqs)) for m in range(n_mels): lo, mid, hi = int(bins[m]), int(bins[m+1]), int(bins[m+2]) for k in range(lo, mid): if mid != lo: fb = fb.at[m, k].set((k - lo) / (mid - lo)) for k in range(mid, hi): if hi != mid: fb = fb.at[m, k].set((hi - k) / (hi - mid)) log_mel = jnp.log(spectra @ fb.T + 1e-10) # --- SpecAugment --- def spec_augment(spec, key, n_freq_masks=2, freq_mask_width=15, n_time_masks=2, time_mask_width=25): """Apply SpecAugment: frequency and time masking.""" augmented = spec.copy() T, F = spec.shape # Frequency masking for _ in range(n_freq_masks): key, k1, k2 = jax.random.split(key, 3) f_width = jax.random.randint(k1, (), 1, freq_mask_width + 1) f_start = jax.random.randint(k2, (), 0, max(1, F - freq_mask_width)) mask = (jnp.arange(F) >= f_start) & (jnp.arange(F) < f_start + f_width) augmented = jnp.where(mask[None, :], 0.0, augmented) # Time masking for _ in range(n_time_masks): key, k1, k2 = jax.random.split(key, 3) t_width = jax.random.randint(k1, (), 1, time_mask_width + 1) t_start = jax.random.randint(k2, (), 0, max(1, T - time_mask_width)) mask = (jnp.arange(T) >= t_start) & (jnp.arange(T) < t_start + t_width) augmented = jnp.where(mask[:, None], 0.0, augmented) return augmented key, subkey = jax.random.split(key) log_mel_aug = spec_augment(log_mel, subkey) # --- Visualise --- fig, axes = plt.subplots(2, 1, figsize=(14, 8)) im0 = axes[0].imshow(log_mel.T, aspect='auto', origin='lower', cmap='inferno', extent=[0, duration, 0, n_mels]) axes[0].set_title('Original Log-Mel Spectrogram') axes[0].set_xlabel('Time (s)'); axes[0].set_ylabel('Mel Band') plt.colorbar(im0, ax=axes[0], label='Log Energy') im1 = axes[1].imshow(log_mel_aug.T, aspect='auto', origin='lower', cmap='inferno', extent=[0, duration, 0, n_mels]) axes[1].set_title('After SpecAugment (frequency + time masking)') axes[1].set_xlabel('Time (s)'); axes[1].set_ylabel('Mel Band') plt.colorbar(im1, ax=axes[1], label='Log Energy') plt.tight_layout(); plt.show()