Phân tích người nói và âm thanh

Phân tích người nói và âm thanh xác định ai đang nói, họ nói lúc nào, và những âm thanh nào không phải lời nói đang hiện diện. File này bao quát xác thực người nói và nhận dạng người nói, i-vector, d-vector, x-vector, phân tách người nói (diarisation), phân loại sự kiện âm thanh, truy xuất thông tin nhạc, và nhận dạng cảm xúc từ giọng nói.

• Trong file 01, ta đã xây dựng nền tảng xử lý tín hiệu: spectrogram, MFCC, và dải lọc mel. Trong file 02, ta đã nhận dạng được nói gì. Giờ đây ta hỏi ai đã nói, họ nói lúc nào, và còn chuyện gì khác đang diễn ra trong đoạn âm thanh. Nhận dạng người nói, phân tách người nói, phân loại âm thanh, và phân tích nhạc đều chia sẻ một điểm chung: học các embedding nhỏ gọn nắm bắt được những tính bất biến phù hợp cho từng tác vụ, gợi nhớ các ý tưởng về embedding từ chương 06.

• Hãy hình dung việc nhận dạng một người nói giống như nhận ra giọng của một người bạn qua điện thoại. Bạn không cần hiểu từ ngữ; cái gì đó về âm sắc, nhịp độ, và chất giọng là duy nhất của người đó. Các hệ thống nhận dạng người nói học cách trích xuất chính cái "dấu vân tay giọng nói" này từ âm thanh thô, bỏ qua nội dung nói và tập trung vào cách nói.

• Nhận dạng người nói (speaker recognition) là thuật ngữ bao trùm cho hai tác vụ liên quan:

  • Xác thực người nói (speaker verification, SV): cho trước một danh tính được claim và một đoạn âm thanh, xác định xem người nói có đúng là người họ claim hay không. Đây là một quyết định nhị phân (chấp nhận hoặc từ chối) và là công nghệ đứng sau xác thực bằng giọng nói ("Hey Siri, đây có phải giọng của tôi không?").
  • Nhận dạng người nói (speaker identification, SI): cho trước một đoạn âm thanh và một tập hợp (gallery) các người nói đã biết, xác định người nói nào đã tạo ra đoạn đó. Đây là một bài toán phân lớp đa lớp.

• Cả hai tác vụ đều chia sẻ cùng một biểu diễn cơ bản: một embedding người nói có số chiều cố định, nắm bắt danh tính của người nói bất kể họ nói gì. Sự khác biệt chỉ nằm ở giai đoạn quyết định: xác thực so sánh hai embedding, còn nhận dạng tìm embedding gần nhất trong số các ứng viên.

• Độ tương đồng cosine (cosine similarity) là thước đo chuẩn để so sánh các embedding người nói. Cho embedding đăng ký \(e\) và embedding kiểm thử \(t\):

\[s = \frac{e \cdot t}{\|e\| \, \|t\|}\]

• Một ngưỡng \(\theta\) quyết định việc chấp nhận/từ chối: nếu \(s > \theta\) thì chấp nhận. Ngưỡng này đánh đổi giữa tỉ lệ chấp nhận sai (FAR) và tỉ lệ từ chối sai (FRR). Tỉ lệ lỗi bằng nhau (EER), tức khi FAR = FRR, là thước đo đánh giá chuẩn. EER càng thấp thì hiệu năng càng tốt. Các hệ thống tiên tiến nhất đạt EER dưới 1% trên các benchmark chuẩn (VoxCeleb).

• i-vector (Dehak et al., 2010) là embedding người nói thống trị trước kỷ nguyên học sâu. Ý tưởng đến từ phân tích nhân tố (phân rã ma trận ở chương 02 và giảm số chiều ở chương 04). Một mô hình nền chung (universal background model, UBM), tức một GMM lớn được huấn luyện trên nhiều người nói đa dạng, định nghĩa một không gian siêu vector (supervector space). Siêu vector GMM của mỗi câu nói được chiếu vào một không gian biến thiên tổng (total variability space) có số chiều thấp:

\[M = m + Tw\]

• trong đó \(M\) là siêu vector GMM của câu nói, \(m\) là siêu vector trung bình UBM, \(T\) là ma trận biến thiên tổng (học từ dữ liệu), và \(w\) là i-vector, một biểu diễn có số chiều thấp (thường 400-600) nắm bắt cả biến thiên của người nói lẫn của kênh.

• Để loại bỏ biến thiên kênh khỏi i-vector, Phân tích phân biệt tuyến tính xác suất (Probabilistic Linear Discriminant Analysis, PLDA) mô hình hóa i-vector như một tổng của các biến latent đặc trưng cho người nói và đặc trưng cho kênh. PLDA cung cấp một điểm số log-likelihood ratio có căn bản lý thuyết cho việc xác thực:

\[\text{score}(w_1, w_2) = \log \frac{P(w_1, w_2 \mid \text{same speaker})}{P(w_1 \mid \text{speaker}_1) \, P(w_2 \mid \text{speaker}_2)}\]

• d-vector (Variani et al., 2014) là các embedding người nói nơ-ron đầu tiên. Một DNN được huấn luyện để phân lớp người nói trên các đặc trưng cấp frame trích xuất một biểu diễn có số chiều cố định bằng cách lấy trung bình các kích hoạt của lớp ẩn cuối trên toàn bộ các frame trong một câu nói. Đơn giản nhưng hiệu quả, d-vector đã chứng minh rằng mạng nơ-ron có thể học các đặc trưng phân biệt người nói mà không cần bộ máy thống kê phức tạp của i-vector.

• x-vector (Snyder et al., 2018) đã tiến bộ vượt bậc các embedding người nói nơ-ron nhờ kiến trúc Mạng nơ-ron trễ theo thời gian (Time Delay Neural Network, TDNN). TDNN là các phép tích chập 1D với các cửa sổ ngữ cảnh cụ thể ở mỗi lớp, liên quan đến các phép tích chập giãn (dilated convolutions) từ WaveNet ở file 03 nhưng áp dụng lên các đặc trưng cấp frame thay vì các mẫu sóng thô.

• Kiến trúc x-vector có ba giai đoạn:
  • Các lớp cấp frame: một chồng các lớp TDNN xử lý MFCC (từ file 01) với ngữ cảnh thời gian mở rộng dần. Mỗi lớp nhìn một cửa sổ ngữ cảnh cố định (ví dụ \(\{t-2, t-1, t, t+1, t+2\}\) cho lớp đầu, rộng hơn ở các lớp sau).
  • Statistics pooling: sau các lớp cấp frame, trung bình và độ lệch chuẩn của các đầu ra cấp frame được tính trên toàn bộ câu nói, tạo ra một vector có số chiều cố định bất kể độ dài câu nói:

\[ \begin{aligned} \mu &= \frac{1}{T} \sum_{t=1}^{T} h_t \\ \sigma &= \sqrt{\frac{1}{T} \sum_{t=1}^{T} (h_t - \mu)^2} \end{aligned} \]

• trong đó \(h_t\) là đầu ra cấp frame tại thời điểm \(t\). Phép nối \([\\mu; \\sigma]\) là biểu diễn đã gộp (pooled).

  • Các lớp cấp đoạn: các lớp kết nối đầy đủ xử lý biểu diễn đã gộp. Đầu ra của lớp cấp đoạn đầu tiên (trước softmax) chính là embedding x-vector.

• x-vector được huấn luyện với hàm mất mát cross-entropy chuẩn trên các danh tính người nói. Mặc dù được huấn luyện để phân lớp, biểu diễn trung gian học được (x-vector) khái quát tốt sang những người nói chưa thấy vì mạng học cách trích xuất các đặc trưng phân biệt người nói thay vì ghi nhớ những người nói cụ thể.

• ECAPA-TDNN (Desplanques et al., 2020) là kiến trúc dựa trên TDNN hiện đại nhất cho nhận dạng người nói. Nó đưa vào ba cải tiến so với x-vector:

  • Khối Squeeze-Excitation (SE): sự chú ý kênh (từ SENet ở chương 08) giúp cân trọng lại các kênh đặc trưng dựa trên ngữ cảnh toàn cục, cho phép mô hình nhấn mạnh các kênh liên quan đến người nói.
  • Đặc trưng đa tỷ lệ kiểu Res2Net: trong mỗi khối TDNN, các kênh được chia thành các nhóm được xử lý phân cấp, tạo ra các đặc trưng ở nhiều độ phân giải thời gian (tương tự trích xuất đa tỷ lệ ở chương 08).
  • Statistics pooling có chú ý: thay vì trung bình cân bằng, một cơ chế chú ý cân trọng đóng góp của mỗi frame vào thống kê đã gộp. Các frame có nội dung phân biệt người nói hơn (ví dụ nguyên âm, mang nhiều thông tin người nói hơn) nhận trọng số chú ý cao hơn:

\[\alpha_t = \frac{\exp(v^T f(h_t))}{\sum_{\tau} \exp(v^T f(h_\tau))}\]

• trong đó \(f\) là một mạng nơ-ron nhỏ và \(v\) là vector chú ý được học. Trung bình và độ lệch chuẩn có chú ý trở thành \(\tilde{\mu} = \sum_t \alpha_t h_t\) và \(\tilde{\sigma} = \sqrt{\sum_t \alpha_t (h_t - \tilde{\mu})^2}\).

• ECAPA-TDNN thường được huấn luyện với AAM-Softmax (Additive Angular Margin Softmax), thêm một hình phạt biên góc (angular margin) vào hàm mất mát phân lớp, đẩy các embedding của cùng một người nói xích lại gần nhau và các người nói khác xa nhau trên siêu cầu:

\[L = -\log \frac{e^{s \cos(\theta_{y_i} + m)}}{e^{s \cos(\theta_{y_i} + m)} + \sum_{j \neq y_i} e^{s \cos \theta_j}}\]

• trong đó \(\theta_{y_i}\) là góc giữa embedding và vector trọng số của lớp đúng, \(m\) là biên (thường 0,2), và \(s\) là hệ số tỷ lệ (thường 30). Hàm mất mát này xuất phát từ nhận dạng khuôn mặt (ArcFace ở chương 08) và rất hiệu quả cho xác thực người nói.

• Phân tách người nói (speaker diarisation) trả lời câu hỏi "ai nói khi nào" trong một bản ghi có nhiều người nói. Hãy hình dung nó như việc tô màu một dòng thời gian: mỗi màu đại diện cho một người nói khác nhau, và hệ thống phải xác định khi nào mỗi người nói hoạt động, kể cả khi có tiếng nói chồng lấp.

• Phân tách dựa trên phân cụm (clustering-based diarisation) là phương pháp đường ống truyền thống:

  • Phân đoạn: chia âm thanh thành các đoạn ngắn (thường 1-2 giây) dùng cửa sổ trượt hoặc phát hiện thay đổi người nói.
  • Trích xuất embedding: trích xuất một embedding người nói (x-vector, ECAPA-TDNN) cho mỗi đoạn.
  • Phân cụm: nhóm các đoạn theo người nói. Phân cụm phân cấp tham lam (Agglomerative Hierarchical Clustering, AHC) là chuẩn: bắt đầu với mỗi đoạn là một cụm riêng, rồi lặp lại gộp hai cụm giống nhau nhất cho đến khi thỏa tiêu chí dừng (dựa trên ngưỡng khoảng cách hoặc số người nói mục tiêu).
  • Phân đoạn lại: tinh chỉnh biên bằng một sự căn chỉnh lại dựa trên Viterbi.

• Số lượng người nói thường không biết trước (a priori), điều này làm bài toán khó hơn phân cụm chuẩn. Phân cụm phổ (spectral clustering) với ngưỡng xác định \(k\) dựa trên trị riêng là một cách tiếp cận phổ biến khác.

• Phân tách nơ-ron đầu cuối (End-to-end neural diarisation, EEND) (Fujita et al., 2019) xem diarisation như một bài toán phân lớp đa nhãn. Một mạng nơ-ron (thường là mô hình dựa trên tự chú ý, transformer ở chương 07) lấy toàn bộ bản ghi làm đầu vào và xuất ra một nhãn hoạt động nhị phân cho mỗi người nói tại mỗi frame. Cách này xử lý trực tiếp tiếng nói chồng lấp, vốn là điểm yếu lớn của các phương pháp dựa trên phân cụm.

• Đầu ra EEND cho \(S\) người nói tại frame \(t\) là:

\[\hat{y}_{t,s} = \sigma(f_s(h_t))\]

• trong đó \(h_t\) là đầu ra transformer tại frame \(t\) và \(f_s\) là một phép chiếu tuyến tính cho người nói \(s\). Hàm mất mát huấn luyện là binary cross-entropy tổng trên người nói và frame. Một thách thức then chốt là số người nói phải cố định hoặc được xử lý bằng một kiến trúc đầu ra biến đổi (EEND-EDA dùng encoder-decoder với các attractor).

• Huấn luyện bất biến hoán vị (Permutation invariant training, PIT) cho diarisation xử lý bài toán nhập nhằng nhãn: vì người nói không có thứ tự cố hữu, mất mát được tính cho mọi phép gán người nói-với-đầu ra có thể có và lấy giá trị nhỏ nhất (đây là cùng PIT dùng trong tách nguồn, xem ở file 05).

• Phân loại âm thanh (audio classification) gán một nhãn cho toàn bộ một đoạn âm thanh. Khác với ASR (file 02) chuyển lời nói thành văn bản, phân loại âm thanh bao phủ phạm vi rộng hơn: âm thanh môi trường (còi hú, mưa, tiếng chó sủa), thể loại nhạc (rock, jazz, cổ điển), và các sự kiện âm thanh nói chung.

• Cách tiếp cận chuẩn đi theo hình mẫu phân lớp ảnh từ chương 08: biểu diễn âm thanh dưới dạng spectrogram (một ảnh 2D thời gian-tần số), rồi áp dụng một bộ phân lớp CNN hoặc transformer. Cách tiếp cận ảnh phổ này tận dụng hàng thập kỷ tiến bộ của thị giác máy tính.

• Phân loại âm thanh môi trường (Environmental sound classification, ESC) dùng các tập dữ liệu như ESC-50 (50 lớp, 2000 đoạn) và UrbanSound8K. Các kiến trúc điển hình là CNN (chương 06) áp dụng lên log-mel spectrogram. Tăng cường dữ liệu rất quan trọng: kéo dãn thời gian, dịch chuyển cao độ, thêm nhiễu nền, và SpecAugment (cách che mặt nạ từ file 02 áp dụng lên spectrogram) đều cải thiện khái quát hóa.

• Phát hiện sự kiện âm thanh (Audio event detection, SED) là hình thức tương tự theo thời gian của phân lớp: không chỉ sự kiện nào hiện diện, mà còn chúng bắt đầu và kết thúc lúc nào. AudioSet (Gemmeke et al., 2017) là benchmark quy mô lớn với 527 lớp sự kiện và hơn 2 triệu đoạn 10 giây trích từ YouTube, mỗi đoạn được gán nhãn yếu (nhãn cấp đoạn, không phải cấp frame).

• SED yếu (weakly-supervised SED) phải học các dự đoán cấp frame từ các nhãn cấp đoạn. Cách tiếp cận chuẩn dùng một CNN sinh ra xác suất lớp cấp frame, rồi tổng hợp chúng thành dự đoán cấp đoạn qua attention pooling:

\[\hat{Y}_c = \sigma\left(\sum_t \alpha_{t,c} \cdot f_{t,c}\right)\]

• trong đó \(f_{t,c}\) là logit cấp frame cho lớp \(c\) tại thời điểm \(t\), và \(\alpha_{t,c}\) là trọng số chú ý. Dự đoán cấp đoạn \(\hat{Y}_c\) được huấn luyện đối với nhãn cấp đoạn.

• Phân loại cảnh âm học (Acoustic scene classification, ASC) phân loại toàn bộ môi trường: "sân bay", "công viên", "ga tàu điện ngầm", "văn phòng". Đây là một tác vụ tổng thể: mô hình phải nắm bắt kết cấu âm học chung thay vì các sự kiện cụ thể. Chuỗi thử thách DCASE đánh giá ASC hàng năm, với các hệ thống thắng thường dùng tập hợp (ensemble) các CNN trên spectrogram đa độ phân giải.

• Embedding âm thanh (audio embeddings) là các biểu diễn đa dụng học từ dữ liệu âm thanh quy mô lớn, tương tự như embedding từ (chương 07) hoặc đặc trưng ảnh (chương 08) có thể chuyển sang các tác vụ hạ nguồn.

• VGGish (Hershey et al., 2017) thích ứng mạng phân lớp ảnh VGG (chương 08) sang âm thanh. Nó xử lý các mảng log-mel spectrogram 0,96 giây qua một CNN kiểu VGG được tiền huấn luyện trên AudioSet, sinh ra một embedding 128 chiều cho mỗi mảng. Embedding VGGish đóng vai trò là đặc trưng âm thanh đa dụng cho các tác vụ hạ nguồn, tương tự như cách các CNN tiền huấn luyện ImageNet cung cấp đặc trưng thị giác.

• PANNs (Pre-trained Audio Neural Networks, Kong et al., 2020) là một họ các kiến trúc CNN (CNN6, CNN10, CNN14) được huấn luyện trên toàn bộ AudioSet để gắn thẻ âm thanh. CNN14, được dùng rộng rãi nhất, là một CNN 14 lớp với các tích chập \(3 \times 3\) áp dụng lên log-mel spectrogram. PANNs sinh ra các embedding 2048 chiều đạt học chuyển (transfer learning) tiên tiến nhất trên nhiều tác vụ âm thanh đa dạng.

• Audio Spectrogram Transformer (AST) (Gong et al., 2021) áp dụng trực tiếp kiến trúc Vision Transformer (ViT, chương 08) lên spectrogram âm thanh. Spectrogram được chia thành các mảng \(16 \times 16\) (giống hệt cách ViT chia ảnh), mỗi mảng được chiếu tuyến tính thành một embedding token, thêm positional embedding, và một transformer encoder chuẩn (chương 07) xử lý chuỗi. Đầu ra của token [CLS] được dùng để phân lớp.

• AST hưởng lợi từ tiền huấn luyện ImageNet: vì spectrogram là ảnh 2D, AST khởi tạo từ một ViT tiền huấn luyện trên ảnh ImageNet, rồi tinh chỉnh trên âm thanh. Sự chuyển giao xuyên phương thức này hiệu quả một cách đáng ngạc nhiên vì cả hai miền chia sẻ các đặc trưng bậc thấp (cạnh, kết cấu) và positional embedding có thể nội suy để xử lý các kích thước spectrogram khác nhau.

• HTS-AT (Chen et al., 2022) cải thiện AST với một kiến trúc Swin Transformer phân cấp (chú ý cửa sổ dịch chuyển ở chương 08), giảm chi phí tính toán trong khi nâng cao hiệu năng qua trích xuất đặc trưng đa tỷ lệ.

• BEATs (Chen et al., 2023) dùng một chiến lược tiền huấn luyện đặc thù cho âm thanh: dự đoán có mặt nạ lặp (iterative masked prediction) với một bộ mã hóa rời rạc (tương tự cách tiếp cận của wav2vec 2.0 từ file 02 nhưng áp dụng cho âm thanh nói chung). Bộ mã hóa được tinh lọc dần, tạo ra các token âm thanh rời rạc ngày càng mang ý nghĩa ngữ nghĩa.

• Phân tách người nói với embedding kết hợp embedding người nói với mô hình hóa theo thời gian. Các hệ thống hiện đại như Pyannote.audio dùng một đường ống 3 giai đoạn: (1) một mô hình phân đoạn nơ-ron phát hiện lượt nói và tiếng nói chồng lấp, (2) một giai đoạn trích xuất embedding (ECAPA-TDNN) áp dụng cho mỗi đoạn được phát hiện, và (3) phân cụm để gán danh tính người nói trên toàn bộ bản ghi.

• Truy xuất thông tin nhạc (Music information retrieval, MIR) áp dụng phân tích âm thanh cho âm nhạc. Các biểu diễn phổ từ file 01 đặc biệt hữu ích ở đây vì nhạc có cấu trúc hài âm phong phú.

• Theo dõi nhịp (beat tracking) phát hiện xung nhịp đệm của âm nhạc. Cách tiếp cận chuẩn tính một bao cường độ onset (onset strength envelope) từ spectrogram (phát hiện sự tăng năng lượng báo hiệu các onset nốt), rồi tìm nhịp độ (tempo) dùng tự tương quan (autocorrelation) hoặc tempogram, và cuối cùng theo dõi vị trí từng nhịp dùng quy hoạch động (dynamic programming) để tìm dãy các thời điểm nhịp khớp nhất với bao onset trong khi duy trì một nhịp độ nhất quán.

• Nhận dạng hợp âm (chord recognition) xác định nội dung hài âm theo thời gian. Đầu vào thường là một chromagram (còn gọi là pitch class profile): một biểu diễn 12 chiều gập tất cả các quãng tám lại với nhau, thể hiện năng lượng trong mỗi lớp pitch trong 12 lớp (C, C#, D, ..., B). Một CNN hoặc RNN (chương 06) phân lớp mỗi frame thời gian thành một trong các nhãn hợp âm chuẩn (C trưởng, A thứ, G7, v.v.).

• Chromagram được tính từ STFT (file 01) bằng cách ánh xạ mỗi bin tần số vào lớp pitch của nó:

\[\text{chroma}(p) = \sum_{k : \text{pitch}(k) \bmod 12 = p} |X(k)|^2\]

• trong đó \(p \in \{0, 1, \ldots, 11\}\) là lớp pitch và \(\text{pitch}(k)\) ánh xạ bin tần số \(k\) sang số nốt MIDI của nó.

• Tách nguồn cơ bản (chi tiết hơn ở file 05) tách một bản thu âm nhạc thành từng nhạc cụ riêng biệt (giọng hát, trống, bass, khác). Đây là cốt lõi của các ứng dụng MIR như phối lại, karaoke, và phiên âm nhạc. Các mô hình như Demucs (file 05) đạt chất lượng tách rất tốt trên benchmark chuẩn MUSDB18.

• Gắn thẻ nhạc (music tagging) gán nhãn cho bài hát (thể loại, tâm trạng, nhạc cụ, thời đại). Vên bản chất là phân loại âm thanh áp dụng cho nhạc, dùng cùng cách tiếp cận CNN-trên-spectrogram. The Million Song Dataset và MagnaTagATune là các benchmark chuẩn.

• Lấy vân tay âm thanh (audio fingerprinting) nhận dạng một bản thu âm cụ thể từ một đoạn trích ngắn, thậm chí có nhiễu, dội âm, hoặc artefact nén. Hệ thống kinh điển là Shazam, băm các điểm chòm sao (các đỉnh nổi bật trong spectrogram). Các cách tiếp cận nơ-ron học các embedding bền vững bất biến với suy thoái âm học trong khi vẫn phân biệt được giữa các bản thu âm khác nhau, gợi nhớ học đặc trưng bất biến từ chương 06 và chương 08.

Bài tập lập trình (dùng CoLab hoặc notebook)

• Bài 1: Trích xuất embedding người nói với statistics pooling. Xây dựng một mô hình kiểu x-vector đơn giản xử lý đặc trưng cấp frame qua các lớp TDNN và statistics pooling để sinh embedding người nói.

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

# Simulate frame-level MFCC features for multiple speakers
def generate_speaker_data(key, n_speakers=5, utterances_per_speaker=20,
                          n_frames=100, n_features=40):
    """Generate synthetic speaker data with speaker-dependent patterns."""
    keys = jr.split(key, 3)
    all_features = []
    all_labels = []

    # Each speaker has a characteristic spectral pattern
    speaker_patterns = jr.normal(keys[0], (n_speakers, n_features)) * 0.5

    for spk in range(n_speakers):
        for utt in range(utterances_per_speaker):
            k = jr.fold_in(keys[1], spk * utterances_per_speaker + utt)
            noise = jr.normal(k, (n_frames, n_features)) * 0.3
            features = speaker_patterns[spk][None, :] + noise
            all_features.append(features)
            all_labels.append(spk)

    perm = jr.permutation(keys[2], len(all_features))
    features = jnp.stack(all_features)[perm]
    labels = jnp.array(all_labels)[perm]
    return features, labels

key = jr.PRNGKey(42)
features, labels = generate_speaker_data(key)
n_speakers = 5
n_features = 40

# x-vector-style model
def init_xvector(key, n_features=40, hidden=128, embed_dim=64, n_speakers=5):
    keys = jr.split(key, 8)
    params = {
        # TDNN layer 1: context [-2, 2]
        'tdnn1_w': jr.normal(keys[0], (5, n_features, hidden)) * jnp.sqrt(2.0 / (5 * n_features)),
        'tdnn1_b': jnp.zeros(hidden),
        # TDNN layer 2: context [-2, 2]
        'tdnn2_w': jr.normal(keys[1], (5, hidden, hidden)) * jnp.sqrt(2.0 / (5 * hidden)),
        'tdnn2_b': jnp.zeros(hidden),
        # TDNN layer 3: context [-3, 3]
        'tdnn3_w': jr.normal(keys[2], (7, hidden, hidden)) * jnp.sqrt(2.0 / (7 * hidden)),
        'tdnn3_b': jnp.zeros(hidden),
        # Segment-level layers (after pooling: 2*hidden -> embed_dim)
        'seg1_w': jr.normal(keys[3], (2 * hidden, embed_dim)) * jnp.sqrt(2.0 / (2 * hidden)),
        'seg1_b': jnp.zeros(embed_dim),
        # Classification head
        'cls_w': jr.normal(keys[4], (embed_dim, n_speakers)) * jnp.sqrt(2.0 / embed_dim),
        'cls_b': jnp.zeros(n_speakers),
    }
    return params

def xvector_forward(params, x, return_embedding=False):
    """x: (batch, frames, features) -> logits or embeddings."""
    # TDNN layers (1D convolutions)
    h = jax.lax.conv_general_dilated(
        x.transpose(0, 2, 1), params['tdnn1_w'].transpose(2, 1, 0),
        window_strides=(1,), padding='SAME'
    ).transpose(0, 2, 1) + params['tdnn1_b']
    h = jax.nn.relu(h)

    h = jax.lax.conv_general_dilated(
        h.transpose(0, 2, 1), params['tdnn2_w'].transpose(2, 1, 0),
        window_strides=(1,), padding='SAME'
    ).transpose(0, 2, 1) + params['tdnn2_b']
    h = jax.nn.relu(h)

    h = jax.lax.conv_general_dilated(
        h.transpose(0, 2, 1), params['tdnn3_w'].transpose(2, 1, 0),
        window_strides=(1,), padding='SAME'
    ).transpose(0, 2, 1) + params['tdnn3_b']
    h = jax.nn.relu(h)

    # Statistics pooling: mean and std over time
    mu = jnp.mean(h, axis=1)
    sigma = jnp.std(h, axis=1)
    pooled = jnp.concatenate([mu, sigma], axis=-1)

    # Segment-level layer -> embedding
    embedding = jax.nn.relu(pooled @ params['seg1_w'] + params['seg1_b'])

    if return_embedding:
        return embedding

    # Classification
    logits = embedding @ params['cls_w'] + params['cls_b']
    return logits

def cross_entropy_loss(params, features, labels):
    logits = xvector_forward(params, features)
    one_hot = jax.nn.one_hot(labels, n_speakers)
    log_probs = jax.nn.log_softmax(logits)
    return -jnp.mean(jnp.sum(one_hot * log_probs, axis=-1))

grad_fn = jax.jit(jax.value_and_grad(cross_entropy_loss))

# Train
params = init_xvector(jr.PRNGKey(0))
lr = 1e-3
losses = []

for epoch in range(300):
    loss_val, grads = grad_fn(params, features, labels)
    params = jax.tree.map(lambda p, g: p - lr * g, params, grads)
    losses.append(float(loss_val))

# Extract embeddings and visualise with t-SNE-style 2D projection (using PCA)
embeddings = xvector_forward(params, features, return_embedding=True)

# Simple PCA to 2D
emb_centered = embeddings - jnp.mean(embeddings, axis=0)
_, _, Vt = jnp.linalg.svd(emb_centered, full_matrices=False)
proj_2d = emb_centered @ Vt[:2].T

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

axes[0].plot(losses, color='#3498db', linewidth=1.5)
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Cross-Entropy Loss')
axes[0].set_title('Speaker Classification Training')
axes[0].set_yscale('log')

colors = ['#3498db', '#e74c3c', '#27ae60', '#f39c12', '#9b59b6']
for spk in range(n_speakers):
    mask = labels == spk
    axes[1].scatter(proj_2d[mask, 0], proj_2d[mask, 1], c=colors[spk],
                    label=f'Speaker {spk}', alpha=0.7, s=30)
axes[1].set_xlabel('PC 1')
axes[1].set_ylabel('PC 2')
axes[1].set_title('Speaker Embeddings (PCA projection)')
axes[1].legend()

plt.tight_layout()
plt.show()

# Verification demo: cosine similarity
emb_norm = embeddings / jnp.linalg.norm(embeddings, axis=-1, keepdims=True)
sim_matrix = emb_norm @ emb_norm.T
print(f"Embedding shape: {embeddings.shape}")
print(f"Avg same-speaker similarity: {jnp.mean(sim_matrix[labels[:, None] == labels[None, :]]):.4f}")
print(f"Avg diff-speaker similarity: {jnp.mean(sim_matrix[labels[:, None] != labels[None, :]]):.4f}")

• Bài 2: Xác thực người nói với điểm độ tương đồng cosine. Cho trước các embedding người nói đã tính sẵn, hiện thực một hệ thống xác thực tính EER (Equal Error Rate) và vẽ đường cong DET.

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

def generate_verification_pairs(key, n_speakers=20, dim=64, n_pairs=2000):
    """Generate speaker embeddings and verification trial pairs."""
    keys = jr.split(key, 5)

    # Speaker centroids with some variance
    centroids = jr.normal(keys[0], (n_speakers, dim))
    centroids = centroids / jnp.linalg.norm(centroids, axis=-1, keepdims=True)

    # Generate enrollment and test embeddings with intra-speaker variance
    enroll_embs = []
    test_embs = []
    trial_labels = []  # 1 = same speaker (target), 0 = different (impostor)

    for i in range(n_pairs):
        k1, k2, k3 = jr.split(jr.fold_in(keys[1], i), 3)
        is_target = jr.bernoulli(k1).astype(int)

        spk1 = jr.randint(k2, (), 0, n_speakers)
        emb1 = centroids[spk1] + jr.normal(jr.fold_in(k3, 0), (dim,)) * 0.15

        if is_target:
            spk2 = spk1
        else:
            spk2 = (spk1 + jr.randint(jr.fold_in(k3, 1), (), 1, n_speakers)) % n_speakers

        emb2 = centroids[spk2] + jr.normal(jr.fold_in(k3, 2), (dim,)) * 0.15

        enroll_embs.append(emb1)
        test_embs.append(emb2)
        trial_labels.append(int(is_target))

    return (jnp.stack(enroll_embs), jnp.stack(test_embs),
            jnp.array(trial_labels))

key = jr.PRNGKey(42)
enroll, test, labels = generate_verification_pairs(key)

# Compute cosine similarity scores
enroll_norm = enroll / jnp.linalg.norm(enroll, axis=-1, keepdims=True)
test_norm = test / jnp.linalg.norm(test, axis=-1, keepdims=True)
scores = jnp.sum(enroll_norm * test_norm, axis=-1)

# Compute FAR and FRR at various thresholds
thresholds = jnp.linspace(-1.0, 1.0, 500)

target_scores = scores[labels == 1]
impostor_scores = scores[labels == 0]

fars = []
frrs = []
for thresh in thresholds:
    far = jnp.mean(impostor_scores >= thresh)  # false accepts
    frr = jnp.mean(target_scores < thresh)     # false rejects
    fars.append(float(far))
    frrs.append(float(frr))

fars = jnp.array(fars)
frrs = jnp.array(frrs)

# Find EER: where FAR ≈ FRR
eer_idx = jnp.argmin(jnp.abs(fars - frrs))
eer = float((fars[eer_idx] + frrs[eer_idx]) / 2)
eer_threshold = float(thresholds[eer_idx])

print(f"Equal Error Rate (EER): {eer:.4f} ({eer*100:.2f}%)")
print(f"EER threshold: {eer_threshold:.4f}")

fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# Score distributions
bins = jnp.linspace(-0.5, 1.0, 60)
axes[0].hist(target_scores, bins=bins, alpha=0.6, color='#27ae60',
             label='Target (same speaker)', density=True)
axes[0].hist(impostor_scores, bins=bins, alpha=0.6, color='#e74c3c',
             label='Impostor (different speaker)', density=True)
axes[0].axvline(eer_threshold, color='#f39c12', linestyle='--', linewidth=2,
                label=f'EER threshold = {eer_threshold:.3f}')
axes[0].set_xlabel('Cosine Similarity Score')
axes[0].set_ylabel('Density')
axes[0].set_title('Score Distributions')
axes[0].legend()

# FAR vs FRR
axes[1].plot(thresholds, fars, color='#e74c3c', linewidth=2, label='FAR')
axes[1].plot(thresholds, frrs, color='#3498db', linewidth=2, label='FRR')
axes[1].axvline(eer_threshold, color='#f39c12', linestyle='--', linewidth=1.5)
axes[1].scatter([eer_threshold], [eer], color='#f39c12', s=100, zorder=5,
                label=f'EER = {eer:.4f}')
axes[1].set_xlabel('Threshold')
axes[1].set_ylabel('Error Rate')
axes[1].set_title('FAR and FRR vs Threshold')
axes[1].legend()

# DET curve (FAR vs FRR)
axes[2].plot(fars, frrs, color='#9b59b6', linewidth=2)
axes[2].plot([0, 1], [0, 1], 'k--', alpha=0.3)
axes[2].scatter([eer], [eer], color='#f39c12', s=100, zorder=5,
                label=f'EER = {eer:.4f}')
axes[2].set_xlabel('False Acceptance Rate')
axes[2].set_ylabel('False Rejection Rate')
axes[2].set_title('DET Curve')
axes[2].set_xlim([0, 0.5])
axes[2].set_ylim([0, 0.5])
axes[2].legend()
axes[2].set_aspect('equal')

plt.tight_layout()
plt.show()

• Bài 3: Embedding mảng spectrogram âm thanh (kiểu AST). Hiện thực lớp trích xuất và embedding mảng của Audio Spectrogram Transformer, minh họa cách một spectrogram được token hóa.

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

# Generate a synthetic spectrogram (harmonic structure + noise)
def generate_spectrogram(key, n_time=128, n_freq=128):
    """Create a synthetic spectrogram with harmonic patterns."""
    k1, k2 = jr.split(key)
    spec = jr.normal(k1, (n_time, n_freq)) * 0.1

    # Add harmonic bands (simulating speech formants)
    for f0 in [15, 30, 45, 70]:
        width = 3
        envelope = jnp.exp(-0.5 * ((jnp.arange(n_freq) - f0) / width) ** 2)
        time_mod = 0.5 + 0.5 * jnp.sin(2 * jnp.pi * jnp.arange(n_time) / 40)
        spec += jnp.outer(time_mod, envelope)

    return jnp.clip(spec, 0, None)

key = jr.PRNGKey(42)
spectrogram = generate_spectrogram(key)
n_time, n_freq = spectrogram.shape

# Patch extraction parameters
patch_h = 16  # time
patch_w = 16  # frequency
stride_h = 16
stride_w = 16
embed_dim = 192  # ViT-Small dimension

n_patches_h = n_time // stride_h
n_patches_w = n_freq // stride_w
n_patches = n_patches_h * n_patches_w

print(f"Spectrogram: {n_time} x {n_freq}")
print(f"Patch size: {patch_h} x {patch_w}")
print(f"Number of patches: {n_patches_h} x {n_patches_w} = {n_patches}")

# Extract patches
def extract_patches(spec, patch_h, patch_w, stride_h, stride_w):
    """Extract non-overlapping patches from spectrogram."""
    patches = []
    positions = []
    for i in range(0, spec.shape[0] - patch_h + 1, stride_h):
        for j in range(0, spec.shape[1] - patch_w + 1, stride_w):
            patch = spec[i:i+patch_h, j:j+patch_w]
            patches.append(patch.flatten())
            positions.append((i, j))
    return jnp.stack(patches), positions

patches, positions = extract_patches(spectrogram, patch_h, patch_w, stride_h, stride_w)
print(f"Patches shape: {patches.shape}")  # (n_patches, patch_h * patch_w)

# Linear projection (patch embedding)
patch_dim = patch_h * patch_w
k1, k2 = jr.split(jr.PRNGKey(0))
W_embed = jr.normal(k1, (patch_dim, embed_dim)) * jnp.sqrt(2.0 / patch_dim)
b_embed = jnp.zeros(embed_dim)

# Learnable positional embeddings
pos_embed = jr.normal(k2, (n_patches + 1, embed_dim)) * 0.02  # +1 for CLS

# CLS token
cls_token = jnp.zeros((1, embed_dim))

# Forward pass
patch_tokens = patches @ W_embed + b_embed  # (n_patches, embed_dim)
tokens = jnp.concatenate([cls_token, patch_tokens], axis=0)  # (n_patches+1, embed_dim)
tokens = tokens + pos_embed  # Add positional embeddings

print(f"Token sequence shape: {tokens.shape}")
print(f"Each token has dimension: {embed_dim}")

# Visualisation
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# Original spectrogram with patch grid
axes[0, 0].imshow(spectrogram.T, aspect='auto', origin='lower', cmap='magma')
for i in range(0, n_time + 1, stride_h):
    axes[0, 0].axvline(i - 0.5, color='white', linewidth=0.5, alpha=0.5)
for j in range(0, n_freq + 1, stride_w):
    axes[0, 0].axhline(j - 0.5, color='white', linewidth=0.5, alpha=0.5)
axes[0, 0].set_title(f'Spectrogram with {patch_h}x{patch_w} Patch Grid')
axes[0, 0].set_xlabel('Time frame')
axes[0, 0].set_ylabel('Frequency bin')

# Individual patches visualised
n_show = min(16, n_patches)
patch_grid = patches[:n_show].reshape(n_show, patch_h, patch_w)
combined = jnp.concatenate([patch_grid[i] for i in range(min(8, n_show))], axis=1)
axes[0, 1].imshow(combined.T, aspect='auto', origin='lower', cmap='magma')
axes[0, 1].set_title(f'First {min(8, n_show)} Patches (concatenated)')
axes[0, 1].set_xlabel('Patch index (horizontal)')
axes[0, 1].set_ylabel('Frequency within patch')

# Token embeddings similarity matrix
token_norms = tokens / jnp.linalg.norm(tokens, axis=-1, keepdims=True)
sim = token_norms @ token_norms.T
im = axes[1, 0].imshow(sim, cmap='RdBu_r', vmin=-1, vmax=1)
axes[1, 0].set_title('Token Similarity Matrix (cosine)')
axes[1, 0].set_xlabel('Token index')
axes[1, 0].set_ylabel('Token index')
plt.colorbar(im, ax=axes[1, 0], fraction=0.046)

# Positional embedding similarity
pos_norms = pos_embed / jnp.linalg.norm(pos_embed, axis=-1, keepdims=True)
pos_sim = pos_norms @ pos_norms.T
im2 = axes[1, 1].imshow(pos_sim, cmap='RdBu_r', vmin=-1, vmax=1)
axes[1, 1].set_title('Positional Embedding Similarity')
axes[1, 1].set_xlabel('Position index')
axes[1, 1].set_ylabel('Position index')
plt.colorbar(im2, ax=axes[1, 1], fraction=0.046)

plt.tight_layout()
plt.show()

• Bài 4: Tính chromagram đơn giản cho phân tích hợp âm. Tính và minh họa một chromagram từ một tín hiệu hài âm tổng hợp, trình diễn phép gập lớp pitch dùng trong truy xuất thông tin nhạc.

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

# Generate a synthetic musical signal: C major chord -> G major chord
sr = 16000
duration = 2.0
t = jnp.linspace(0, duration, int(sr * duration))

# C major (C4=261.6, E4=329.6, G4=392.0) for first half
# G major (G3=196.0, B3=246.9, D4=293.7) for second half
half = len(t) // 2

c_major = (0.5 * jnp.sin(2 * jnp.pi * 261.63 * t[:half]) +
           0.4 * jnp.sin(2 * jnp.pi * 329.63 * t[:half]) +
           0.3 * jnp.sin(2 * jnp.pi * 392.00 * t[:half]))

g_major = (0.5 * jnp.sin(2 * jnp.pi * 196.00 * t[:half]) +
           0.4 * jnp.sin(2 * jnp.pi * 246.94 * t[:half]) +
           0.3 * jnp.sin(2 * jnp.pi * 293.66 * t[:half]))

signal = jnp.concatenate([c_major, g_major])

# Compute STFT
n_fft = 4096  # high resolution for pitch accuracy
hop_length = 512
window = jnp.hanning(n_fft)

def stft(signal, n_fft, hop_length, window):
    n_frames = 1 + (len(signal) - n_fft) // hop_length
    frames = jnp.stack([
        signal[i * hop_length : i * hop_length + n_fft] * window
        for i in range(n_frames)
    ])
    return jnp.fft.rfft(frames, n=n_fft)

S = stft(signal, n_fft, hop_length, window)
power_spec = jnp.abs(S) ** 2
freqs = jnp.fft.rfftfreq(n_fft, 1.0 / sr)

# Compute chromagram by mapping frequency bins to pitch classes
# MIDI note number from frequency: 69 + 12 * log2(f / 440)
note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']

def freq_to_chroma(freq):
    """Map frequency to pitch class (0-11). Returns -1 for freq <= 0."""
    midi = 69 + 12 * jnp.log2(jnp.clip(freq, 1e-10, None) / 440.0)
    return jnp.round(midi).astype(int) % 12

# Build chromagram: sum power spectrum energy for each pitch class
chromagram = jnp.zeros((power_spec.shape[0], 12))
valid_freqs = freqs[1:]  # skip DC
valid_power = power_spec[:, 1:]

for p in range(12):
    # Find frequency bins belonging to this pitch class
    chroma_bins = freq_to_chroma(valid_freqs)
    mask = (chroma_bins == p).astype(jnp.float32)
    chromagram = chromagram.at[:, p].set(
        jnp.sum(valid_power * mask[None, :], axis=1)
    )

# Normalise each frame
chromagram = chromagram / (jnp.max(chromagram, axis=1, keepdims=True) + 1e-8)

# Visualisation
fig, axes = plt.subplots(3, 1, figsize=(14, 10))

# Waveform
axes[0].plot(t[:3000], signal[:3000], color='#3498db', linewidth=0.5,
             label='C major')
axes[0].plot(t[half:half+3000], signal[half:half+3000], color='#e74c3c',
             linewidth=0.5, label='G major')
axes[0].set_title('Waveform: C major → G major')
axes[0].set_ylabel('Amplitude')
axes[0].set_xlabel('Time (s)')
axes[0].legend()

# Spectrogram (log scale)
time_axis = jnp.arange(power_spec.shape[0]) * hop_length / sr
axes[1].imshow(jnp.log1p(power_spec[:, :500].T), aspect='auto', origin='lower',
               cmap='magma', extent=[0, time_axis[-1], 0, freqs[500]])
axes[1].set_title('Power Spectrogram')
axes[1].set_ylabel('Frequency (Hz)')
axes[1].set_xlabel('Time (s)')

# Chromagram
im = axes[2].imshow(chromagram.T, aspect='auto', origin='lower', cmap='YlOrRd',
                     extent=[0, time_axis[-1], -0.5, 11.5])
axes[2].set_yticks(range(12))
axes[2].set_yticklabels(note_names)
axes[2].set_title('Chromagram (pitch class energy over time)')
axes[2].set_ylabel('Pitch class')
axes[2].set_xlabel('Time (s)')
plt.colorbar(im, ax=axes[2], fraction=0.046, label='Normalised energy')

# Mark expected active pitch classes
mid_frame = chromagram.shape[0] // 2
print(f"C major region - expected: C, E, G")
print(f"  Chroma values: {dict(zip(note_names, [f'{v:.2f}' for v in chromagram[mid_frame//2]]))}")
print(f"G major region - expected: G, B, D")
print(f"  Chroma values: {dict(zip(note_names, [f'{v:.2f}' for v in chromagram[mid_frame + mid_frame//2]]))}")

plt.tight_layout()
plt.show()