Transformers documentation

xLSTM

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

该模型于 2024-05-07 在 HF papers 上发布,并于 2025-07-25 贡献至 Hugging Face Transformers。

xLSTM

概述

xLSTM 模型由 Maximilian Beck、Korbinian Pöppel、Markus Spanring、Andreas Auer、Oleksandra Prudnikova、Michael Kopp、Günter Klambauer、Johannes Brandstetter 和 Sepp Hochreiter 在 xLSTM: Extended Long Short-Term Memory 中提出。 xLSTM 通过引入指数门控、矩阵记忆扩展以及可并行化的训练与数据摄入,对原始 LSTM 架构进行了更新,使其能够与 Transformer 模型相竞争。

7B 模型变体由 NXAI 的 xLSTM 团队训练,团队成员包括 Maximilian Beck、Korbinian Pöppel、Phillip Lippe、Richard Kurle、Patrick Blies、Sebastian Böck 和 Sepp Hochreiter。

该论文的摘要如下:

20 世纪 90 年代,恒定误差传送带(constant error carousel)和门控机制被提出,成为长短期记忆网络(LSTM)的核心思想。自那以后,LSTM 经受住了时间的考验,为众多深度学习成功案例做出了贡献,尤其是它们构成了首批大语言模型(LLM)。然而,以可并行化的自注意力机制为核心的 Transformer 技术的出现标志着一个新纪元的开端,在大规模场景下超越了 LSTM。现在我们提出一个简单的问题:当我们将 LSTM 扩展到数十亿参数规模,并借助现代 LLM 的最新技术,同时缓解 LSTM 的已知局限性时,在语言建模方面能走多远?首先,我们引入了带有适当归一化和稳定化技术的指数门控。其次,我们修改了 LSTM 的记忆结构,从而得到:(i) sLSTM,具有标量记忆、标量更新和新的记忆混合机制;(ii) mLSTM,具有矩阵记忆和协方差更新规则,且可完全并行化。将这些 LSTM 扩展集成到残差块骨干网络中,便得到了 xLSTM 块,随后通过残差堆叠形成 xLSTM 架构。指数门控和修改后的记忆结构提升了 xLSTM 的能力,使其在性能和扩展性方面与最先进的 Transformer 和状态空间模型相比均表现良好。

该模型由 NX-AI 贡献。 原始代码可在此处找到。

备注

  • 在批量生成时使用左侧填充。请参阅 填充侧。

xLSTMConfig

class transformers.xLSTMConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonevocab_size: int = 50304hidden_size: int = 4096embedding_dim: int | None = Nonenum_hidden_layers: int = 32num_blocks: int | None = Nonenum_heads: int = 8use_bias: bool = Falsenorm_reduction_force_float32: bool = Truetie_word_embeddings: bool = Falseadd_out_norm: bool = Truenorm_eps: float = 1e-06qk_dim_factor: float = 0.5v_dim_factor: float = 1.0chunkwise_kernel: typing.Literal['chunkwise--native_autograd', 'parallel--native_autograd'] = 'chunkwise--native_autograd'sequence_kernel: typing.Literal['native_sequence__native'] = 'native_sequence__native'step_kernel: typing.Literal['native'] = 'native'mode: typing.Literal['train', 'train_with_padding', 'inference'] = 'inference'chunk_size: int = 64return_last_states: bool = Trueautocast_kernel_dtype: typing.Literal['float32', 'bfloat16', 'float16'] = 'bfloat16'eps: float = 1e-06inference_state_dtype: typing.Literal['float32', 'bfloat16', 'float16'] = 'float32'ffn_proj_factor: float = 2.667ffn_round_up_to_multiple_of: int = 64gate_soft_cap: float = 15.0output_logit_soft_cap: float = 30.0weight_mode: typing.Literal['single', 'fused'] = 'single'use_cache: bool = Truepad_token_id: int | None = 1bos_token_id: int | None = 0eos_token_id: int | list[int] | None = 2max_inference_chunksize: int = 16384 )

参数

  • vocab_size (int,可选,默认为 50304) — 模型的词汇表大小。定义了 input_ids 可以表示的不同 token 的数量。
  • hidden_size (int,可选,默认为 4096) — 隐藏表示的维度。
  • embedding_dim (int, 可选) — 嵌入和隐藏状态的维度。
  • num_hidden_layers (int,可选,默认为 32) — Transformer 解码器中的隐藏层数量。
  • num_blocks (int,可选,optional,默认值为 32) — xLSTM 模型的块数,如果为 None,则使用 num_hidden_layers。
  • num_heads (int,可选,optional,默认值为 8) — xLSTM 层/单元的头数。
  • use_bias (bool,可选,optional,默认为 False) — 是否在 xLSTM 模型中使用偏置。
  • norm_reduction_force_float32 (bool,可选,optional,默认为 True) — 是否强制以 fp32 精度执行 float32 范数归约操作。
  • tie_word_embeddings (bool,可选,默认为 False) — 是否根据模型的 tied_weights_keys 映射来绑定权重嵌入。
  • add_out_norm (bool,可选,optional,默认为 True) — 是否在 LMHead 之前的 blocks 之后添加输出归一化。
  • norm_eps (float,可选,默认为 1e-06) — 层归一化层所使用的 epsilon 值。
  • qk_dim_factor (float,可选,optional,默认值为 0.5) — 查询和键维度的缩放因子。
  • v_dim_factor (float,可选,optional,默认值为 1.0) — 值维度的缩放因子。
  • chunkwise_kernel (ChunkwiseKernelType,可选,可选,默认为 "chunkwise--native_autograd") — 分块处理模式的内核类型。
  • sequence_kernel (SequenceKernelType,可选,optional,默认为 "native_sequence__native") — 序列处理模式的 Kernel 类型。
  • step_kernel (StepKernelType,可选,optional,默认为 "native") — 步骤处理模式的内核类型。
  • mode (BackendModeType,可选,可选,默认为 "inference") — 操作模式(生成需要推理模式)。
  • chunk_size (int,可选,optional,默认值为 64) — 内部块大小。
  • return_last_states (bool,可选,optional,默认为 True) — 是否返回内部的最后状态/缓存。生成时需要设置为 True。
  • autocast_kernel_dtype (DtypeType,可选,optional,默认为 "bfloat16") — 用于 states 的 kernel 数据类型。
  • eps (float,可选,默认为 1e-06) — 层归一化层所使用的 epsilon 值。
  • inference_state_dtype (DtypeType,可选,optional,默认为 "float32") — 推理时状态使用的 Kernel 数据类型。
  • ffn_proj_factor (float,可选,optional,默认值为 2.667) — 门控前馈网络中后上投影的大小因子。
  • ffn_round_up_to_multiple_of (int,可选,optional,默认值为 64) — 后上投影门控前馈网络的尺寸因子舍入值。
  • gate_soft_cap (float,可选,optional,默认值为 15.0) — 门控软上限的缩放系数。
  • output_logit_soft_cap (float,可选,optional,默认值为 30.0) — 输出 logit 软上限的缩放系数。
  • weight_mode(Literal,可选,默认为 "single")—— 是否将并行线性层分离或融合(single)。
  • use_cache (bool,可选,默认为 True) — 模型是否应返回最后一次键/值注意力(并非所有模型都使用)。仅当 config.is_decoder=True 或模型为仅解码器生成式模型时相关。
  • pad_token_id (int,可选,默认为 1) — 词表中用于填充的 token id。
  • bos_token_id (int,可选,默认为 0) — 词汇表中用于表示流起始的 Token id。
  • eos_token_id (Union[int, list[int]],可选,默认为 2) — 词汇表中用于表示流结束(end-of-stream)的 token id。
  • max_inference_chunksize (int,可选,optional,默认值为 16384) — 限制推理时的块大小以节省内存。

这是用于存储 xLSTMModel 配置的类。它根据指定的参数实例化 Xlstm 模型,并定义模型架构。使用默认值实例化配置,将得到与 NX-AI/xLSTM-7b 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

示例:

>>> from transformers import xLSTMConfig, xLSTMModel

>>> # Initializing a xLSTM configuration
>>> configuration = xLSTMConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = xLSTMModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

xLSTMModel

class transformers.xLSTMModel

< >

( config: xLSTMConfig )

参数

  • config (xLSTMConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

未添加任何特定头部的裸 Xlstm 模型,直接输出原始隐藏状态。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.LongTensor] = Nonecache_params: transformers.models.xlstm.modeling_xlstm.xLSTMCache | None = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → xLSTMOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.LongTensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • cache_params (xLSTMCache,可选) — 用于承载 RNN 状态的 xLSTMCache。
  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。

返回值

xLSTMOutput 或 tuple(torch.Tensor)

xLSTMOutput 或 torch.Tensor 元组(如果传入 return_dict=False 或 config.return_dict=False),包含根据配置(xLSTMConfig)和输入而定的各种元素。

xLSTMModel 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • last_hidden_state(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选)— 模型最后一层输出的隐藏状态序列。

  • cache_params (~models.xlstm.modeling_xlstm.xLSTMCache, 可选) — 模型在最后一个时间步的状态。可在前向方法中与下一个 input_ids 一起使用,以避免提供旧的 input_ids。

  • hidden_states (tuple[torch.Tensor],可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层输出的一个张量,以及每个层输出的一个张量),形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

xLSTMLMHeadModel

class transformers.xLSTMForCausalLM

< >

( config: xLSTMConfig )

参数

  • config (xLSTMConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

用于因果语言建模的 Xlstm 模型。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonecache_params: transformers.models.xlstm.modeling_xlstm.xLSTMCache | None = Nonelabels: typing.Optional[torch.LongTensor] = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → xLSTMCausalLMOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • cache_params (xLSTMCache,可选) — 用于承载 RNN 状态的 xLSTMCache。
  • labels (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 用于计算因果语言建模损失的标签。索引应在 [-100, 0, ..., config.vocab_size] (参见 input_ids 文档字符串)。索引设置为 -100 的 token 将被忽略(掩码),损失仅针对标签在 [0, ..., config.vocab_size].

    标签在内部会进行移位,因此使用labels[..., 1:]来预测logits[..., :-1, :]。 请传入与input_ids匹配的未移位标签(即labels = input_ids)。在传入之前,请勿自行对标签进行移位。

    如果标签已经移位(例如在使用序列并行时),请将其作为 shift_labels 关键字参数传入,系统将直接使用该值。为了计算损失,仍然必须传入 labels。

  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。

返回值

xLSTMCausalLMOutput 或 tuple(torch.Tensor)

xLSTMCausalLMOutput 或 torch.Tensor 元组(如果传入 return_dict=False 或 config.return_dict=False),包含根据配置(xLSTMConfig)和输入而定的各种元素。

xLSTMForCausalLM 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • loss (torch.Tensor,形状为 (1,),可选,当提供 labels 时返回) — 语言建模损失(用于下一个 token 预测)。

  • logits(形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.Tensor)— 语言建模头的预测分数(SoftMax 之前每个词汇表 token 的分数)。

  • cache_params (xLSTMCache,可选,承载 RNN 状态) — 模型在最后一个时间步的状态。可在前向方法中与下一个 input_ids 一起使用,从而避免提供旧的 input_ids。

  • hidden_states (tuple[torch.Tensor],可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层输出的一个张量,以及每个层输出的一个张量),形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

示例:

Update on GitHub