Bài Lab tổng hợp — Qwen3-Omni & Qwen3-TTS-12Hz Architecture

Giới thiệu bài lab

Lab Qwen3-Omni và TTS-12Hz: chạy thử toàn trình thu nhỏ từ phổ Mel đến 48 mẫu sóng âm bằng NumPy

Bài lab là một script Python dùng NumPy (thư viện tính toán trên mảng số) chạy ngay trong trình duyệt. Nó dựng lại cả 10 khối 1A đến 6B ở kích thước thu nhỏ: spectrogram Mel 8 kênh × 16 khung, vector 8 chiều ở Thinker và 6 chiều ở Talker, 6 chuyên gia chọn 2, 4 sổ mã RVQ mỗi sổ 8 mục, và 4 khung 80 ms. Kết quả cuối là 48 mẫu sóng âm.

Nhiều bước của kiến trúc, như dịch chỉ số bằng code_offset, cửa sổ trượt hay giải mã theo từng đoạn, khó hình dung chỉ bằng chữ. Với số nhỏ, bạn tự cộng, tự đếm và đối chiếu được, ví dụ bộ mã [4, 1, 6, 4] cộng độ lệch [0, 8, 16, 24] ra chỉ số dòng [4, 9, 22, 28].

Bài lab gồm

Sơ đồ 4 cột có khung tương tác ở đầu trang, để biết mỗi bảng số thuộc khối nào.
Bảng ánh xạ 1:1 cho từng mã 1A đến 6B: lab dựng gì và nhìn số nào để kiểm chứng.
Trình chạy trực tiếp với mã qwen3-omni-speech.py sửa được, và các nút Chạy cell, Dừng, Đặt lại.
Output chia thành các khối [1A] đến [6B], rồi mục [KIỂM TRA & THỬ NGHIỆM] với 3 phép kiểm chứng tự động.
Trang bài đọc Lab đầy đủ (liên kết ở thanh công cụ) có hướng dẫn chi tiết và 4 bài tập kèm đáp án.

Làm xong, bạn có thể

Giải thích bằng số vì sao Mel (8, 16) còn (1, 2) sau 3 tầng Conv2d bước nhảy 2 (mỗi tầng giảm một nửa).
Đọc 4 khung Dual-Track: khung 0 và 1 dùng chữ thật từ lớp 24, khung 2 và 3 chuyển sang vector đệm tts_pad_embed vì chữ đã hết.
Tự tính chỉ số dòng trong bảng RVQ gộp từ mã và code_offset.
Đọc mặt nạ cửa sổ trượt 4×4 và giải thích vì sao giải mã theo từng đoạn cần ghép thêm một khung ngữ cảnh bên trái.
Nói được lab chứng minh gì (đường đi của tensor và các phép tính đúng) và không chứng minh gì (48 mẫu sóng âm không phải giọng nói thật).

Phạm vi: lab chỉ chạy suy luận (forward) của một mô hình thu nhỏ và không có bước huấn luyện; trọng số là số ngẫu nhiên cố định bằng seed 30, nên chạy lại cho ra đúng các số cũ nhưng sóng âm đầu ra không phát ra lời nói. Muốn học huấn luyện, hãy xem khóa LLM Training (Text), vốn chỉ dạy mô hình văn bản; huấn luyện mô hình đa phương thức (âm thanh, giọng nói) chưa có bài nào trên site.

Cách dùng bài lab

Xem sơ đồ 4 cột ở đầu trang: rê chuột hoặc bấm vào khung để mở pop-over, các nút lọc theo cột làm mờ khung của những cột còn lại, và nút Ẩn khung overlay cho bạn đọc ảnh gốc.
Đọc bảng ánh xạ bên dưới: cột thứ ba cho biết lab dựng gì ở mỗi mã, cột thứ tư cho biết nhìn số nào để kiểm chứng.
Đọc cấu hình thu nhỏ ở đầu mã, nơi mỗi hằng số có chú thích kích thước thật: spectrogram Mel (8, 16), 2 token chữ do Thinker sinh, người nói speaker_id = 2302 (Ethan), 4 khung ở 12,5 Hz và seed 30.
Dự đoán trước rồi bấm Chạy cell. Ở [1A] đến [2A], kiểm tra Mel giảm từ (8, 16) xuống (1, 2) và chuẩn vector không đổi sau khi xoay TM-RoPE.
Ở [3A] đến [4A], theo dõi 4 khung Dual-Track, hai trọng số chuyên gia Top-2 cộng lại bằng 1, và bộ mã [4, 1, 6, 4] của khung 0. Lưu ý output in [6A] trước [5A], vì 5A cần vector trung bình mà 6A tạo ra.
Ở [6A], [5A], [6B] và [KIỂM TRA & THỬ NGHIỆM], đối chiếu phép dịch chỉ số, mặt nạ cửa sổ trượt 4×4, 48 mẫu sóng âm và 3 phép kiểm chứng. Bạn có thể sửa mã và chạy lại; nếu đổi cấu hình, các assert ở cuối có thể báo lỗi, và nút Đặt lại khôi phục mã gốc.

Đọc kết quả đúng cách: 48 mẫu sóng âm chỉ dùng để kiểm tra đường đi của tensor và độ đúng của phép tính; trọng số ngẫu nhiên không tạo ra giọng nói thật.

Sơ đồ luồng dữ liệu 4 cột · 10 thành phần (1A → 6B)

Kiến trúc Qwen3-Omni & Qwen3-TTS-12Hz · Ánh xạ 1:1 với Bài Lab tổng hợp

Bắt đầu từ 1A ở góc dưới Cột 1 đi lên 1B → 2A, sang Cột 2 (3A → 3B → 3C), qua Cột 3 (4A → 6A) và kết thúc ở Cột 4 (5A → 6B). Di chuột hoặc bấm vào từng khối để xem pop-over và đối chiếu với Lab.

Sơ đồ luồng dữ liệu 4 cột của kiến trúc Qwen3-Omni & Qwen3-TTS-12Hz gồm 10 thành phần từ 1A đến 6B

1A1 · AuT Audio, TM-RoPE & Thinker MoE

Mel [B, 128, T_mel] → 3× Conv2d(k=3, s=2) (8× Downsample) → Windowed Attn (cu_seqlens) → [B, T_aud, 2048]

1A · AuT Audio Encoder (3× Conv2d Stride-2)

Điểm xuất phát ở đáy Cột 1: nén spectrogram 128 kênh Mel giảm 8 lần theo cả tần số lẫn thời gian qua 3 tầng Conv2d stride-2 và Windowed Attention.

Bắt đầu đọc sơ đồ từ góc dưới bên trái (1A): Qwen3OmniMoeAudioEncoder chia dải Mel spectrogram 128 kênh thành từng cửa sổ n_window, đưa qua 3 lớp tích chập 2 chiều Conv2d(kernel=3, stride=2) liên tiếp để giảm 2³ = 8 lần độ dài thời gian và số kênh tần số. Chuỗi đặc trưng sau đó cộng SinusoidsPositionEmbedding, chạy Attention cục bộ theo ranh giới cu_seqlens và chiếu qua MLP (1280 → GELU → 2048) để khớp với không gian ẩn của Thinker.

Bài 1A · AuT Audio Encoder & 3× Conv2d Stride-2 Xem bước [1A] trong Bài Lab tổng hợp

BẢNG ÁNH XẠ 1:1 · MỖI KHỐI KIẾN TRÚC Ở TAB 01 ĐƯỢC DỰNG VÀ KIỂM CHỨNG TRONG LAB NHƯ THẾ NÀO?

Mã	Thành phần ở Tab 01 (bấm để mở bài)	Lab NumPy dựng gì (`qwen3-omni-speech.py`)	Cách đọc số liệu & phép kiểm chứng
1A	AuT Audio Encoder (3× Conv2d Stride-2 & Windowed Attn) ↗	Nén spectrogram Mel qua 3 tầng tích chập stride-2 (giảm 2³ = 8× thời gian & tần số), cộng SinusoidsPositionEmbedding và chạy Attention theo cu_seqlens.	Kiểm tra kích thước khung Mel giảm đúng 8 lần trước khi chiếu sang chiều ẩn của Thinker (2048).
1B	Interleaved TM-RoPE 3D (24, 20, 20) @ 12.5 Hz ↗	Đan xen các cặp tần số theo tỷ lệ [24, 20, 20] (Temporal, Height, Width) đồng bộ theo nhịp 12.5 Hz (80 ms/khung).	Xác nhận bảo toàn chuẩn Euclid ‖RoPE(q)‖ = ‖q‖ và tính bất biến góc tương đối giữa khung âm thanh và khung hình.
2A	Thinker MoE & Trích ngang trạng thái ẩn Lớp 24 ↗	Chạy chuỗi lớp Thinker MoE để phát văn bản ở tầng đỉnh, đồng thời trích ngang accept_hidden_layer = 24 làm luồng trailing_text_hidden.	Đối chiếu vector tại Lớp 24 giàu thông tin ngữ điệu hơn so với trạng thái đã ép sát từ vựng ở lớp cuối.
3A	Cầu nối ResizeMLP kép & Điều kiện hóa Speaker ID ↗	Chiếu text_projection và hidden_projection qua Linear(2048→2048) → SiLU → Linear(2048→1024) kèm nhúng speaker_id (Ethan=2302, Chelsie=2301).	Kiểm tra phép hạ chiều 2048 → 1024 và sự thay đổi phân phối mã âm thanh khi đổi speaker_id.
3B	Đồng bộ Dual-Track Streaming 12.5 Hz (97 ms Latency) ↗	Tại mỗi bước generation_step (80 ms), cộng vector tổng 16 sổ mã bước trước với trailing_text_hidden[step] (hoặc tts_pad_embed khi hết chữ).	Xác nhận Talker phát khung âm thanh đầu tiên ngay sau token chữ thứ nhất mà không cần chờ Thinker sinh hết câu.
3C	Talker MoE (20 Lớp, 128 Experts Top-6) & Mã ngữ nghĩa c_0 ↗	Chạy GQA (16Q/2KV) + MoE Top-6/128 chuyên gia + 1 chuyên gia chung, xuất h_talker và chiếu codec_head dự đoán mã ngữ nghĩa c_0.	Kiểm tra tổng trọng số Top-6 sau chuẩn hóa bằng 1.0 và mã ngữ nghĩa c_0 được chọn cho từng khung 80 ms.
4A	5-Layer Depth Code Predictor (Tự hồi quy c_1 .. c_15) ↗	Đưa cặp prefill [h_talker, embed(c_0)] vào 5 lớp Depth Transformer và chạy 15 bước tự hồi quy qua 15 đầu lm_head[0..14] riêng biệt.	Kiểm chứng toàn bộ 15 mã âm học c_1..c_15 được sinh ngay trong cùng một khung 80 ms mà không phải gọi lại Talker 20 lớp.
5A	Code2Wav Pre-Transformer (Window=72) & ConvNeXt 4× ↗	Áp dụng mặt nạ nhân quả cửa sổ trượt 72 khung + LayerScale (0.01) và 2 tầng CausalTransConv(×2) + ConvNeXt nâng tần số 4× (12.5 Hz → 50 Hz).	Kiểm chứng mặt nạ cửa sổ trượt chặn tuyệt đối các khung tương lai (cô lập nhân quả 0.00e+00) và nhân đúng 4× số bước thời gian.
6A	Bảng nhúng 16-Group RVQ (Offset 2048) & Mean Pooling ↗	Cộng chỉ số 16 sổ mã với code_offset = arange(16) × K để tra cứu một lần trên bảng gộp rồi lấy trung bình .mean(dim=1).	Chứng minh tra bảng gộp với code_offset cho kết quả khớp 0.00e+00 so với 16 vòng lặp tra cứu rời rạc.
6B	SnakeBeta Causal ConvNet 480× (Tổng 1920×) → Sóng âm 24 kHz ↗	Kích hoạt tuần hoàn SnakeBeta(x) = x + (1/e^β)·sin²(e^α·x), phóng đại 4 tầng [×8, ×5, ×4, ×3] = 480× (tổng 4×480 = 1920 mẫu/khung) và cắt left_context khi giải mã theo chunk.	Xác nhận 1 khung 80 ms (12.5 Hz) nở ra đúng 1920 mẫu sóng âm 24 kHz nằm trong [-1, 1].

Trình chạy Lab trực tiếp · `qwen3-omni-speech.py`

qwen3-omni-speech.pyPYTHON

# BÀI LAB TỔNG HỢP QWEN3-OMNI & QWEN3-TTS-12HZ: TOÀN TRÌNH 1A → 6B BẰNG NUMPY
# Chạy một lượt từ nén spectrogram AuT (1A), đan xen TM-RoPE 3D (1B), trích ngang
# Lớp 24 của Thinker (2A), hạ chiều ResizeMLP & Speaker ID (3A), đồng bộ Dual-Track
# 12.5 Hz (3B), Talker MoE sinh mã ngữ nghĩa c_0 (3C), 5-Layer Code Predictor tự
# hồi quy các mã âm học c_1..c_3 (4A), bảng nhúng gộp RVQ code_offset (6A),
# Pre-Transformer cửa sổ trượt (5A) đến giải mã sóng âm tuần hoàn SnakeBeta (6B).

import numpy as np

rng = np.random.default_rng(30)

# ── Cấu hình thu nhỏ (ánh xạ 1:1 với Qwen3-Omni-30B-A3B & Qwen3-TTS-12Hz) ────
MEL_BINS = 8              # prod: 128 mel bins
MEL_FRAMES = 16           # 16 khung Mel → sau 3× stride-2 còn 1 × 2 khung
THINKER_DIM = 8           # prod: thinker_hidden_size = 2048
TALKER_DIM = 6            # prod: talker hidden_size = 1024
HEAD_DIM = 6              # 3 cặp tần số (T, H, W) đan xen (prod: 64 cặp = [24, 20, 20])
NUM_EXPERTS = 6           # prod: 128 routed experts
TOP_K_EXPERTS = 2         # prod: num_experts_per_tok = 6
NUM_QUANTIZERS = 4        # prod: num_code_groups = 16 (1 semantic + 15 acoustic)
CODEBOOK_SIZE = 8         # prod: codebook_size = 2048
SLIDING_WINDOW = 2        # prod: sliding_window = 72 frames
PRE_UPSAMPLE = 2          # prod: upsampling_ratios = [2, 2] (4×)
DEC_UPSAMPLES = (2, 3)    # prod: upsample_rates = [8, 5, 4, 3] (480×)
TOTAL_UPSAMPLE = PRE_UPSAMPLE * DEC_UPSAMPLES[0] * DEC_UPSAMPLES[1]  # 2 × 6 = 12× (prod: 1920×)

def silu(x):
    x = np.asarray(x, dtype=float)
    return x / (1.0 + np.exp(-x))

def gelu_tanh(x):
    x = np.asarray(x, dtype=float)
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3)))

def rmsnorm(x, weight, eps=1e-6):
    x = np.asarray(x, dtype=float)
    ms = np.mean(x * x, axis=-1, keepdims=True)
    return (x / np.sqrt(ms + eps)) * weight

def softmax(x, axis=-1):
    x = np.asarray(x, dtype=float)
    s = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(s)
    return e / np.sum(e, axis=axis, keepdims=True)

def fmt_vec(v, nd=3):
    arr = np.asarray(v, dtype=float).ravel()
    return "[" + ", ".join(f"{x:.{nd}f}" for x in arr) + "]"

def say(code, title, lines):
    print(f"\n[{code}] {title}")
    for ln in lines:
        print("  ", ln)

# ── [1A] AuT Audio Encoder: 3× Conv2d Stride-2 + Windowed Attention ──────────
def conv2d_stride2_mean(mat):
    # Tích chập 2D bước nhảy 2 (kernel 2×2 trung bình có trọng số cố định)
    h, w = mat.shape
    out = np.zeros((h // 2, w // 2), dtype=float)
    kernel = np.array([[0.35, 0.25], [0.20, 0.20]], dtype=float)
    for i in range(h // 2):
        for j in range(w // 2):
            patch = mat[2 * i : 2 * i + 2, 2 * j : 2 * j + 2]
            out[i, j] = float(np.sum(patch * kernel))
    return out

MEL_INPUT = rng.normal(0.0, 0.8, size=(MEL_BINS, MEL_FRAMES))
W_AUT_PROJ = rng.normal(0.0, 0.35, size=(1, THINKER_DIM))
W_AUT_Q = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
W_AUT_K = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
W_AUT_V = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))

def run_aut_encoder(mel_input, cu_seqlens=(0, 1, 2)):
    c1 = conv2d_stride2_mean(mel_input)   # (4, 8)
    c2 = conv2d_stride2_mean(c1)          # (2, 4)
    c3 = conv2d_stride2_mean(c2)          # (1, 2) -> 2 khung âm thanh sau 8× downsample
    frames = c3.T @ W_AUT_PROJ            # (2, 8)
    pos = np.array([[np.sin(t / 4.0 + i * 0.2) for i in range(THINKER_DIM)] for t in range(frames.shape[0])])
    x = frames + 0.15 * pos
    out = np.zeros_like(x)
    for idx in range(len(cu_seqlens) - 1):
        a, b = cu_seqlens[idx], cu_seqlens[idx + 1]
        seg = x[a:b]
        q = seg @ W_AUT_Q
        k = seg @ W_AUT_K
        v = seg @ W_AUT_V
        att = softmax((q @ k.T) / np.sqrt(THINKER_DIM), axis=-1)
        out[a:b] = seg + att @ v
    return {"c1": c1, "c2": c2, "c3": c3, "out": out}

# ── [1B] Interleaved TM-RoPE 3D: đan xen 3 trục (Temporal, Height, Width) ────
def apply_interleaved_tm_rope(vec, pos_thw):
    # vec: (..., 6) gồm 3 cặp đan xen: cặp 0 = Temporal, cặp 1 = Height, cặp 2 = Width
    thetas = np.array([1.0, 0.1, 0.01], dtype=float)
    angles = np.asarray(pos_thw, dtype=float) * thetas
    out = np.array(vec, dtype=float, copy=True)
    for pair_idx in range(3):
        phi = angles[pair_idx]
        c, s = np.cos(phi), np.sin(phi)
        x0 = out[..., 2 * pair_idx].copy()
        x1 = out[..., 2 * pair_idx + 1].copy()
        out[..., 2 * pair_idx] = x0 * c - x1 * s
        out[..., 2 * pair_idx + 1] = x0 * s + x1 * c
    return out

# ── [2A] Thinker MoE & Layer-24 Hidden Handoff ───────────────────────────────
W_THINKER_L12 = rng.normal(0.0, 0.3, size=(THINKER_DIM, THINKER_DIM))
W_THINKER_L24 = rng.normal(0.0, 0.3, size=(THINKER_DIM, THINKER_DIM))
W_THINKER_TOP = rng.normal(0.0, 0.3, size=(THINKER_DIM, THINKER_DIM))
THINKER_TEXT_EMB = rng.normal(0.0, 0.5, size=(2, THINKER_DIM))  # 2 token văn bản sinh ra

def run_thinker_handoff(text_emb):
    h12 = np.tanh(text_emb @ W_THINKER_L12)
    h24 = h12 + silu(h12 @ W_THINKER_L24)      # accept_hidden_layer = 24
    htop = h24 @ W_THINKER_TOP                 # tầng cuối sát LM Head
    return {"h12": h12, "h24": h24, "htop": htop}

# ── [3A] ResizeMLP Bridge (2048 → 1024) & Speaker Conditioning ───────────────
W_RESIZE_TXT1 = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
B_RESIZE_TXT1 = rng.normal(0.0, 0.05, size=(THINKER_DIM,))
W_RESIZE_TXT2 = rng.normal(0.0, 0.35, size=(THINKER_DIM, TALKER_DIM))
B_RESIZE_TXT2 = rng.normal(0.0, 0.05, size=(TALKER_DIM,))

W_RESIZE_HID1 = rng.normal(0.0, 0.35, size=(THINKER_DIM, THINKER_DIM))
B_RESIZE_HID1 = rng.normal(0.0, 0.05, size=(THINKER_DIM,))
W_RESIZE_HID2 = rng.normal(0.0, 0.35, size=(THINKER_DIM, TALKER_DIM))
B_RESIZE_HID2 = rng.normal(0.0, 0.05, size=(TALKER_DIM,))

SPEAKER_TABLE = {
    2301: rng.normal(0.0, 0.4, size=(TALKER_DIM,)),  # Chelsie
    2302: rng.normal(0.0, 0.4, size=(TALKER_DIM,)),  # Ethan
    2303: rng.normal(0.0, 0.4, size=(TALKER_DIM,)),  # Aiden
}
TTS_PAD_EMBED = rng.normal(0.0, 0.25, size=(TALKER_DIM,))
CODEC_BOS_EMBED = rng.normal(0.0, 0.25, size=(TALKER_DIM,))

def resize_text_mlp(x):
    return silu(x @ W_RESIZE_TXT1 + B_RESIZE_TXT1) @ W_RESIZE_TXT2 + B_RESIZE_TXT2

def resize_hidden_mlp(x):
    return silu(x @ W_RESIZE_HID1 + B_RESIZE_HID1) @ W_RESIZE_HID2 + B_RESIZE_HID2

# ── [3C] Talker MoE (Top-2 of 6 Routed + 1 Shared) & Semantic Head c_0 ───────
W_TALKER_ROUTER = rng.normal(0.0, 0.4, size=(TALKER_DIM, NUM_EXPERTS))
W_TALKER_EXPERTS = rng.normal(0.0, 0.3, size=(NUM_EXPERTS, TALKER_DIM, TALKER_DIM))
W_TALKER_SHARED = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_CODEC_HEAD = rng.normal(0.0, 0.4, size=(TALKER_DIM, CODEBOOK_SIZE))
TALKER_RMS_W = np.ones(TALKER_DIM, dtype=float)

def talker_moe_step(dual_track_in, speaker_vec):
    x = rmsnorm(dual_track_in + 0.25 * speaker_vec, TALKER_RMS_W)
    router_probs = softmax(x @ W_TALKER_ROUTER, axis=-1)
    top_idx = np.argsort(-router_probs)[:TOP_K_EXPERTS]
    top_vals = router_probs[top_idx]
    norm_weights = top_vals / np.sum(top_vals)  # norm_topk_prob = True
    routed_out = np.zeros(TALKER_DIM, dtype=float)
    for w, e_id in zip(norm_weights, top_idx):
        routed_out += w * silu(x @ W_TALKER_EXPERTS[int(e_id)])
    shared_out = silu(x @ W_TALKER_SHARED)
    h_talker = dual_track_in + routed_out + shared_out
    logits_c0 = h_talker @ W_CODEC_HEAD
    c0 = int(np.argmax(logits_c0))
    return {
        "h_talker": h_talker,
        "router_probs": router_probs,
        "top_idx": top_idx,
        "norm_weights": norm_weights,
        "logits_c0": logits_c0,
        "c0": c0,
    }

# ── [6A] Unified 16-Group RVQ Table (thu nhỏ: 4 × 8 = 32 dòng) ──────────────
UNIFIED_RVQ_TABLE = rng.normal(0.0, 0.45, size=(NUM_QUANTIZERS * CODEBOOK_SIZE, TALKER_DIM))
CODE_OFFSET = np.arange(NUM_QUANTIZERS, dtype=int).reshape(NUM_QUANTIZERS, 1) * CODEBOOK_SIZE

def lookup_rvq_offset_mean(codes_matrix):
    # codes_matrix: (NUM_QUANTIZERS, T_frames)
    shifted = codes_matrix + CODE_OFFSET
    emb = UNIFIED_RVQ_TABLE[shifted]  # (NUM_QUANTIZERS, T_frames, TALKER_DIM)
    return np.mean(emb, axis=0), shifted  # (T_frames, TALKER_DIM)

def lookup_rvq_separate_tables_mean(codes_matrix):
    # Tra cứu rời rạc từng bảng con để đối chứng với bảng gộp
    separate = [
        UNIFIED_RVQ_TABLE[q * CODEBOOK_SIZE : (q + 1) * CODEBOOK_SIZE][codes_matrix[q]]
        for q in range(NUM_QUANTIZERS)
    ]
    return np.mean(np.stack(separate, axis=0), axis=0)

# ── [4A] 5-Layer Depth Code Predictor (tự hồi quy c_1 .. c_{Q-1}) ────────────
W_DEPTH_ATTN = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_DEPTH_MLP = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
DEPTH_STEP_HEADS = [
    rng.normal(0.0, 0.4, size=(TALKER_DIM, CODEBOOK_SIZE))
    for _ in range(NUM_QUANTIZERS - 1)
]

def run_depth_code_predictor(h_talker, c0):
    # Prefill 2 vị trí: [h_talker, embed_0(c_0)]
    c0_emb = UNIFIED_RVQ_TABLE[0 * CODEBOOK_SIZE + int(c0)]
    seq = [h_talker, c0_emb]
    codes = [int(c0)]
    step_logs = []
    for step in range(NUM_QUANTIZERS - 1):
        mat = np.stack(seq, axis=0)  # (len_seq, TALKER_DIM)
        scores = (mat[-1:] @ mat.T) / np.sqrt(TALKER_DIM)
        ctx = softmax(scores, axis=-1)[0] @ (mat @ W_DEPTH_ATTN)
        h_depth = mat[-1] + ctx + silu((mat[-1] + ctx) @ W_DEPTH_MLP)
        logits_k = h_depth @ DEPTH_STEP_HEADS[step]
        ck = int(np.argmax(logits_k))
        codes.append(ck)
        step_logs.append({"step": step, "seq_len": len(seq), "code_idx": step + 1, "code": ck, "logits": logits_k})
        # Tra bảng nhúng riêng của tầng (step + 1) cho bước vi mô kế tiếp
        next_emb = UNIFIED_RVQ_TABLE[(step + 1) * CODEBOOK_SIZE + ck]
        seq.append(next_emb)
    return np.array(codes, dtype=int), step_logs

# ── [3B] Vòng lặp Dual-Track 12.5 Hz qua 4 khung 80 ms ───────────────────────
def run_dual_track_frames(trailing_text_hidden, speaker_id=2302, n_frames=4):
    speaker_vec = SPEAKER_TABLE[speaker_id]
    prev_audio_emb = CODEC_BOS_EMBED.copy()
    all_codes = []
    frame_records = []
    for t in range(n_frames):
        used_pad = t >= len(trailing_text_hidden)
        text_cond = TTS_PAD_EMBED if used_pad else trailing_text_hidden[t]
        dual_in = prev_audio_emb + text_cond
        t_out = talker_moe_step(dual_in, speaker_vec)
        codes_t, depth_logs = run_depth_code_predictor(t_out["h_talker"], t_out["c0"])
        all_codes.append(codes_t)
        # Nhúng trung bình của 4 sổ mã khung t làm phản hồi cho khung t+1
        frame_emb, _ = lookup_rvq_offset_mean(codes_t.reshape(NUM_QUANTIZERS, 1))
        prev_audio_emb = frame_emb[0]
        frame_records.append({
            "frame": t,
            "used_pad": used_pad,
            "dual_in": dual_in,
            "talker": t_out,
            "codes": codes_t,
            "depth_logs": depth_logs,
        })
    codes_matrix = np.stack(all_codes, axis=1)  # (NUM_QUANTIZERS, n_frames)
    return codes_matrix, frame_records

# ── [5A] Pre-Transformer (Sliding-Window Causal Mask) & ConvNeXt 2× ──────────
W_PRE_Q = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_PRE_K = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
W_PRE_V = rng.normal(0.0, 0.3, size=(TALKER_DIM, TALKER_DIM))
LAYER_SCALE = np.full(TALKER_DIM, 0.01, dtype=float)
W_UP_PRE = rng.normal(0.0, 0.35, size=(PRE_UPSAMPLE, TALKER_DIM, TALKER_DIM))

def build_sliding_causal_mask(T, window=SLIDING_WINDOW):
    mask = np.zeros((T, T), dtype=bool)
    for q in range(T):
        for k in range(T):
            if 0 <= q - k < window:
                mask[q, k] = True
    return mask

def run_pre_transformer_and_upsample(frame_embeds):
    T = frame_embeds.shape[0]
    mask = build_sliding_causal_mask(T, SLIDING_WINDOW)
    q = frame_embeds @ W_PRE_Q
    k = frame_embeds @ W_PRE_K
    v = frame_embeds @ W_PRE_V
    raw_scores = (q @ k.T) / np.sqrt(TALKER_DIM)
    masked_scores = np.where(mask, raw_scores, -1e9)
    attn_w = softmax(masked_scores, axis=-1)
    h_pre = frame_embeds + LAYER_SCALE * (attn_w @ v)
    # Causal Transposed Upsample ×2: mỗi khung t nở thành 2 bước con nhân quả
    up_list = []
    for t in range(T):
        for s in range(PRE_UPSAMPLE):
            up_list.append(h_pre[t] @ W_UP_PRE[s])
    h_up = np.stack(up_list, axis=0)  # (T * PRE_UPSAMPLE, TALKER_DIM)
    return {"mask": mask, "attn_w": attn_w, "h_pre": h_pre, "h_up": h_up}

# ── [6B] SnakeBeta Causal ConvNet & Chunked 24kHz Wave Decode ────────────────
SNAKE_ALPHA = np.linspace(0.2, 0.8, TALKER_DIM)
SNAKE_BETA = np.linspace(0.5, 1.1, TALKER_DIM)
W_DEC_STAGE1 = rng.normal(0.0, 0.3, size=(DEC_UPSAMPLES[0], TALKER_DIM, TALKER_DIM))
W_DEC_STAGE2 = rng.normal(0.0, 0.3, size=(DEC_UPSAMPLES[1], TALKER_DIM, 1))

def snake_beta(x, alpha=SNAKE_ALPHA, beta=SNAKE_BETA):
    a = np.exp(alpha)
    b = np.exp(beta)
    return x + (1.0 / (b + 1e-9)) * (np.sin(x * a) ** 2)

def run_snakebeta_decoder(h_up):
    # Stage 1: SnakeBeta -> upsample ×2
    s1_in = snake_beta(h_up)
    s1_list = []
    for t in range(s1_in.shape[0]):
        for r in range(DEC_UPSAMPLES[0]):
            s1_list.append(s1_in[t] @ W_DEC_STAGE1[r])
    s1_out = np.stack(s1_list, axis=0)
    # Stage 2: SnakeBeta -> upsample ×3 -> 1 kênh sóng âm
    s2_in = snake_beta(s1_out)
    wav_list = []
    for t in range(s2_in.shape[0]):
        for r in range(DEC_UPSAMPLES[1]):
            val = float((s2_in[t] @ W_DEC_STAGE2[r])[0])
            wav_list.append(val)
    wav = np.clip(np.array(wav_list, dtype=float), -1.0, 1.0)
    return wav

def decode_codes_to_wav(codes_matrix):
    frame_embeds, _ = lookup_rvq_offset_mean(codes_matrix)
    pre = run_pre_transformer_and_upsample(frame_embeds)
    return run_snakebeta_decoder(pre["h_up"])

def chunked_decode_codes(codes_matrix, chunk_size=2, left_context_size=1):
    total_frames = codes_matrix.shape[1]
    wav_chunks = []
    start = 0
    while start < total_frames:
        end = min(start + chunk_size, total_frames)
        ctx = left_context_size if start - left_context_size > 0 else start
        sub_codes = codes_matrix[:, start - ctx : end]
        sub_wav = decode_codes_to_wav(sub_codes)
        wav_chunks.append(sub_wav[ctx * TOTAL_UPSAMPLE :])
        start = end
    return np.concatenate(wav_chunks, axis=0)

# ── Chạy toàn trình và in báo cáo 1A → 6B ────────────────────────────────────
print("=== LAB QWEN3-OMNI & QWEN3-TTS-12HZ: TOÀN TRÌNH 1A → 6B BẰNG NUMPY ===")
print("MỤC TIÊU: chạy khép kín đường ống Thinker–Talker MoE + 16-RVQ + SnakeBeta Code2Wav,")
print("kiểm chứng hạ mẫu AuT 8×, TM-RoPE 3D, đồng bộ Dual-Track, Depth Predictor và bảng RVQ gộp.")

aut = run_aut_encoder(MEL_INPUT, cu_seqlens=(0, 1, 2))
say("1A", "AuT AUDIO ENCODER (3× CONV2D STRIDE-2 & WINDOWED ATTENTION)", [
    f"Mel spectrogram đầu vào: shape {MEL_INPUT.shape} (8 dải Mel × 16 khung thời gian).",
    f"Sau conv2d1 (stride=2): {aut['c1'].shape} → sau conv2d2: {aut['c2'].shape} → sau conv2d3: {aut['c3'].shape}.",
    f"Tỷ lệ nén thời gian: {MEL_FRAMES} → {aut['c3'].shape[1]} khung (giảm đúng 2³ = 8×).",
    f"Đặc trưng âm thanh sau Windowed Attn (cu_seqlens=(0,1,2)): shape {aut['out'].shape}",
    f"   khung âm thanh 0: {fmt_vec(aut['out'][0])}",
])

q_sample = aut["out"][0, :HEAD_DIM]
q_rot = apply_interleaved_tm_rope(q_sample, pos_thw=(3, 3, 3))
norm_before = float(np.linalg.norm(q_sample))
norm_after = float(np.linalg.norm(q_rot))
say("1B", "INTERLEAVED TM-ROPE 3D (TEMPORAL, HEIGHT, WIDTH @ 12.5 HZ)", [
    "Prod phân bổ 64 cặp theo mrope_section = [24, 20, 20]; lab đan xen 3 cặp (T, H, W).",
    f"Vector trước xoay : {fmt_vec(q_sample)} (‖q‖ = {norm_before:.6f})",
    f"Vector sau xoay   : {fmt_vec(q_rot)} (‖RoPE(q)‖ = {norm_after:.6f})",
    f"Độ lệch chuẩn Euclid |‖RoPE(q)‖ − ‖q‖| = {abs(norm_after - norm_before):.2e}.",
])

thinker = run_thinker_handoff(THINKER_TEXT_EMB)
say("2A", "THINKER MOE BACKBONE & TRÍCH NGANG TRẠNG THÁI ẨN LỚP 24", [
    f"Thinker sinh 2 token văn bản; trích ngang tại accept_hidden_layer = 24 (shape {thinker['h24'].shape}):",
    f"   trailing_text_hidden[0] (Lớp 24): {fmt_vec(thinker['h24'][0])}",
    f"   trailing_text_hidden[1] (Lớp 24): {fmt_vec(thinker['h24'][1])}",
    f"   So với tầng đỉnh sát LM Head    : {fmt_vec(thinker['htop'][0])}",
])

txt_proj = resize_text_mlp(THINKER_TEXT_EMB)
hid_proj = resize_hidden_mlp(thinker["h24"])
say("3A", "CẦU NỐI RESIZEMLP KÉP (2048 → 1024) & ĐIỀU KIỆN HÓA SPEAKER ID", [
    f"text_projection(embed)   shape {THINKER_TEXT_EMB.shape} → {txt_proj.shape}: {fmt_vec(txt_proj[0])}",
    f"hidden_projection(L24)   shape {thinker['h24'].shape} → {hid_proj.shape}: {fmt_vec(hid_proj[0])}",
    f"Hồ sơ giọng đọc speaker_id = 2302 (Ethan)  : {fmt_vec(SPEAKER_TABLE[2302])}",
    f"Hồ sơ giọng đọc speaker_id = 2301 (Chelsie): {fmt_vec(SPEAKER_TABLE[2301])}",
])

codes_matrix, frame_records = run_dual_track_frames(hid_proj, speaker_id=2302, n_frames=4)
say("3B", "ĐỒNG BỘ DUAL-TRACK STREAMING 12.5 HZ (4 KHUNG 80 MS, 2 TOKEN CHỮ)", [
    "Khung {} (80ms): nguồn chữ = {} → dual_in[:3] = {}".format(
        r["frame"],
        "tts_pad_embed" if r["used_pad"] else f"trailing_text_hidden[{r['frame']}]",
        fmt_vec(r["dual_in"][:3]),
    )
    for r in frame_records
] + [
    "Kết luận: Khung 0 phát ngay khi có chữ đầu tiên (độ trễ 97 ms); Khung 2–3 tự động dùng tts_pad_embed.",
])

f0_talker = frame_records[0]["talker"]
say("3C", "TALKER MOE (TOP-K ROUTING + SHARED EXPERT) & MÃ NGỮ NGHĨA c_0", [
    f"Khung 0: xác suất 6 chuyên gia = {fmt_vec(f0_talker['router_probs'])}",
    f"Chọn Top-{TOP_K_EXPERTS} chuyên gia {list(map(int, f0_talker['top_idx']))} → trọng số chuẩn hóa = {fmt_vec(f0_talker['norm_weights'])} (tổng = {np.sum(f0_talker['norm_weights']):.4f}).",
    f"Logits codec_head (8 mã): {fmt_vec(f0_talker['logits_c0'])} → chọn mã ngữ nghĩa c_0 = {f0_talker['c0']}.",
])

say("4A", "5-LAYER DEPTH CODE PREDICTOR (TỰ HỒI QUY MÃ ÂM HỌC TRONG 80 MS)", [
    f"Bước depth {d['step']} (chiều dài chuỗi nội bộ = {d['seq_len']}): dùng lm_head[{d['step']}] → sinh c_{d['code_idx']} = {d['code']}"
    for d in frame_records[0]["depth_logs"]
] + [
    f"Bộ mã đầy đủ của Khung 0 [c_0 .. c_3]: {codes_matrix[:, 0].tolist()} (prod sinh 16 mã [c_0 .. c_15]).",
])

rvq_mean, shifted_indices = lookup_rvq_offset_mean(codes_matrix)
say("6A", "BẢNG NHÚNG 16-GROUP RVQ HỢP NHẤT (CODE_OFFSET) & MEAN POOLING", [
    f"Ma trận mã 4 tầng × 4 khung:\n      {codes_matrix.tolist()}",
    f"Cộng code_offset {CODE_OFFSET.ravel().tolist()} → chỉ số tra bảng gộp 32 dòng:\n      {shifted_indices.tolist()}",
    f"Vector trung bình .mean(dim=0) của Khung 0: {fmt_vec(rvq_mean[0])}",
])

pre_res = run_pre_transformer_and_upsample(rvq_mean)
mask_rows = [
    f"   q={q} | " + " ".join("1" if pre_res["mask"][q, k] else "." for k in range(4))
    for q in range(4)
]
say("5A", f"CODE2WAV PRE-TRANSFORMER (SLIDING WINDOW = {SLIDING_WINDOW}) & CONVNEXT {PRE_UPSAMPLE}×", [
    f"Mặt nạ nhân quả cửa sổ trượt 4×4 (window = {SLIDING_WINDOW}, 1 = cho phép, . = chặn):",
] + mask_rows + [
    f"Sau Pre-Transformer (LayerScale=0.01): shape {pre_res['h_pre'].shape} → sau TransConv {PRE_UPSAMPLE}×: shape {pre_res['h_up'].shape}.",
])

wav_full = run_snakebeta_decoder(pre_res["h_up"])
wav_chunked = chunked_decode_codes(codes_matrix, chunk_size=2, left_context_size=1)
say("6B", f"SNAKEBETA CAUSAL CONVNET ({DEC_UPSAMPLES[0]}×{DEC_UPSAMPLES[1]} = {DEC_UPSAMPLES[0]*DEC_UPSAMPLES[1]}×, TỔNG {TOTAL_UPSAMPLE}×) → SÓNG ÂM", [
    f"4 khung mã @ 12.5 Hz × {TOTAL_UPSAMPLE} mẫu/khung = {len(wav_full)} mẫu sóng âm (prod: 4 × 1920 = 7680 mẫu @ 24 kHz).",
    f"8 mẫu sóng âm đầu tiên (nằm trong [-1, 1]): {fmt_vec(wav_full[:8])}",
    f"Giải mã theo từng chunk (chunk_size=2, left_context=1): thu được {len(wav_chunked)} mẫu liền mạch.",
])

# ── [KIỂM TRA & THỬ NGHIỆM] ─────────────────────────────────────────────────
print("\n[KIỂM TRA & THỬ NGHIỆM]")

# (1) Kiểm chứng bảng gộp code_offset vs 4 bảng rời rạc
rvq_sep = lookup_rvq_separate_tables_mean(codes_matrix)
diff_rvq = float(np.max(np.abs(rvq_mean - rvq_sep)))
print(f"  (1) Bảng gộp code_offset vs tra cứu từng bảng RVQ rời: lệch tối đa = {diff_rvq:.2e} (ngưỡng 0.00e+00).")
assert diff_rvq == 0.0

# (2) Kiểm chứng bảo toàn chuẩn TM-RoPE 3D và chuẩn hóa Top-K MoE
rope_err = abs(norm_after - norm_before)
moe_sum_err = abs(float(np.sum(f0_talker["norm_weights"])) - 1.0)
print(f"  (2) Sai số chuẩn 3D TM-RoPE = {rope_err:.2e}; sai số tổng trọng số Top-K Talker MoE = {moe_sum_err:.2e}.")
assert rope_err < 1e-12 and moe_sum_err < 1e-12

# (3) Kiểm chứng cô lập nhân quả & khớp Chunked Streaming Decode với Left Context
codes_perturbed = codes_matrix.copy()
codes_perturbed[:, 3] = (codes_perturbed[:, 3] + 3) % CODEBOOK_SIZE  # đổi khung cuối (frame 3)
wav_perturbed = decode_codes_to_wav(codes_perturbed)
causal_leak = float(np.max(np.abs(wav_full[: 3 * TOTAL_UPSAMPLE] - wav_perturbed[: 3 * TOTAL_UPSAMPLE])))
chunk_diff = float(np.max(np.abs(wav_full - wav_chunked)))
print(f"  (3) Đổi mã khung cuối (frame 3): rò rỉ ngược về 3 khung đầu = {causal_leak:.2e} (nhân quả tuyệt đối).")
print(f"      Full decode vs Chunked decode (chunk=2, left_context=1): lệch tối đa = {chunk_diff:.2e} (< 1e-12).")
assert causal_leak == 0.0
assert chunk_diff < 1e-12

print("PASS: toàn trình 1A→6B của Qwen3-Omni & Qwen3-TTS-12Hz chạy đúng trên số thật.")

Python outputSẵn sàng.

Nguồn công bố & mã nguồn tham chiếu

Qwen/Qwen3-Omni-30B-A3B-Instruct (Hugging Face) ↗ Model card, config.json (thinker_config, talker_config, code2wav_config) modeling_qwen3_omni_moe.py (HuggingFace Transformers) ↗ Mã nguồn AuT, Thinker, Talker MoE, 5-Layer Code Predictor & SnakeBeta Code2Wav Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice (Hugging Face) ↗ Tokenizer 12.5 Hz 16-codebook, Dual-Track 97ms streaming & điều khiển giọng nói

Danh sách bài học (11)