大模型中top-p&top-k&temperature如何共同使用——Gemma为例子

参考：
大模型文本生成——解码策略（Top-k & Top-p & Temperature）
大模型源码理解-以Gemma为例子

摘要

之前系统学习了大模型的解码方式，Top-p, Top-k, Beam-search, Greedy, temperature等等，具体使用的时候，也清楚采用这些方式混合使用，但是具体怎么混合，有些模糊。看了一篇相关文章大模型文本生成——解码策略（Top-k & Top-p & Temperature），如下图所示，解决了我一些理解方面的问题，但是感觉还有有些模糊，仔细研究了一下Gemma，记录一下。

结论

先说一下结论，Gemma是怎么解码的，主要研究top-p & top-k &temperature是如何使用的，temperature>top-p>top-k（其实top-p和top-k可以算作并行，同时使用）。与上图的结论略有区别，可能不同得模型策略方便略有区别吧。

代码分析

通过代码进行分析，是如何实现三种策略混合使用。
我觉得，top-p和top-k，一起使用，保留同时满足top-p&top-k的概率值。

class Sampler(nn.Module):

    def __init__(self, vocab_size: int):
        super().__init__()
        self.vocab_size = vocab_size

    @torch.no_grad()
    def forward(
        self,
        embedding: torch.Tensor,
        hidden_states: torch.Tensor,
        output_positions: torch.Tensor,
        temperatures: Union[torch.Tensor, None],
        top_ps: torch.Tensor,
        top_ks: torch.Tensor,
        embedding_bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # Select the last element for each sequence.
        # (batch_size, input_len, hidden_size) -> (batch_size, hidden_size)
        hidden_states = hidden_states.index_select(#
            1, output_positions).squeeze(dim=1)
        logits = torch.matmul(hidden_states, embedding.t())#计算不同得Token得分情况
        if embedding_bias is not None:#是否增加bisa，这个无关紧要
            logits += embedding_bias

        if temperatures is None:#temperature为空则设置贪婪匹配则选择最大得概率，GPT中好像是0进行贪婪匹配
            return torch.argmax(logits, dim=-1).squeeze(dim=-1)

        # Apply temperature scaling.
        logits.div_(temperatures.unsqueeze(dim=1))#预测得结果去除以temperature，修改分布

        # Calculate probabilities with softmax.
        probs = torch.softmax(logits, dim=-1, dtype=torch.float)#进行softmax归一化
        probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)#从大到小进行排序

        # Apply top-p, top-k.
		#这里进行top-p
        probs_sum = torch.cumsum(probs_sort, dim=-1)
        top_ps_mask = (probs_sum - probs_sort) > top_ps.unsqueeze(dim=1)
        probs_sort = torch.where(top_ps_mask, 0, probs_sort)
		
		#这里进行top-k
        top_ks_mask = torch.arange(probs_idx.shape[-1],
                                   device=probs_idx.device)
        top_ks_mask = top_ks_mask.expand(probs_idx.shape[0], -1)
        top_ks_mask = top_ks_mask >= top_ks.unsqueeze(dim=1)
      	#这里top-p和top-k同时起作用，同时满足top-p和top-k得结果才有概率值，否则就为0
      	#因为top_ks_mask判断大于top_ks得为True， 大于得不是我们得范围，所以top_ks_mask就补0, 不为True，则补top-p得结果
        probs_sort = torch.where(top_ks_mask, 0, probs_sort)

        # Re-normalization.
        probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
        probs = torch.gather(probs_sort,
                             dim=-1,
                             index=torch.argsort(probs_idx, dim=-1))

        next_token_ids = torch.multinomial(probs,
                                           num_samples=1,
                                           replacement=True).squeeze(dim=-1)
        return next_token_ids