传送门

模型解读

我们知道Transformer的复杂度是平方的, 就很丑; 而RNN的复杂度虽然是线性的, 但是因为隐状态的大小是固定受限的, 长序列处理一直是弱项. 用可训练模型作为隐状态对RNN进行魔改就有了TTT.

连连看

Quick Start是个好东西

点我查看完整代码

from transformers import AutoTokenizer
from ttt import TTTForCausalLM, TTTConfig, TTT_STANDARD_CONFIGS

# Initializing a TTT ttt-1b style configuration
# configuration = TTTConfig(**TTT_STANDARD_CONFIGS['1b']) is equivalent to the following
configuration = TTTConfig()

# Initializing a model from the ttt-1b style configuration
model = TTTForCausalLM(configuration)
model.eval()

# Accessing the model configuration
configuration = model.config

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b-hf')

# Prefill
input_ids = tokenizer("Greeting from TTT!", return_tensors="pt").input_ids
logits = model(input_ids=input_ids)
print(logits)

# Decoding
out_ids = model.generate(input_ids=input_ids, max_length=50)
out_str = tokenizer.batch_decode(out_ids, skip_special_tokens=True)
print(out_str)

config 先跳过, 顺着 TTTForCausalLM 一路找下去, 先看 __init__()

def __init__(self, config):
    super().__init__(config)
    self.model = TTTModel(config)
    self.vocab_size = config.vocab_size
    self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    # Initialize weights and apply final processing
    self.post_init()

很明显, 重要的是TTTModel, 同样是先看 __init__()

def __init__(self, config: TTTConfig):
    super().__init__(config)
    self.padding_idx = config.pad_token_id
    self.vocab_size = config.vocab_size

    self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
    self.layers = nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
    self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
    self.gradient_checkpointing = False

    # Initialize weights and apply final processing
    self.post_init()

从 Block 接着DFS, 可以发现有两个分支: TTTLinear 和 TTTMLP, 它们都继承了 TTTBase, 先看最重要的forward():

点我查看完整代码

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    cache_params: Optional[TTTCache] = None,
):
    B, L = hidden_states.shape[:2]
    reminder_len = L % self.mini_batch_size
    num_mini_batch = L // self.mini_batch_size
    last_mini_batch_params_dict = None

    XQ, XK, XV = self.get_qkv_projections(hidden_states, cache_params=cache_params) 
    #// X_{Q,K,V} = \theta_{Q,K,V} * X_{t}

    # [B, L, C] -> [B, L, num_heads, head_dim] -> [B, num_heads, L, head_dim]
    XQ = XQ.reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)
    XK = XK.reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)
    XV = XV.reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)

    cos, sin = self.rotary_emb(XV, position_ids % self.mini_batch_size)

    # permute_qk and undo_permute_qk is just for aligning pytorch with jax pre-training
    XQ, XK = permute_qk(XQ, XK)
    XQ, XK = apply_rotary_pos_emb(XQ, XK, cos, sin) #// 旋转位置编码
    XQ, XK = undo_permute_qk(XQ, XK)

    output_hidden_states = []
    # when input sequence length is not a multiple of mini_batch_size
    # we need to compute them seperately, when computing the reminder,
    # we will need the last_mini_batch_params_dict to continue TTT learning
    if num_mini_batch > 0:
        inputs = {
            "XQ": XQ[:, :, : num_mini_batch * self.mini_batch_size],
            "XK": XK[:, :, : num_mini_batch * self.mini_batch_size],
            "XV": XV[:, :, : num_mini_batch * self.mini_batch_size],
            "X": hidden_states[:, : num_mini_batch * self.mini_batch_size],
        }
        output_mod, last_mini_batch_params_dict = self.ttt(
            self.get_ttt_inputs(inputs, self.mini_batch_size, cache_params),
            mini_batch_size=self.mini_batch_size,
            last_mini_batch_params_dict=last_mini_batch_params_dict,
            cache_params=cache_params,
        )
        output_hidden_states.append(output_mod)
    if reminder_len > 0:
        inputs = {
            "XQ": XQ[:, :, -reminder_len:],
            "XK": XK[:, :, -reminder_len:],
            "XV": XV[:, :, -reminder_len:],
            "X": hidden_states[:, -reminder_len:],
        }
        output_reminder, _ = self.ttt(
            self.get_ttt_inputs(inputs, reminder_len, cache_params),
            mini_batch_size=reminder_len,
            last_mini_batch_params_dict=last_mini_batch_params_dict,
            cache_params=cache_params,
        )
        output_hidden_states.append(output_reminder)

    output_hidden_states = torch.cat(output_hidden_states, dim=1)
    output_hidden_states = self.post_norm(output_hidden_states) #// LayerNorm
    if self.use_gate: #// Mamba Only
        output_hidden_states = self.apply_gate(hidden_states, output_hidden_states)
    output_hidden_states = self.o_proj(output_hidden_states)

    return output_hidden_states

实现细节先跳过, 跳到 ttt 的具体实现, 先看 TTTLinear 里的实现:

点我查看完整代码

class TTTLinear(TTTBase):
    def __init__(self, config: TTTConfig, layer_idx: Optional[int] = None):
        super().__init__(config, layer_idx)
        # TTT model initialization for TTT-Linear
        self.W1 = nn.Parameter(torch.normal(0, 0.02, size=(self.num_heads, self.head_dim, self.head_dim)))
        self.b1 = nn.Parameter(torch.zeros(self.num_heads, 1, self.head_dim))

    def ttt(
        self,
        inputs,
        mini_batch_size,
        last_mini_batch_params_dict,
        cache_params: Optional[TTTCache] = None,
    ):
        if mini_batch_size is None:
            mini_batch_size = self.mini_batch_size

        # in this case, we are decoding
        if last_mini_batch_params_dict is None and cache_params is not None:
            last_mini_batch_params_dict = cache_params.ttt_params_to_dict(self.layer_idx)

        # [B, num_heads, num_mini_batch, mini_batch_size, head_dim]
        B = inputs["XV"].shape[0]
        num_mini_batch = inputs["XV"].shape[2]
        L = inputs["XV"].shape[2] * inputs["XV"].shape[3]
        device = inputs["XV"].device
        dtype = inputs["XV"].dtype

        # NOTE:
        # for prefilling, we will always use dual form for faster computation
        # we need to use primal form if mini_batch_size is not a multiple of self.mini_batch_size
        # since we need store the gradient for the next mini-batch computation
        use_dual_form = cache_params is None or mini_batch_size % self.mini_batch_size == 0

        def compute_mini_batch(params_dict, inputs): #// TL;DR

        if last_mini_batch_params_dict is not None:
            init_params_dict = last_mini_batch_params_dict
        else:
            init_params_dict = {
                "W1_states": torch.tile(self.W1.unsqueeze(0), dims=(B, 1, 1, 1)),
                "b1_states": torch.tile(self.b1.unsqueeze(0), dims=(B, 1, 1, 1)),
            }
            init_params_dict.update(W1_grad=torch.zeros_like(init_params_dict["W1_states"]))
            init_params_dict.update(b1_grad=torch.zeros_like(init_params_dict["b1_states"]))

        # [B,num_heads, num_mini_batch, mini_batch_size, f] -> [num_mini_batch, B, num_heads, mini_batch_size, f]
        inputs = tree_map(lambda x: x.permute(2, 0, 1, 3, 4), inputs)

        # allocate output tensor
        XQW_batch = torch.empty(
            (num_mini_batch, B, self.num_heads, mini_batch_size, self.head_dim),
            device=device,
            dtype=dtype,
        )
        # XQW_batch: [num_mini_batch, B, num_heads, mini_batch_size, head_dim]
        batch_params_dict, XQW_batch = scan( #// foreach (var x in inputs) { compute_mini_batch(x) }
            compute_mini_batch,
            init_params_dict,
            inputs,
            XQW_batch,
            self.config.scan_checkpoint_group_size if self.training else 0,
        )

        # [B, num_heads, L, C]
        if cache_params is not None:
            cache_params.update(batch_params_dict, self.layer_idx, L)

        # [num_mini_batch, B, num_heads, mini_batch_size, head_dim] -> [B, num_mini_batch, mini_batch_size, num_heads, head_dim]
        XQW_batch = XQW_batch.permute(1, 0, 3, 2, 4)
        # [B, L, C]
        XQW_batch = XQW_batch.reshape(B, L, self.width)
        return XQW_batch, batch_params_dict

继续研究 compute_mini_batch() 是干什么的

点我查看完整代码

def compute_mini_batch(params_dict, inputs):
    # [B, nh, f, f], nh=num_heads, f=head_dim
    W1_init = params_dict["W1_states"]
    # [B, nh, 1, f]
    b1_init = params_dict["b1_states"]

    # [B,nh,K,f], K=mini_batch_size
    XQ_mini_batch = inputs["XQ"]
    XV_mini_batch = inputs["XV"]
    XK_mini_batch = inputs["XK"]
    # [B, nh, K, 1]
    eta_mini_batch = inputs["eta"]
    token_eta_mini_batch = inputs["token_eta"]
    ttt_lr_eta_mini_batch = inputs["ttt_lr_eta"]

    X1 = XK_mini_batch
    # [B,nh,K,f] @ [B,nh,f,f] -> [B,nh,K,f]
    Z1 = X1 @ W1_init + b1_init
    reconstruction_target = XV_mini_batch - XK_mini_batch

    ln_weight = self.ttt_norm_weight.reshape(self.num_heads, 1, self.head_dim)
    ln_bias = self.ttt_norm_bias.reshape(self.num_heads, 1, self.head_dim)
    # [B,nh,K,f]
    grad_l_wrt_Z1 = ln_fused_l2_bwd(Z1, reconstruction_target, ln_weight, ln_bias) 
    #// Batch backward for LayerNorm fused with L2 loss.

    if use_dual_form:
        # [B,nh,K,K]
        Attn1 = torch.tril(XQ_mini_batch @ X1.transpose(-2, -1)) #// 经典的 Q*K^{T}
        # [B,nh,1,f] - [B,nh,K,K] @ [B,nh,K,f] -> [B,nh,K,f]
        b1_bar = b1_init - torch.tril(eta_mini_batch) @ grad_l_wrt_Z1
        # [B,nh,K,f] @ [B,nh,f,f] - ([B,nh,K,1] * [B,nh,K,K]) @ [B,nh,K,f] + [B,nh,K,f]
        Z1_bar = XQ_mini_batch @ W1_init - (eta_mini_batch * Attn1) @ grad_l_wrt_Z1 + b1_bar 
        # Z = X * W - 2 \eta * \delta + B

        #// Gradient Descent and Update Gradient
        last_eta_mini_batch = eta_mini_batch[:, :, -1, :, None]
        # [B,nh,f,f] - [B,nh,f,K] @ [B,nh,K,f]
        W1_last = W1_init - (last_eta_mini_batch * X1).transpose(-1, -2) @ grad_l_wrt_Z1
        # [B,nh,1,f]
        b1_last = b1_init - torch.sum(last_eta_mini_batch * grad_l_wrt_Z1, dim=-2, keepdim=True)
        grad_W1_last = torch.zeros_like(W1_last)
        grad_b1_last = torch.zeros_like(b1_last)
    else:
        ttt_lr_eta_mini_batch = torch.broadcast_to(
            ttt_lr_eta_mini_batch,
            (
                *ttt_lr_eta_mini_batch.shape[:2],
                mini_batch_size,
                mini_batch_size,
            ),
        )

        # [B, nh, K, f, f]
        grad_W1 = torch.einsum("bhki,bhkj->bhkij", X1, grad_l_wrt_Z1)
        grad_W1 = torch.einsum("bhnk,bhkij->bhnij", torch.tril(ttt_lr_eta_mini_batch), grad_W1)
        grad_W1 = grad_W1 + params_dict["W1_grad"].unsqueeze(2)
        # [B, nh, K, f]
        grad_b1 = torch.einsum("bhnk,bhki->bhni", torch.tril(ttt_lr_eta_mini_batch), grad_l_wrt_Z1)
        grad_b1 = grad_b1 + params_dict["b1_grad"]

        W1_bar = W1_init.unsqueeze(2) - grad_W1 * token_eta_mini_batch.unsqueeze(-1)
        b1_bar = b1_init - grad_b1 * token_eta_mini_batch

        # [B, nh, K, 1, f] @ [B, nh, K, f, f]
        Z1_bar = (XQ_mini_batch.unsqueeze(3) @ W1_bar).squeeze(3) + b1_bar
        #// Z = X * W + B

        W1_last = W1_bar[:, :, -1]
        b1_last = b1_bar[:, :, -1:]
        grad_W1_last = grad_W1[:, :, -1]
        grad_b1_last = grad_b1[:, :, -1:]

    Z1_bar = ln_fwd(Z1_bar, ln_weight, ln_bias) 
    #// Batch forward for LayerNorm.

    XQW_mini_batch = XQ_mini_batch + Z1_bar

    last_param_dict = {
        "W1_states": W1_last,
        "b1_states": b1_last,
        "W1_grad": grad_W1_last,
        "b1_grad": grad_b1_last,
    }
    return last_param_dict, XQW_mini_batch

TTTMLP 实际上就是两层 TTTLinear 叠叠乐, 这里就不做了

钾肥喵的窝

代码连连看——TTT

传送门

模型解读

连连看