token剪枝实现细节方面的疑问

[SomeOtherScenery]

作者好，感谢你们对这项工作的开源，我最近也在研究VLA剪枝方法，lightVLA这个方法给了我很多启发。我在阅读代码的过程中产生了一些疑问，希望可以得到解答。

• 在TokenPruner的代码实现中，gumbel-softmax不是使用gumbel噪声，而是使用均匀分布噪声实现的，这有什么特殊原因吗？
  

  LightVLA/prismatic/extern/hf/modeling_prismatic.py

  Lines 92 to 94 in a4680fd

  	def score_to_indices(self, score, patches):
  	if self.noise_scale is not None:
  	score = score + torch.rand_like(score) * self.noise_scale

• VLA-cache，ADP等剪枝方法中，为了保证训练时每个batch间token长度一致会事先固定剪枝比例，但在lightVLA训练使用的score_to_indices函数中，似乎是通过让部分patch重复占位来保证token总长度在剪枝前后不发生变化。这意味着lightVLA只会降低推理显存，而不会降低训练时的显存占用吗？

  LightVLA/prismatic/extern/hf/modeling_prismatic.py

  Lines 95 to 98 in a4680fd

  	hard_score = F.one_hot(score.argmax(dim=-1), num_classes=self.num_patches)
  	soft_score = torch.softmax(score, dim=-1)
  	score = hard_score + soft_score - soft_score.detach()
  	return score.argmax(dim=-1), score @ patches

[Promisery]

• 在训练前期，TokenPruner不太能准确地选出有用的token，所以我们在score上加了一个采样噪声来帮助采样。在训练后期，TokenPruner已经可以准确选出有用的token了，就不需要采样噪声了，所以我们给这个噪声加了一个系数，让噪声大小随着训练进行逐渐减小。torch.nn.functional.gumbel_softmax的噪声大小是固定的，在训练后期很干扰token选择的效果，所以我们没用gumbel噪声。

LightVLA/vla-scripts/finetune.py

Line 1032 in a4680fd

vla.module.language_model.model.pruner.set_noise_scale(1 - log_step / cfg.max_steps)

• 对，LightVLA在训练的时候token长度不变，这个原因也很简单，LightVLA每次保留的token数量是不一样的，如果我们在训练的时候也做剪枝，同一个batch里的数据长度不一样，要padding才能重新拼成一个batch。而且这样做的话在训练的过程中显存占用不稳定。所以我们在训练的时候是保留原长度，只在推理的时候剪枝。