用VQ-VAE离散标记重塑人体姿态：关键点组合关系编码新思路

DatapiTHU · 2025 年5 月 8 日 22:17

探索人体姿态估计新范式：利用VQ-VAE将姿态转化为离散标记组合，有效捕捉关键点间关系，提升模型泛化能力。

原文标题：用离散标记重塑人体姿态：VQ-VAE实现关键点组合关系编码

原文作者：数据派THU

原文链接： http://mp.weixin.qq.com/s?__biz=MzI1MjQ2OTQ3Ng==&mid=2247656833&idx=2&sn=142afab91faa3940b18c6461c4c552d9&

冷月清谈：

本文介绍了一种使用向量量化自编码器(VQ-VAE)以离散标记重塑人体姿态的方法。传统人体姿态估计方法通常将关键点作为独立单元处理，忽略了人体骨骼结构中固有的关键点间组合关系。该方法将姿态表示为一组学习到的离散标记组合，能够有效捕获关键点之间的语义和空间依赖关系，从而更好地泛化到未见过的姿态配置，并便于整合到下游任务中。文章详细介绍了该方法第一阶段的实现过程，包括组合编码器、向量量化码本和姿态解码器的训练。通过实验验证，采用分阶段训练策略或结合承诺损失调整与EMA更新等技术，可以有效缓解码本崩溃问题，确保码本的充分利用和模型的稳健学习。这种基于离散标记的姿态表示方法为后续的姿态分析和理解任务提供了新的可能性。

怜星夜思：

1、文章中提到直接端到端训练 VQ-VAE 容易出现“码本崩溃”现象，导致模型只学习到平均姿态。除了文章中提到的预训练编码器外，还有哪些方法可以缓解或者避免码本崩溃？
2、文章中使用了合成的火柴人数据集，这虽然简化了实验，但与真实人体姿态相比存在差距。你认为在真实人体姿态数据集上应用这种方法，会遇到哪些新的挑战？又该如何应对？
3、文章中提到基于离散标记的姿态表示方法，为后续的姿态分析和理解任务提供了新的可能性。除了文章中提到的动作识别和姿态分类，你还能想到哪些其他的下游任务可以利用这种离散的姿态表示？这种表示方法在这些任务中可能有哪些优势？

原文内容

来源：Deephub Imba
本文共5500字，建议阅读6分钟
本文介绍了使用离散标记重塑人体姿态的过程。

在人体姿态估计领域，传统方法通常将关键点作为基本处理单元，这些关键点在人体骨架结构上代表关节位置（如肘部、膝盖和头部）的空间坐标。现有模型对这些关键点的预测主要采用两种范式：直接通过坐标回归或间接通过热图（heat map，即图像空间中的密集概率分布）进行估计。尽管这些方法在实际应用中取得了显著效果，但它们往往将每个关键点作为独立单元处理，未能充分利用人体骨架结构中固有的关键点间组合关系。

如果我们转换思路，将姿态表示为一组学习到的、离散的标记（token）组合，这些标记不仅仅编码原始坐标或热图信息，而是捕获关键点之间的共享模式、对称性和结构化关系，会带来怎样的优势？

受2023年发表的研究论文《Human Pose As Compositional Tokens》启发，本文构建了一个姿态重建模型，实现了上述概念。我们将详细介绍该方法第一阶段的实现过程：训练组合编码器（Compositional Encoder）、向量量化（Vector Quantization, VQ）码本和姿态解码器（Pose Decoder）——这些组件共同构成了一个用于学习姿态结构紧凑、离散表示的系统。

与传统的热图或坐标回归方法相比，基于标记的表征方法能够有效捕获关键点之间的语义和空间依赖关系。这种表征形式使模型能够更好地泛化到未见过的姿态配置，并且便于将这些离散标记整合到下游任务中，如动作识别或姿态分类。

合成火柴人数据集

为简化实验过程并专注于方法本身，我们创建了一个火柴人的合成数据集。每个火柴人实例由13个二维关键点精确定义，这些关键点包括头部、颈部、肩膀、肘部、手部、腰部、膝盖和脚部。该数据集采用即时生成方式，支持随机但符合人体结构约束的肢体配置、微小的姿态变化以及基于旋转的数据增强。

图1：火柴人示例，其中标注了关键点位置

基于组合编码的VQ-VAE

本文的核心目标是将每个姿态压缩为一组离散的标记（token），每个标记对应于从共享码本中学习的潜在表示。

图2：PCT（Pose as Compositional Tokens）架构的训练阶段 - 来源：https://arxiv.org/abs/2303.11638

辅助模块

 
class MLPBlock(nn.Module):
def __init__(self, dim, inter_dim, dropout_ratio):
super().__init__()
self.ff = nn.Sequential(

nn.Linear(dim, inter_dim),

nn.GELU(),

nn.Dropout(dropout_ratio),

nn.Linear(inter_dim, dim),

nn.Dropout(dropout_ratio)

)
def forward(self, x):

return self.ff(x)
class MixerLayer(nn.Module):

def init(self,

hidden_dim,

hidden_inter_dim,

token_dim,

token_inter_dim,

dropout_ratio):

super().init()
self.layernorm1 = nn.LayerNorm(hidden_dim)

self.MLP_token = MLPBlock(token_dim, token_inter_dim, dropout_ratio)

self.layernorm2 = nn.LayerNorm(hidden_dim)

self.MLP_channel = MLPBlock(hidden_dim, hidden_inter_dim, dropout_ratio)
def forward(self, x):

y = self.layernorm1(x)

y = y.transpose(2, 1)

y = self.MLP_token(y)

y = y.transpose(2, 1)

z = self.layernorm2(x + y)

z = self.MLP_channel(z)

out = x + y + z

return out

MLPMixer层是《Human Pose as Compositional Tokens》论文作者提出的用于将关键点信息融合成潜在向量的核心组件。

组合编码器

class CompositionalEncoder(nn.Module):
def __init__(self, numberOfKeypoints=11, dimensionOfKeypoints=2, linearProjectionSize=128, numberOfMixerBlocks=4, codebookTokenDimension=64, internalMixerSize=64, internalMixerTokenSize=32, mixerDropout=0.1):
super(CompositionalEncoder, self).__init__()
self.numberOfKeypoints = numberOfKeypoints # K

self.dimensionOfKeypoints = dimensionOfKeypoints # D

self.linearProjectionSize = linearProjectionSize # H

self.numberOfMixerBlocks = numberOfMixerBlocks # N

self.codebookTokenDimension = codebookTokenDimension # M

self.internalMixerSize = internalMixerSize

self.internalMixerTokenSize = internalMixerTokenSize

self.mixerDropout = mixerDropout
self.initial_linear = nn.Linear(self.dimensionOfKeypoints,

self.linearProjectionSize) # 从BxKxD投影到BxKxH
self.mixer_layers = nn.ModuleList([MixerLayer(self.linearProjectionSize,

self.internalMixerSize,

self.numberOfKeypoints,

self.internalMixerTokenSize,

self.mixerDropout) for _ in range(self.numberOfMixerBlocks)]) # BxKxH
self.mixer_layer_norm = nn.LayerNorm(self.linearProjectionSize) # BxKxH
self.token_linear = nn.Linear(self.numberOfKeypoints,

self.codebookTokenDimension) # BxHxK -> BxHxM
self.feature_embed = nn.Linear(self.linearProjectionSize,

self.codebookTokenDimension)
def forward(self, x):
之前: BxDxK
x = x.transpose(2,1)
之后: BxKxD
之前: BxKxD
x = self.initial_linear(x)
之后: BxKxH
之前: BxKxH
for mixer in self.mixer_layers:

x = mixer(x)
之后: BxKxH
之前: BxKxH
x = self.mixer_layer_norm(x)
之后: BxKxH
之前: BxKxH
x = x.transpose(2,1)
之后: BxHxK
之前: BxHxK
x = self.token_linear(x)
之后: BxHxM
之前: BxHxM
x = x.transpose(2,1)
之后: BxMxH
之前: BxMxH
x = self.feature_embed(x)
之后: BXMxM
return x

编码器接收一组二维关键点坐标，通过基于MLP-Mixer架构设计的网络结构将这些坐标转换为M个潜在标记特征。具体而言，关键点首先被嵌入到高维空间，然后在关节和特征维度之间进行混合处理，最终投影到形状为B × M × D（批量大小×标记数量×标记维度）的输出特征空间。

EMA码本（VQ层）

class CodebookVQ(nn.Module):
def __init__(self, codebookDimension, numberOfCodebookTokens, decay=0.99, epsilon=1e-5):
super(CodebookVQ, self).__init__()
self.codebookDimension = codebookDimension

self.numberOfCodebookTokens = numberOfCodebookTokens

self.decay = decay

self.epsilon = epsilon
self.register_buffer(‘codebook’, torch.empty(numberOfCodebookTokens, codebookDimension))

self.codebook.data.normal_()
self.register_buffer(‘ema_cluster_size’, torch.zeros(numberOfCodebookTokens))

self.register_buffer(‘ema_w’, torch.empty(numberOfCodebookTokens, codebookDimension))

self.ema_w.data.normal_()
def forward(self, encode_feat):
M = encode_feat.shape[1]

B = encode_feat.shape[0]

encode_feat = encode_feat.view(-1, self.codebookDimension)  # [B*M, M]
计算与码本条目的距离
distances = (

encode_feat.pow(2).sum(1, keepdim=True)

2 * encode_feat @ self.codebook.t()


self.codebook.pow(2).sum(1)

)  # [B*M, num_tokens]

找到最近的码本索引
encoding_indices = torch.argmin(distances, dim=1)  # [BM]

encodings = F.one_hot(encoding_indices, self.numberOfCodebookTokens).type(encode_feat.dtype)  # [BM, num_tokens]
量化输出
quantized = encodings @ self.codebook  # [B*M, M]

quantized = quantized.view_as(encode_feat)  # 重塑回原始输入形状
if self.training:
EMA更新
ema_counts = encodings.sum(0)  # [num_tokens]

dw = encodings.t() @ encode_feat  # [num_tokens, M]
self.ema_cluster_size.mul_(self.decay).add_(ema_counts, alpha=1 - self.decay)

self.ema_w.mul_(self.decay).add_(dw, alpha=1 - self.decay)
n = self.ema_cluster_size.sum()

cluster_size = (

(self.ema_cluster_size + self.epsilon)

/ (n + self.numberOfCodebookTokens * self.epsilon) * n

)
self.codebook.data = self.ema_w / cluster_size.unsqueeze(1)
quantized = quantized.view(B, M, M)

encoding_indices = encoding_indices.view(B, M)
return quantized, encoding_indices

潜在标记通过向量量化层进行离散化处理，该层采用指数移动平均（EMA）方法更新码本。在这一过程中，每个连续的标记向量都被码本中最相近的离散代码向量替换，从而将姿态表示转化为一组符号化的离散表示。具体实现中：

码本包含num_codes个代码向量条目
每个输入标记根据L2距离独立选择最相近的码本向量
码本在训练过程中通过EMA机制进行自我更新，确保码本适应训练数据分布

姿态解码器

class PoseDecoder(nn.Module):
def __init__(self, codebookTokenDimension=64, numberOfKeypoints=11, keypointDimension=2, hiddenDimensionSize=128, numberOfMixerBlocks=4, mixerInternalDimensionSize=64, mixerTokenInternalDimensionSize=128, mixerDropout=0.1):
super(PoseDecoder, self).__init__()
self.codebookTokenDimension = codebookTokenDimension

self.numberOfKeypoints = numberOfKeypoints

self.keypointDimension = keypointDimension

self.hiddenDimensionSize = hiddenDimensionSize

self.mixerInternalDimensionSize = mixerInternalDimensionSize

self.mixerTokenInternalDimensionSize = mixerTokenInternalDimensionSize

self.mixerDropout = mixerDropout

self.numberOfMixerBlocks = numberOfMixerBlocks
self.linear_token = nn.Linear(self.codebookTokenDimension, self.numberOfKeypoints)

self.initial_linear = nn.Linear(self.codebookTokenDimension, self.hiddenDimensionSize)
self.mixer_layers = nn.ModuleList([MixerLayer(self.hiddenDimensionSize, self.mixerInternalDimensionSize, self.numberOfKeypoints, self.mixerTokenInternalDimensionSize, self.mixerDropout) for _ in range(self.numberOfMixerBlocks)])
self.decoder_layer_norm = nn.LayerNorm(self.hiddenDimensionSize)
self.recover_embed = nn.Linear(self.hiddenDimensionSize, self.keypointDimension)
def forward(self, x):
之前: BxMxM
x = self.linear_token(x)
之后: BxMxK
之前: BxMxK
x = x.transpose(2,1)
之后: BxKxM
之前: BxKxM
x = self.initial_linear(x)
之后: BxKxH
之前: BxKxH
for mixer in self.mixer_layers:

x = mixer(x)
之后: BxKxH
之前: BxKxH
x = self.decoder_layer_norm(x)
之后: BxKxH
之前: BxKxH
x = self.recover_embed(x)
之后: BxKxD
之后: BxKxD
x = x.transpose(2,1)
之后: BxDxK
return x

解码器模块负责接收量化后的标记并重建原始关键点坐标。其结构设计与编码器形成镜像对称，通过多层MLP处理和标记-特征混合操作，最终将离散表示投影回每个关节的二维坐标空间。

预训练码本和姿态解码器

在这一阶段，我们采用自监督重建策略来训练码本和姿态解码器。训练过程中使用两类关键损失函数：姿态重建损失和码本承诺损失：

图3：PCT模型的损失函数设计 - 来源：https://arxiv.org/abs/2303.11638

重建损失采用目标姿态与预测姿态之间的平滑L1距离度量。承诺损失（Commitment loss）则确保编码器输出的连续向量与其被映射到的量化码本条目保持接近，这对于防止编码器忽略码本至关重要。如果缺少承诺损失，编码器可能会生成与实际码本条目相距甚远的任意向量，导致模型无法有效学习离散表示。

初始实验中，我们尝试同时训练组合编码器、码本和姿态解码器：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
seed = 1

random.seed(seed)

np.random.seed(seed)

torch.manual_seed(seed)

torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True

torch.backends.cudnn.benchmark = False
— 数据集 —
dataset = StickFigureDataset(

num_samples=10000,

image_size=64,

core_radius=1,

limb_radius=5

)

loader = DataLoader(dataset, batch_size=64, shuffle=True)
numberOfKeypoints = 13

dimensionOfKeypoints = 2

linearProjectionSize = 256

numberOfMixerBlocks = 16

codebookTokenDimension = 64

internalMixerSize = 64

internalMixerTokenSize = 32

mixerDropout = 0.1
encoder = CompositionalEncoder(numberOfKeypoints=numberOfKeypoints,

dimensionOfKeypoints=dimensionOfKeypoints,

linearProjectionSize=linearProjectionSize,

numberOfMixerBlocks=numberOfMixerBlocks,

codebookTokenDimension=codebookTokenDimension,

internalMixerSize=internalMixerSize,

internalMixerTokenSize=internalMixerTokenSize,

mixerDropout=mixerDropout).to(device)

codebook = CodebookVQ(codebookDimension=codebookTokenDimension,

numberOfCodebookTokens=codebookTokenDimension,

decay=0.99,

epsilon=1e-5).to(device)

decoder = PoseDecoder(codebookTokenDimension=codebookTokenDimension,

numberOfKeypoints=numberOfKeypoints,

keypointDimension=dimensionOfKeypoints,

hiddenDimensionSize=linearProjectionSize,

numberOfMixerBlocks=numberOfMixerBlocks,

mixerInternalDimensionSize=internalMixerSize,

mixerTokenInternalDimensionSize=internalMixerTokenSize,

mixerDropout=mixerDropout).to(device)
optimizer = torch.optim.Adam(

list(encoder.parameters()) +

list(decoder.parameters()),

lr=1e-4

)
encoder.train()

codebook.train()

decoder.train()
num_epochs = 20

beta = 0.25
for epoch in range(num_epochs):

epoch_loss = 0.0

num_batches = 0
for imgs, gt_keypoints in loader:

keypoints = gt_keypoints.permute(0, 2, 1).to(device)
optimizer.zero_grad()
token_feats = encoder(keypoints)

quantized, _ = codebook(token_feats)

reconstructed = decoder(quantized)
recon_loss = F.smooth_l1_loss(reconstructed, keypoints)

commitment_loss = F.mse_loss(quantized.detach(), token_feats)

loss = recon_loss + beta * commitment_loss
loss.backward()

optimizer.step()
epoch_loss += loss.item()

num_batches += 1
avg_loss = epoch_loss / num_batches

print(f"Epoch {epoch+1}/{num_epochs} - Average Loss: {avg_loss:.4f}")

然而，我们观察到学习过程很快趋于饱和，模型仅学习了"平均"二维姿态关键点。深入分析发现，这是因为组合编码器的潜在向量对所有输入姿态都映射到码本中的相同条目，即发生了所谓的"码本崩溃"（codebook collapse）现象。

Epoch 1/20 - Average Loss: 18.5585
Epoch 2/20 - Average Loss: 14.4645
Epoch 3/20 - Average Loss: 11.6697
Epoch 4/20 - Average Loss: 9.7948
Epoch 5/20 - Average Loss: 8.3735
Epoch 6/20 - Average Loss: 7.2084
Epoch 7/20 - Average Loss: 6.5090
Epoch 8/20 - Average Loss: 6.0753
Epoch 9/20 - Average Loss: 5.6844
Epoch 10/20 - Average Loss: 5.4609
Epoch 11/20 - Average Loss: 5.3141
Epoch 12/20 - Average Loss: 5.2014
Epoch 13/20 - Average Loss: 5.1606
Epoch 14/20 - Average Loss: 5.1018
Epoch 15/20 - Average Loss: 5.1005
Epoch 16/20 - Average Loss: 5.0874
Epoch 17/20 - Average Loss: 5.0735
Epoch 18/20 - Average Loss: 5.0267
Epoch 19/20 - Average Loss: 5.0190
Epoch 20/20 - Average Loss: 5.0247

图4：姿态编码器/解码器的端到端训练结果。注意所有输入都产生相似的解码结果，表明发生了"码本崩溃"现象。

为避免编码器输出被强制映射到单一码本条目的问题，我们采用了分阶段训练策略：首先在不使用码本的情况下训练编码器和解码器，然后冻结编码器权重，并使用码本重新训练新的解码器：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
seed = 1

random.seed(seed)

np.random.seed(seed)

torch.manual_seed(seed)

torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True

torch.backends.cudnn.benchmark = False
— 数据集 —
dataset = StickFigureDataset(

num_samples=10000,

image_size=64,

core_radius=1,

limb_radius=5

)

loader = DataLoader(dataset, batch_size=64, shuffle=True)
numberOfKeypoints = 13

dimensionOfKeypoints = 2

linearProjectionSize = 256

numberOfMixerBlocks = 16

codebookTokenDimension = 64

internalMixerSize = 64

internalMixerTokenSize = 32

mixerDropout = 0.1
encoder = CompositionalEncoder(numberOfKeypoints=numberOfKeypoints, dimensionOfKeypoints=dimensionOfKeypoints, linearProjectionSize=linearProjectionSize, numberOfMixerBlocks=numberOfMixerBlocks, codebookTokenDimension=codebookTokenDimension, internalMixerSize=internalMixerSize, internalMixerTokenSize=internalMixerTokenSize, mixerDropout=mixerDropout).to(device)

codebook = CodebookVQ(codebookDimension=codebookTokenDimension, numberOfCodebookTokens=codebookTokenDimension, decay=0.99, epsilon=1e-5).to(device)

decoder = PoseDecoder(codebookTokenDimension=codebookTokenDimension, numberOfKeypoints=numberOfKeypoints, keypointDimension=dimensionOfKeypoints, hiddenDimensionSize=linearProjectionSize, numberOfMixerBlocks=numberOfMixerBlocks, mixerInternalDimensionSize=internalMixerSize, mixerTokenInternalDimensionSize=internalMixerTokenSize, mixerDropout=mixerDropout).to(device)
optimizer = torch.optim.Adam(

list(encoder.parameters()) +

list(decoder.parameters()),

lr=1e-4

)
encoder.train()

codebook.train()

decoder.train()
num_epochs = 20

beta = 0.25
skipQuantization = True
print(“Encoder pretraining”)

for epoch in range(num_epochs):

epoch_loss = 0.0

num_batches = 0
for imgs, gt_keypoints in loader:

keypoints = gt_keypoints.permute(0, 2, 1).to(device)  # [B, 2, 11]
optimizer.zero_grad()
token_feats = encoder(keypoints)                     # (B, M, M)

reconstructed = decoder(token_feats)                 # (B, K, D)
loss = F.smooth_l1_loss(reconstructed, keypoints)
loss.backward()

optimizer.step()
epoch_loss += loss.item()

num_batches += 1
avg_loss = epoch_loss / num_batches

print(f"Epoch {epoch+1}/{num_epochs} - Average Loss: {avg_loss:.4f}")
初始训练后冻结编码器
for param in encoder.parameters():

param.requires_grad = False
重置解码器（这一行已经做到了）
decoder = PoseDecoder(codebookTokenDimension=codebookTokenDimension, numberOfKeypoints=numberOfKeypoints, keypointDimension=dimensionOfKeypoints, hiddenDimensionSize=linearProjectionSize, numberOfMixerBlocks=numberOfMixerBlocks, mixerInternalDimensionSize=internalMixerSize, mixerTokenInternalDimensionSize=internalMixerTokenSize, mixerDropout=mixerDropout).to(device)
更新优化器，只包括解码器（如有需要，可选择包括码本）
optimizer = torch.optim.Adam(

list(decoder.parameters()),

lr=1e-4

)
print(“Codebook and Decoder training”)

for epoch in range(num_epochs):

epoch_loss = 0.0

num_batches = 0
for imgs, gt_keypoints in loader:

keypoints = gt_keypoints.permute(0, 2, 1).to(device)
optimizer.zero_grad()
token_feats = encoder(keypoints)                     # (B, M, M)

quantized, _ = codebook(token_feats)                 # (B, M, M)

reconstructed = decoder(quantized)                   # (B, K, D)
recon_loss = F.smooth_l1_loss(reconstructed, keypoints)

commitment_loss = F.mse_loss(quantized.detach(), token_feats)

loss = recon_loss + beta * commitment_loss
loss.backward()

optimizer.step()
epoch_loss += loss.item()

num_batches += 1
avg_loss = epoch_loss / num_batches

print(f"Epoch {epoch+1}/{num_epochs} - Average Loss: {avg_loss:.4f}")

分阶段训练策略的训练日志显示了明显改善的学习曲线：

Encoder pretraining
Epoch 1/20 - Average Loss: 18.5178
Epoch 2/20 - Average Loss: 14.1350
Epoch 3/20 - Average Loss: 10.7014
Epoch 4/20 - Average Loss: 8.3755
Epoch 5/20 - Average Loss: 6.5254
Epoch 6/20 - Average Loss: 4.8045
Epoch 7/20 - Average Loss: 3.7144
Epoch 8/20 - Average Loss: 2.9114
Epoch 9/20 - Average Loss: 2.2571
Epoch 10/20 - Average Loss: 1.7662
Epoch 11/20 - Average Loss: 1.4548
Epoch 12/20 - Average Loss: 1.2346
Epoch 13/20 - Average Loss: 1.0853
Epoch 14/20 - Average Loss: 0.9722
Epoch 15/20 - Average Loss: 0.9048
Epoch 16/20 - Average Loss: 0.8413
Epoch 17/20 - Average Loss: 0.7932
Epoch 18/20 - Average Loss: 0.7520
Epoch 19/20 - Average Loss: 0.7124
Epoch 20/20 - Average Loss: 0.6845
Codebook and Decoder training

Epoch 1/20 - Average Loss: 18.7635

Epoch 2/20 - Average Loss: 14.3140

Epoch 3/20 - Average Loss: 10.7931

Epoch 4/20 - Average Loss: 8.4169

Epoch 5/20 - Average Loss: 6.4152

Epoch 6/20 - Average Loss: 4.8894

Epoch 7/20 - Average Loss: 3.9022

Epoch 8/20 - Average Loss: 3.1704

Epoch 9/20 - Average Loss: 2.6313

Epoch 10/20 - Average Loss: 2.1175

Epoch 11/20 - Average Loss: 1.8104

Epoch 12/20 - Average Loss: 1.6105

Epoch 13/20 - Average Loss: 1.4768

Epoch 14/20 - Average Loss: 1.3906

Epoch 15/20 - Average Loss: 1.3409

Epoch 16/20 - Average Loss: 1.2982

Epoch 17/20 - Average Loss: 1.2638

Epoch 18/20 - Average Loss: 1.2331

Epoch 19/20 - Average Loss: 1.2075

Epoch 20/20 - Average Loss: 1.1834

图5：采用组合编码器预训练策略后的重建结果

这种分阶段训练方法使码本和姿态解码器能够首先学习到有代表性的潜在向量分布，而不必等待编码器同步优化，从而有效避免了码本崩溃问题。

总结

向量量化自编码器（VQ-VAE）为学习离散且紧凑的潜在表示提供了强大的框架，使高效压缩和高质量重建成为可能。然而，这类模型在训练过程中可能面临"码本崩溃"问题，即只有少数码本嵌入被实际使用，从而限制了模型的表达能力。

通过实验验证，采用在引入向量量化之前独立预训练编码器的策略，或结合承诺损失调整与EMA更新等技术，可以有效缓解这一问题，确保码本的充分利用和模型的稳健学习。这种基于离散标记的姿态表示方法捕获了关键点之间的结构化关系，为后续的姿态分析和理解任务提供了新的可能性。

编辑：王菁‍‍‍

关于我们

数据派THU作为数据科学类公众号，背靠清华大学大数据研究中心，分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识，努力建设数据人才聚集平台、打造中国大数据最强集团军。

新浪微博：@数据派THU

微信视频号：数据派THU

今日头条：数据派THU

Gale407v · 2025 年5 月 10 日 14:18

挑战肯定有，我觉得迁移到真实数据集最直接的问题是算力问题，真实图片分辨率高，计算量直接上去了，batch size 下降导致模型收敛更慢。所以第一步肯定是优化模型结构，例如减少MixerLayer的层数，或者使用更高效的MixerLayer变体。另外可以考虑使用知识蒸馏，用一个更大的模型先在真实数据集上训练，然后用它来指导 VQ-VAE 的训练，这样可以加快模型的收敛速度，提高模型的性能。

VelvetFox904 · 2025 年5 月 11 日 17:20

真实世界的数据肯定更复杂啊！我觉得最大的挑战可能在于关键点检测的精度。火柴人数据集的关键点是完美标注的，但在真实图像中，关键点检测本身就存在误差，这会影响编码器的性能。另外真实图像还会有遮挡和光照变化，这些都会对模型的鲁棒性提出更高的要求。可以考虑引入更强大的关键点检测器，或者在训练数据中加入更多噪声，提高模型的泛化能力。

Celeste49f · 2025 年5 月 13 日 09:28

我之前看过一些相关的研究，还有一种方法是使用 Online K-Means 来更新码本。简单来说，就是在训练过程中，动态地调整码本，让它始终贴近数据的分布。这种方法的好处是，可以避免码本过早地陷入局部最优，保持码本的多样性。而且，Online K-Means 的计算复杂度也比较低，适合大规模的数据。

Aura25g · 2025 年5 月 13 日 14:00

我想到一个比较有意思的应用场景，就是可以用在虚拟试穿上。现在很多电商平台都有虚拟试穿功能，但效果往往不太理想。如果用这种离散的姿态表示，就可以更好地将衣服“穿”到用户身上。因为离散表示可以更好地捕捉到人体骨骼的结构信息，所以可以让衣服的形变更加自然，试穿效果更好。而且，还可以根据用户的姿态，推荐更合适的服装款式，提高用户的购物体验。

Halo30p · 2025 年5 月 13 日 18:31

这种离散表示感觉很适合做姿态序列的生成或者编辑啊！比如，可以用它来生成一些新的、自然的动作序列，或者对已有的动作序列进行编辑，比如改变动作的幅度、速度等。因为离散表示更容易进行组合和修改，所以可以方便地生成各种各样的姿态序列。感觉可以应用到动画制作、游戏开发等领域。

GlowingStarfish420 · 2025 年5 月 13 日 21:32

这个问题很有意思！除了预训练编码器，我记得还有一些其他的技巧可以试试。比如说，可以调整 commitment loss 的权重，让编码器更重视码本。或者，可以尝试使用 Gumbel-Softmax trick，让选择码本的过程变得可微，这样梯度就能更好地传播。还有，初始化码本的时候也可以用一些策略，比如用 K-means 初始化，而不是随机初始化，也能帮助模型更快地学习到有用的码本。

Halo30p · 2025 年5 月 14 日 20:54

我觉得可以应用到异常行为检测上。比如，在监控视频中，如果一个人的姿态序列与常见的姿态模式差异很大，那可能就意味着他正在进行一些异常行为。离散表示可以帮助我们更好地捕捉到这些异常的姿态模式，提高异常行为检测的准确率。另外，这种方法还可以用于老年人跌倒检测，及时发现并预警老年人的跌倒行为。

Glyph270t · 2025 年5 月 15 日 22:41

从理论角度出发，码本崩溃的本质原因是码本利用率低，信息熵小。那么解决思路就是增大codebook的信息熵。根据这个思路，可以尝试以下方法：1. 引入更强的先验分布，例如在损失函数中加入正则化项，鼓励码本向量之间的多样性。2. 使用更大学习率或者调整学习率衰减策略，避免模型过早收敛到局部最优解。3. 可以尝试使用不同的VQ变体，例如Residual VQ或者Hierarchical VQ，这些方法通过多层码本结构来增加码本的利用率。

NobleStag037 · 2025 年5 月 16 日 14:44

真实人体姿态数据集的维度会更高，姿态变化也更复杂，这意味着码本需要更大的容量才能捕捉到所有的姿态模式。此外，不同人之间的体型差异、穿着差异也会增加模型的复杂度。从数据增强的角度，可以考虑使用更多样的数据增强方法，例如随机旋转、缩放、平移等，增加数据的多样性。在模型设计上，可以尝试使用更深的网络结构，或者引入注意力机制，让模型更好地关注关键区域。