如何优化建模496×2多维输入以最小化KL散度？

问题翻译 #

我的输入维度为496×2，代表496个不同特征及其对应的两个参数值；输出为这496个特征的分布，需满足总和为1.0。我尝试用多层感知器（MLP）建模，当前可行模型如下：

model.add(Dense(16, activation='relu', kernel_initializer='he_normal', input_shape=(496,2)))
model.add(Flatten())
model.add(Dense(128, activation='relu', kernel_initializer='he_normal'))
model.add(Dense(128, activation='relu', kernel_initializer='he_normal'))
model.add(Dense(496, activation='softmax'))
model.compile(optimizer='adam', loss='kl_divergence')

选择KL散度作为损失函数，是为了最小化两个含496个特征的分布之间的差异。该模型虽能运行，但存在不足：未考虑输入与输出的位置对应关系，且Flatten层会混合每个特征的两个参数。请问是否存在更优的建模方案？

解决方案 #

你的核心痛点是输入的每个2维特征与输出对应位置的概率值是一一绑定的，原模型的Flatten层破坏了这种位置关联，导致模型无法学习到特征与输出的对应逻辑。以下是几种针对性的优化方案：

方案1：共享权重的单特征编码器（最直接匹配需求）

针对每个2维特征单独做编码，严格保留输入输出的位置对应，同时用共享参数避免过拟合：

from tensorflow.keras import layers, Model

# 输入层：(批量大小, 496, 2)
inputs = layers.Input(shape=(496, 2))
# 对每个2维特征用共享的小型网络编码，输出(批量大小, 496, 1)
per_feature = layers.Dense(8, activation='relu', kernel_initializer='he_normal')(inputs)
per_feature = layers.Dense(1, activation='linear', kernel_initializer='he_normal')(per_feature)
# 展平后直接做全局softmax得到概率分布
flattened = layers.Flatten()(per_feature)
outputs = layers.Dense(496, activation='softmax')(flattened)

model = Model(inputs=inputs, outputs=outputs)
model.compile(optimizer='adam', loss='kl_divergence')

• 优势：完全保留输入输出的位置对应，每个特征的编码逻辑一致，参数总量远小于原模型，不易过拟合。

方案2：1D卷积特征提取（支持特征间关联）

如果496个特征存在序列/局部关联（比如按某种逻辑排序），可以用1D卷积捕捉相邻特征的交互，同时保留位置信息：

from tensorflow.keras import layers, Model

inputs = layers.Input(shape=(496, 2))
# 1D卷积提取特征，kernel_size=1仅关注当前特征，>1则捕捉相邻特征交互
x = layers.Conv1D(16, kernel_size=1, activation='relu', kernel_initializer='he_normal')(inputs)
x = layers.Conv1D(8, kernel_size=1, activation='relu', kernel_initializer='he_normal')(x)
# 映射到单维特征后做softmax
x = layers.Conv1D(1, kernel_size=1, activation='linear')(x)
x = layers.Flatten()(x)
outputs = layers.Dense(496, activation='softmax')(x)

model = Model(inputs=inputs, outputs=outputs)
model.compile(optimizer='adam', loss='kl_divergence')

• 优势：灵活支持特征间的局部交互，若用kernel_size=1则效果与方案1一致，但扩展性更强。

补充说明

• 原模型的问题本质：Flatten层将496×2的结构拆成992维的无序向量，后续全连接层会混合所有特征的参数，导致模型无法建立“第i个输入特征对应第i个输出概率”的关联。
• KL散度的选择是合理的，完全匹配你“最小化两个分布差异”的目标。

内容的提问来源于stack exchange，提问作者ymc