跑完 10 个 epoch，按下生成按钮，听到那段 AI 写出来的钢琴旋律时，我愣了好几秒。这不是什么商业级作品，但那种"机器真的学会了音乐"的感觉，挺震撼的。

---

一、为什么是 Transformer？

之前折腾完 EnCodec，把 72 首钢琴曲切成了 token 序列。当时就在想：这玩意儿跟 GPT 处理文本有什么区别？

其实没有本质区别。

GPT 把"今天天气不错"变成 [2345, 6789, 1234]，然后预测下一个词；Music Transformer 把一段钢琴旋律变成 [42, 156, 893]，然后预测下一个音频 token。

核心就一件事：自回归建模。

给你前 N 个 token，让你猜第 N+1 个。训练的时候用交叉熵损失，生成的时候逐个采样。就是这么简单，但也就是这么优雅。

---

二、训练样本怎么造？

拿到 EnCodec 编码出来的数据，形状是 [4, 750]：

• 4：4 个 codebook（量化层）
• 750：10 秒音频的时间步

第一步是转置，变成 [750, 4]。为什么？因为 Transformer 处理序列是按时间步来的，每个时间步有 4 个 token（分别对应 4 个量化层）。

然后就是经典的自回归 shift：

x = frames[:-1]   # [T-1, K]  输入
y = frames[1:]    # [T-1, K]  标签

x 的第 2 行等于 y 的第 1 行，x 的第 3 行等于 y 的第 2 行。模型的任务就是：看到 x 的第 1 行，预测出 y 的第 1 行。

这个设计跟 GPT 一模一样，只不过 GPT 的每个 token 是一个词，我们这里的每个"token"是4 个 codebook 的组合。

---

三、模型架构：多 Codebook 的坑

最核心的代码在这里：

class MusicTransformer(nn.Module):
    def __init__(self, vocab_size=1024, num_codebooks=4, 
                 embed_dim=512, max_seq_len=8192,
                 num_layers=6, num_heads=8, dropout=0.2):
        super().__init__()
        self.num_codebooks = num_codebooks
        self.embed_dim = embed_dim

        # 关键：每个 codebook 独立 embedding
        self.embeddings = nn.ModuleList([
            nn.Embedding(vocab_size, embed_dim)
            for _ in range(num_codebooks)
        ])
        
        self.pos_embedding = nn.Embedding(max_seq_len, embed_dim)
        
        decoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=num_heads,
            dim_feedforward=embed_dim * 4, dropout=dropout, 
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(decoder_layer, num_layers=num_layers)
        
        self.output_head = nn.Linear(embed_dim, vocab_size)

为什么每个 codebook 要独立的 embedding？

一开始我以为可以共享一个 embedding 表，后来想明白了：4 个 codebook 根本不是同一个语义空间。

• 第 1 个 codebook 记录的是低频主旋律（最重要）
• 第 2 个 codebook 补充中频细节
• 第 3、4 个 codebook 捕捉高频泛音

它们虽然都是 0-1023 的整数，但含义完全不同。就像你不能把"苹果"和"红色"的 embedding 混用一样。

---

四、Forward 的逻辑：展平与掩码

def forward(self, x):
    B, T, K = x.shape  # B=batch, T=时间步, K=codebook 数
    
    # 1. 每个 codebook 单独 embedding
    h_list = []
    for k in range(K):
        h_list.append(self.embeddings[k](x[:, :, k]))
    
    # 2. 堆叠并展平：[B, T, K, D] → [B, T*K, D]
    h = torch.stack(h_list, dim=2).reshape(B, T * K, -1)
    
    # 3. 加位置编码
    pos_ids = torch.arange(T * K, device=x.device)
    h = h + self.pos_embedding(pos_ids)[None, :, :]
    
    # 4. 构造因果掩码（上三角 -inf）
    mask = torch.triu(
        torch.full((T * K, T * K), float('-inf'), device=x.device),
        diagonal=1
    )
    
    # 5. Transformer 前向
    h = self.transformer(h, mask)
    
    # 6. 分类头
    logits = self.output_head(h)
    
    # 7. 还原形状：[B, T*K, vocab] → [B, T, K, vocab]
    logits = logits.reshape(B, T, K, -1)
    
    return logits

几个关键点：

1. 展平：把 [B, T, K] 变成 [B, T*K]，这样 Transformer 就能当普通序列处理了。一个 10 秒音频、4 个 codebook，展平后就是 3000 个"词"。
2. 位置编码：直接用一个 [max_seq_len, embed_dim] 的 embedding 表，按展平后的位置索引查表。简单粗暴，但有效。
3. 因果掩码：torch.triu 生成上三角矩阵，对角线以上全是 -inf。这样第 i 个 token 就只能看到 i 之前的 token，看不到未来。这是自回归模型的铁律。

---

五、损失函数：为什么加权？

训练部分的代码：

weights = [1.0, 0.5, 0.25, 0.1]
loss = 0

for k in range(4):
    l_k = criterion(
        logits[:, :, k, :].reshape(-1, 1024),
        y[:, :, k].reshape(-1)
    )
    loss += weights[k] * l_k

为什么第一个 codebook 权重最高？

因为 EnCodec 的 RVQ 是逐层细化的：

• 第 1 层抓的是整体轮廓（低频、主旋律）
• 第 2 层补的是细节纹理
• 第 3、4 层补的是高频泛音

如果 4 个 codebook 权重一样，模型可能会被高频噪声带偏，学不到音乐的主干结构。加权损失函数本质上是在告诉模型：先学会旋律，再学音色。

---

六、训练策略：AdamW + Warmup

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    weight_decay=1e-2
)

warmup_steps = int(total_steps * 0.05)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

为什么用 AdamW 而不是 Adam？

AdamW 把权重衰减（weight decay）和梯度更新解耦了。对于 Transformer 这种大模型，解耦的正则化效果更好，不容易过拟合。

为什么需要 Warmup？

训练初期梯度不稳定，直接上大学习率容易爆炸。Warmup 让学习率从 0 慢慢升到目标值（前 5% 步数），然后再线性衰减。这是 GPT 的标准配置，搬过来直接用。

---

七、训练日志解读

跑完 10 个 epoch，loss 曲线是这样的：

CB1（第一个 codebook）的 loss 从 7.1 降到 3.3。

这个曲线很健康：前期下降快，后期趋于平稳，没有震荡。说明模型确实在学东西，而不是在梯度爆炸的边缘反复横跳。

---

八、生成：自回归采样

@torch.no_grad()
def generate(model, start_tokens, max_new_tokens=200, temperature=1.0):
    model.eval()
    x = start_tokens.to(device)
    
    for _ in range(max_new_tokens):
        x_cond = x[:, -model.max_seq_len:]  # 只保留最近 max_seq_len 个 token
        
        logits = model(x_cond)  # [B, T, K, vocab]
        
        next_tokens = []
        for k in range(model.num_codebooks):
            last_logits = logits[:, -1, k, :] / temperature
            probs = F.softmax(last_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            next_tokens.append(next_token)
        
        next_tokens = torch.stack(next_tokens, dim=-1)  # [B, 1, K]
        x = torch.cat([x, next_tokens], dim=1)
    
    return x

几个关键设计：

1. 上下文窗口限制：x_cond = x[:, -model.max_seq_len:]。Transformer 的注意力计算是 O(T²) 的，不能无限长。只保留最近的 8192 个 token。
2. Temperature 采样：last_logits / temperature。temperature > 1 会让分布更平，生成更多样；temperature < 1 会让分布更尖，生成更保守。我试过 1.1，效果不错。
3. 逐个 codebook 采样：每个 codebook 独立采样，然后拼起来。不能一起采样，因为每个 codebook 的词汇表是独立的。

---

九、解码：从 Token 回到音频

生成的 token 形状是 [B, T, K]，需要转成 EnCodec 能接受的格式：

codes = generated_tokens[0].permute(1, 0).unsqueeze(0).to(device)
# [T, K] → [K, T] → [1, K, T]

encoded_frames = [(codes, None)]  # 关键！

with torch.no_grad():
    wav = codec.decode(encoded_frames)

为什么是 [(codes, None)]？

EnCodec 的 decode 函数期望的输入是 [(codes, scale), ...] 的列表，每个元素是一个时间帧。scale 是编码时的缩放因子，生成时没有，所以填 None。

然后用 soundfile 保存：

audio = wav[0, 0].cpu().numpy()
sf.write("generated_music.wav", audio, 24000)

十、跑出来的效果怎么样？

说实话，不惊艳，但能听。

生成的钢琴曲有明显的旋律结构，不会是一团噪声。但：

• 和声进行比较单调
• 节奏有时候会乱
• 高音区偶尔会有刺耳的泛音

原因分析：

1. 数据量不够：72 首曲子，10 个 epoch，模型还没见过足够多的音乐模式
2. 模型太小：6 层、512 维，跟 GPT-3 比就是玩具级别
3. 没有 fine-tuning：直接从头训练，没有预训练基础

但方向是对的。加数据、加层数、调超参，效果会越来越好。

---

十一、踩过的坑

1. OOM（显存爆炸）

第一次跑 batch_size=16，直接 OOM。改成 8 就好了。Transformer 的显存占用是 O(B × T²)，序列长度翻倍，显存要翻 4 倍。

2. 梯度爆炸

没加梯度裁剪的时候，loss 突然跳到 nan。加上 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) 就稳了。

3. 生成全是同一个 token

temperature 设太低（0.5），模型变得过于保守，一直重复同一个 token。调到 1.1 之后多样性就上来了。

---

十二、下一步怎么搞？

1. 加数据：Maestro 全量数据有 200 小时，现在只用了 72 首
2. 加模型容量：层数从 6 加到 12，embed_dim 从 512 加到 768
3. 加条件控制：比如输入"欢快的 C 大调"，生成对应风格的音乐
4. 加 MIDI 对齐：用 Maestro 的 MIDI 标注做辅助监督，让模型学到更清晰的音乐结构

---

十三、写在最后

这套流程走下来，我对"AI 音乐生成"的理解深了一层。

以前觉得音乐生成是黑魔法，现在知道它就是：音频→token→自回归建模→生成 token→解码回音频。没有什么神秘的，就是工程 + 算法。

但工程细节决定上限。怎么切数据、怎么加权 loss、怎么调 temperature、怎么处理 OOM，这些琐碎的东西堆起来，才有一段能听的音乐。

代码已经跑通了，接下来就是堆算力、调参数、加数据。AI 写歌这条路，才刚开始。

---

有问题欢迎提 issue 或者评论区留言。