LLM 推理中处理超出上下文长度的外推方法

3. NTK-aware 系列

NTK-aware Scaled RoPE

目前基于 RoPE 的大模型非常多，基于该位置编码研究外推的算法也很多。NTK-aware Scaled RoPE 是一个网友提出来的，思路很简单，就是将 sin(m/base{-2i/d}) 中的 base 变为 sin(m/(base*alpha){-2i/d})。

从苏剑林等人的分析来看，这种方式类似于进制转换，将 b 进制转换为 lambda b 进制，其中 lambda=k^(2/d)。且理论上分析，这种方法可以实现高频外推、低频内插的目的。苏剑林提出：NTK-aware Scaled RoPE'平摊'就不是最优的，应该是低位（比如第 35 维到 64 维）要分摊更多，高位（第 10-35 维）分摊更少，这就导致了混合进制的产生。

代码如下：

import transformers
old_init = transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.__init__

def ntk_scaled_init(self, dim, max_position_embeddings=2048, base=10000, device=None):
    # The method is just these three lines
    max_position_embeddings = 16384
    a = 8  # Alpha value
    base = base * a ** (dim / (dim - 2))  # Base change formula
    old_init(self, dim, max_position_embeddings, base, device)

transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.__init__ = ntk_scaled_init

Dynamically Scaled RoPE

对于动态 NTK，设置了一个 alpha，其中 alpha 的缩放设置为 (alpha * 当前序列长度 / 原始模型上下文长度) - (alpha - 1)。随着序列长度的增加动态缩放超参数。

def _compute_dynamic_ntk_parameters(
    config: Optional[PretrainedConfig] = None,
    device: Optional["torch.device"] = None,
    seq_len: Optional[int] = None,
    **rope_kwargs,
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies with NTK scaling.
    Credits to the Reddit users /u/bloc97 and /u/emozilla
    """
    if config is not None and len(rope_kwargs) > 0:
        raise ValueError(
            "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
            f"`_compute_dynamic_ntk_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
        )
    if len(rope_kwargs) > 0:
        base = rope_kwargs["base"]
        dim = rope_kwargs["dim"]
        max_position_embeddings = rope_kwargs["max_position_embeddings"]
        factor = rope_kwargs["factor"]
    elif config is not None:
        base = config.rope_theta
        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        dim = int(head_dim * partial_rotary_factor)
        max_position_embeddings = config.max_position_embeddings
        factor = config.rope_scaling["factor"]

    attention_factor = 1.0
    seq_len = seq_len if seq_len is not None and seq_len > max_position_embeddings else max_position_embeddings

    # Compute the inverse frequencies core code
    base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2))
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
    return inv_freq, attention_factor

此外，苏剑林还提出了 ReRoPE 和 Leaky ReRoPE，结合外推和内插的方法，先设定一个窗口大小 w，在窗口内使用大小为 1 的位置间隔，在窗口外使用大小为 1/k 的位置间隔。此方式支持任意长度的 content，但会增加计算量。他还提出了 log n 的缩放因子，在测试阶段每个 q 都乘以 log n 以稳定注意力分数的分布。

4. YaRN 方法

YaRN 有两份部分工作。第一部分是对灵活进制的设计，即不同的位置用不同的进制。对于低维度，表示数字高位，外推就是角度不变；对于中间维度，表示数字的中间的位，需要逐步内插，相当于逐步增大'进制'；对于高维度，表示数字的低位，彻底内插，就是把进制变成原来的进制的 40 倍。

计算公式涉及对频率的修正，具体实现较为复杂，核心在于区分 extrapolation 和 interpolation 的频率。

第二份工作是在相关研究观察到，在对 logits 进行 softmax 操作之前引入温度 t 可以统一地影响困惑度。我们将注意力权重的计算修改为 softmax(qt/(tsqrt(d)))，这个 t 就是参数。YaRN 方法在微调和非微调场景中均超过以前所有方法，由于其占用空间较小，YaRN 与修改注意力机制库（如 Flash Attention 2）直接兼容，且在对不到 0.1% 的原始预训练数据进行微调后，YaRN 在上下文窗口扩展中达到了最先进的性能。如果 YaRN 与动态缩放的推理技术相结合而得到的 Dynamic-yarn，其允许在超过 2 倍的上下文窗口扩展，而无需任何微调。

def _compute_yarn_parameters(
    config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies with NTK scaling.
    Please refer to the original paper.
    """
    if len(rope_kwargs) > 0:
        raise ValueError(
            f"Unexpected arguments: `**rope_kwargs` should be unset in `_compute_yarn_parameters`, got {rope_kwargs}"
        )

    base = config.rope_theta
    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    dim = int(head_dim * partial_rotary_factor)
    max_position_embeddings = config.max_position_embeddings
    factor = config.rope_scaling["factor"]

    attention_factor = config.rope_scaling.get("attention_factor")
    if attention_factor is None:
        attention_factor = 0.1 * math.log(factor) + 1.0

    beta_fast = config.rope_scaling.get("beta_fast") or 32
    beta_slow = config.rope_scaling.get("beta_slow") or 1

    def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
        return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))

    def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
        low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
        high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
        return max(low, 0), min(high, dim - 1)

    def linear_ramp_factor(min_val, max_val, dim):
        if min_val == max_val:
            max_val += 0.001
        linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val)
        ramp_func = torch.clamp(linear_func, 0, 1)
        return ramp_func

    pos_freqs = base ** (torch.arange(0, dim, 2).float().to(device) / dim)
    inv_freq_extrapolation = 1.0 / pos_freqs
    inv_freq_interpolation = 1.0 / (factor * pos_freqs)

    low, high = find_correction_range(beta_fast, beta_slow, dim, base, max_position_embeddings)

    inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).float().to(device)
    inv_freq = (
        inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
        + inv_freq_extrapolation * inv_freq_extrapolation_factor
    )

    return inv_freq, attention_factor

总结与选择建议

上述方法各有优劣。ALiBi 实现简单且无需微调即可外推，适合快速部署；位置内插法（PI）在特定场景下有效，但可能损失精度；NTK-aware 及其变体通过调整基频实现了较好的平衡，是目前主流大模型常用的策略；YaRN 则在保持精度的同时提供了极高的灵活性，尤其适合超长上下文任务。在实际工程中，应根据具体的硬件资源、上下文需求及微调成本选择合适的方案。