請問有實現的程式碼嗎?

[ILoveSorasakiHina]

親愛的研究團隊 您好:

非常感謝您的研究，拜讀完您的論文後，我有一個小小的疑問:

於論文中，我有看到關於Block Attention Residuals的虛擬碼，但好像沒Full Attention Residuals的虛擬碼

雖然理論上Full Attention Residuals照著公式實現應該就好，但我想確認一下是否實現正確

所以想請問能提供您們實驗的code嗎?

再次感謝您的貢獻，期待您的回覆

祝 研究順利

[zhuygln]

Look forward to official implementation too!

[chilljudaoren]

親愛的研究團隊 您好:

非常感謝您的研究，拜讀完您的論文後，我有一個小小的疑問:

於論文中，我有看到關於Block Attention Residuals的虛擬碼，但好像沒Full Attention Residuals的虛擬碼

雖然理論上Full Attention Residuals照著公式實現應該就好，但我想確認一下是否實現正確

所以想請問能提供您們實驗的code嗎?

再次感謝您的貢獻，期待您的回覆

祝 研究順利

可以看看你实现的吗？

[ILoveSorasakiHina]

"""
Attention Residuals (AttnRes) — Full 版本實現
適用於 ViT / Swin Transformer 等同尺度層的殘差替換
 
論文: Attention Residuals (Kimi Team)
參考: https://github.com/MoonshotAI/Attention-Residuals
 
核心思想:
  標準殘差: h_l = h_{l-1} + f_{l-1}(h_{l-1})  (固定均勻累加)
  AttnRes:  h_l = Σ α_{i→l} · v_i              (softmax 加權選擇性聚合)
 
其中:
  - Q = w_l (每層一個可學習的 d 維向量, 不依賴輸入)
  - K = RMSNorm(v_i) (各層輸出經 RMSNorm)
  - V = v_i (各層原始輸出)
  - α_{i→l} = softmax(w_l^T · RMSNorm(v_i))  沿深度維度歸一化
"""
 
import torch
import torch.nn as nn
from typing import List, Optional
 
 
# ============================================================
# 1. RMSNorm (論文中對 Key 使用的歸一化)
# ============================================================
class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization"""
 
    def __init__(self, dim: int, eps: float = 1e-8):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: 任意形狀, 最後一維為 d
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight
 
 
# ============================================================
# 2. Full Attention Residuals 核心運算
# ============================================================
class FullAttnRes(nn.Module):
    """
    Full Attention Residuals 運算模組
 
    每一層呼叫此模組, 對 history 中所有先前層的輸出做 softmax 注意力加權,
    產生該層的輸入 h_l。
 
    參數:
        dim: 隱藏維度 d
    """
 
    def __init__(self, dim: int):
        super().__init__()
        # 每層一個可學習的偽查詢向量 w_l ∈ R^d
        # 論文強調: 必須初始化為零, 使初始 α 均勻 → 等價於等權平均 → 避免訓練初期不穩定
        self.pseudo_query = nn.Parameter(torch.zeros(dim))
        # 對 Key 做 RMSNorm, 防止輸出幅值大的層主導注意力權重
        self.key_norm = RMSNorm(dim)
 
    def forward(self, history: List[torch.Tensor]) -> torch.Tensor:
        """
        Args:
            history: list of tensors, 每個形狀為 [B, T, d]
                     history[0] = token embedding (h_1)
                     history[i] = f_i(h_i), 第 i 層的輸出 (i >= 1)
 
        Returns:
            h_l: [B, T, d], 當前層的輸入 (取代傳統殘差的 h_{l-1} + f(h_{l-1}))
        """
        # 堆疊所有先前層的輸出: [n, B, T, d], n = len(history)
        V = torch.stack(history, dim=0)
 
        # Key = RMSNorm(V): [n, B, T, d]
        K = self.key_norm(V)
 
        # 計算注意力分數: w_l^T · K_i (對每個 token 獨立)
        # pseudo_query: [d]  ×  K: [n, B, T, d]  →  logits: [n, B, T]
        logits = torch.einsum('d, n b t d -> n b t', self.pseudo_query, K)
 
        # 沿深度維度 (dim=0) 做 softmax 歸一化
        alpha = torch.softmax(logits, dim=0)  # [n, B, T]
 
        # 加權求和: α_i · v_i
        # alpha: [n, B, T]  ×  V: [n, B, T, d]  →  h_l: [B, T, d]
        h_l = torch.einsum('n b t, n b t d -> b t d', alpha, V)
 
        return h_l



# ============================================================
# 3. 使用 AttnRes 的 Transformer Layer (適用於 ViT)
# ============================================================
class AttnResTransformerLayer(nn.Module):
    """
    使用 Full AttnRes 替換標準殘差連接的 Transformer Layer
 
    論文中, 每個 self-attention 和 MLP 被視為獨立的「層」,
    各自擁有獨立的 AttnRes 模組。
 
    標準 ViT Layer:
        x = x + attention(norm(x))
        x = x + mlp(norm(x))
 
    AttnRes ViT Layer:
        h = attn_res_attn(history)          # 用 AttnRes 取代 residual 得到 attention 的輸入
        attn_out = attention(norm(h))       # attention 計算
        history.append(attn_out)            # 記錄 attention 的輸出到歷史
 
        h = attn_res_mlp(history)           # 用 AttnRes 取代 residual 得到 MLP 的輸入
        mlp_out = mlp(norm(h))              # MLP 計算
        history.append(mlp_out)             # 記錄 MLP 的輸出到歷史
    """
 
    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        drop: float = 0.0,
    ):
        super().__init__()
 
        # === 標準 Transformer 組件 ===
        self.norm_attn = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(
            embed_dim=dim, num_heads=num_heads,
            dropout=drop, batch_first=True
        )
 
        self.norm_mlp = nn.LayerNorm(dim)
        mlp_hidden = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_hidden),
            nn.GELU(),
            nn.Dropout(drop),
            nn.Linear(mlp_hidden, dim),
            nn.Dropout(drop),
        )
 
        # === AttnRes 模組: attention 和 MLP 各一個 ===
        self.attn_res_attn = FullAttnRes(dim)
        self.attn_res_mlp = FullAttnRes(dim)
 
    def forward(self, history: List[torch.Tensor]) -> List[torch.Tensor]:
        """
        Args:
            history: 到目前為止所有層輸出的 list
                     初始時 history = [token_embedding]
 
        Returns:
            history: 更新後的 history (新增了本層 attention 和 MLP 的輸出)
        """
        # --- Pre-Attention ---
        # 用 AttnRes 聚合歷史, 取代標準殘差的 h_{l-1}
        h = self.attn_res_attn(history)
 
        # Self-Attention (PreNorm)
        h_normed = self.norm_attn(h)
        attn_out, _ = self.attn(h_normed, h_normed, h_normed)
 
        # 記錄 attention 輸出 (這是 f_l(h_l), 不包含殘差)
        history.append(attn_out)
 
        # --- Pre-MLP ---
        # 再次用 AttnRes 聚合歷史 (現在包含了剛才的 attn_out)
        h = self.attn_res_mlp(history)
 
        # MLP (PreNorm)
        mlp_out = self.mlp(self.norm_mlp(h))
 
        # 記錄 MLP 輸出
        history.append(mlp_out)
 
        return history

我叫claude實現並稍微修改的版本

但我只有實現full版本

而且我沒有實際測試過

主要是因為我是做語意分割問題

我想套用在swin-transformer這類模型時發現圖片尺寸跟維度會隨層數改變

這樣算attetion residual感覺會出問題

所以沒做後續實驗

[Smu-Tan]

親愛的研究團隊 您好:

非常感謝您的研究，拜讀完您的論文後，我有一個小小的疑問:

於論文中，我有看到關於Block Attention Residuals的虛擬碼，但好像沒Full Attention Residuals的虛擬碼

雖然理論上Full Attention Residuals照著公式實現應該就好，但我想確認一下是否實現正確

所以想請問能提供您們實驗的code嗎?

再次感謝您的貢獻，期待您的回覆

祝 研究順利

I've implemented one based on Megatron-LM, with llama3 dense arch pretrained 1k steps, see https://github.com/Smu-Tan/Megatron-LM