AI生成Minecraft(open-oasis) 代码浅读 - 1(Attention and Embedding)

%20%E4%BB%A3%E7%A0%81%E6%B5%85%E8%AF%BB%20-%201(Attention%20and%20E/arch.png)

按照步骤看模块

VIT-VAE

1. Attention

这里用了两种 TemporalAxialAttention (Time dim in sequence of video)

和 SpatialAxialAttention

D → h(heads) * d

The attention mechanism is applied along the temporal axis (T), meaning that the model learns how the information at different time steps relates to each other

两种 rotary_emb 使用方法不同

# Get rotary emb with [t, 1, d(rotary freq dim)]
q = self.rotary_emb.rotate_queries_or_keys(q, self.rotary_emb.freqs)
k = self.rotary_emb.rotate_queries_or_keys(k, self.rotary_emb.freqs)

def rotate_queries_or_keys(self, t, freqs, seq_dim=None, offset=0, scale=None):
    
    seq_dim = default(seq_dim, self.default_seq_dim)
		#self.default_seq_dim = -3 if seq_before_head_dim else -2 对应T
		#q = rearrange(q, "B T H W (h d) -> (B H W) h T d", h=self.heads)
		
    assert not self.use_xpos or exists(scale), "you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings"

    device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim]

    seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)

    seq_freqs = self.forward(seq, freqs, seq_len=seq_len, offset=offset)

    if seq_dim == -3:
        seq_freqs = rearrange(seq_freqs, "n d -> n 1 d")

  return apply_rotary_emb(seq_freqs, t, scale=default(scale, 1.0), seq_dim=seq_dim)

def get_seq_pos(self, seq_len, device, dtype, offset=0):
    return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor

• rotary_emb module is used to embed the queries (q) and keys (k) to incorporate the frequency information, which is helpful in capturing periodic or sequential dependencies.

is_causal == True: Only attend to past time steps

Spatial —> Attention computed on H, W

dependencies between different spatial locations

freqs = self.rotary_emb.get_axial_freqs(H, W)
#calculate pos_emb(axial freq_emb) based on image size

q = rearrange(q, "(B T) h H W d -> (B T) h (H W) d", B=B, T=T, h=self.heads)
#这里(H W) 在-位

q = apply_rotary_emb(freqs, q)
k = apply_rotary_emb(freqs, k)

is_causal == False: Attend to the whole image

def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2):
    dtype = t.dtype

    if t.ndim == 3:
        seq_len = t.shape[seq_dim]
        freqs = freqs[-seq_len:]

    rot_dim = freqs.shape[-1] #rotary_dim
    end_index = start_index + rot_dim
		
		#feature dim 需要足够大
    assert rot_dim <= t.shape[-1], f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"

    # Split t into three parts: left, middle (to be transformed), and right
    # start 默认是0 就把前rotary_dim拿出来做middle 
    t_left = t[..., :start_index]
    t_middle = t[..., start_index:end_index]
    t_right = t[..., end_index:]

    # Apply rotary embeddings without modifying t in place
    t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)
		# middle * freqs.cos() + 1/2 * middle * freqs.sin()
		
		# Rotating: The rotate_half function is applied to t_middle, 
		# which likely splits the tensor along its feature dimension 
		# and rotates one half by 90 degrees 
		# to apply a complementary sine-based transformation.
		
    out = torch.cat((t_left, t_transformed, t_right), dim=-1)

    return out.type(dtype)

推荐: diffusion-forcing 的作者 这里的rotary_emb attention代码来源

https://boyuan.space/blogs/jushenzhineng.html

关键词： 生成式3d