NaViT：解决图像大小问题的革命性大模型架构

NaViT是一种创新的可变分辨率视觉Transformer架构，通过引入"Patch n’ Pack"技术，将不同图像的多个补丁打包到单个序列中，使模型能处理不同分辨率和长宽比的图像，无需调整大小或填充。核心创新包括掩码自注意力防止图像间标记相互关注、分解位置嵌入支持任意分辨率、以及连续标记丢弃和分辨率采样等训练技术。这种方法解决了传统ViT处理不同尺寸图像时存在的性能损失和计算资源浪费问题。

背景
=====

面对不同图像大小，传统做法（ViT）是调整图像大小或填充以适应固定输入尺寸，但这种做法有严重缺陷：

调整大小会损害性能（扭曲原始信息）
填充是计算资源的浪费

图3清晰展示了这个问题：ImageNet中85.9%、LVIS中92.2%和WebLI中57.3%的图像都不是方形的。

NaViT的核心：Patch n’ Pack
NaViT的创新点在于将自然语言处理中使用的序列打包(seq packing)技术引入视觉Transformer。它允许模型处理不同分辨率和长宽比的原始图像，无需调整大小或填充。

关键思想：将不同图像的多个补丁打包到单个序列中（即"Patch n’ Pack"），让Transformer可以处理可变长度的输入序列。

方法详解
=======

2.1 架构变化

2.1.1 掩码自注意力和掩码池化

传统Transformer中，序列中所有标记相互关注。但在打包多张图像的序列中，我们需要防止不同图像的标记相互关注。
解决方案：引入额外的自注意力掩码，确保每个标记只能关注同属一张图像的其他标记（上图清晰展示了这一机制）
同样，掩码池化被引入，确保池化操作只在每个示例的标记内部进行，从而为序列中的每个示例产生一个向量表示
填充标记在注意力和池化过程中被屏蔽，不会影响真实数据

1. Masked Self Attention（掩码自注意力）

工作原理

传统Transformer自注意力：计算时，序列中所有位置都可以关注所有其他位置
NaViT的掩码自注意力：在计算前应用块对角掩码矩阵，确保：

每个图像内部的token可以自由相互关注
不同图像的token完全隔离（注意力权重=0）
填充token被完全屏蔽

掩码矩阵M的结构如下（假设序列包含3张图像）：

M = [ [1,1,...,1, 0,0,...,0, 0,0,...,0], # Image1关注Image1 [1,1,...,1, 0,0,...,0, 0,0,...,0], ... [0,0,...,0, 1,1,...,1, 0,0,...,0], # Image2关注Image2 [0,0,...,0, 1,1,...,1, 0,0,...,0], ... [0,0,...,0, 0,0,...,0, 1,1,...,1], # Image3关注Image3 [0,0,...,0, 0,0,...,0, 1,1,...,1]]

对于长度为N的序列，掩码矩阵是N×N的块对角矩阵
每个块的大小等于对应图像的token数量
填充token所在行/列全为0

实现方式：通过创建一个二进制掩码矩阵，将不同图像之间的注意力权重设为负无穷

# 伪代码示例：构建NaViT的注意力掩码def create_attention_mask(packed_sequence, example_boundaries): """ packed_sequence: 打包后的序列 [seq_length, embedding_dim] example_boundaries: 每个图像在序列中的起始和结束位置 """ seq_length = packed_sequence.shape[0] mask = torch.full((seq_length, seq_length), float('-inf')) # 为每个图像创建一个块 for start, end in example_boundaries: # 许可图像内部的关注 mask[start:end, start:end] = 0 # 应用于注意力计算：attention_score = QK^T / sqrt(d) + mask return mask

2. Masked Pooling（掩码池化）

问题背景

在标准ViT中：

通常使用[CLS]token作为整个图像的表示
或对所有patch token进行全局平均池化

但在打包序列中：

每个序列包含多张图像
需要为每张图像生成独立的表示
传统池化会混合不同图像的特征

掩码池化工作原理：

对序列中每个图像创建掩码
只考虑属于同一图像的tokens进行池化
每个图像独立产生输出表示

# 伪代码示例def masked_pooling(sequence, example_boundaries): pooled_outputs = [] for start, end in example_boundaries: # 提取属于同一图像的tokens image_tokens = sequence[start:end] # 对这些tokens进行池化（例如取平均） pooled = torch.mean(image_tokens, dim=0) pooled_outputs.append(pooled) return torch.stack(pooled_outputs)

这确保每个图像产生独立的表示，可以直接用于分类或其他下游任务。

2.1.2 分解位置嵌入

这是处理任意分辨率和长宽比的关键创新：

传统ViT的位置嵌入问题：对于尺寸R×R的方形图像，ViT学习长度为(R/P)²的一维位置嵌入，这无法自然适应不同长宽比的图像
Pix2Struct的2D绝对位置嵌入：学习尺寸为[maxLen, maxLen]的位置嵌入，但存在泛化问题——当分辨率超出训练范围时表现不佳
NaViT的解决方案：分解位置嵌入

绝对嵌入：基于绝对补丁索引
分数嵌入：基于相对位置( )
将(x,y)坐标分解为单独的嵌入和
然后通过相加等方式组合:
支持两种表示：
可与多种嵌入函数结合：学习嵌入、正弦嵌入或傅里叶嵌入

这种分解使模型能够处理任意分辨率的图像，并能更好地泛化到训练中未见过的分辨率（实验表明比Pix2Struct的2D嵌入表现更好）。

2.2 训练

得益于序列打包能力，NaViT引入了两项重要的训练技术：

2.2.1. 连续标记丢弃(Continuous Token Dropping)

传统做法：所有示例使用相同比例的标记Dropping
NaViT创新：每个图像的标记Dropping率可以独立变化
优势：

通过丢弃部分标记加速训练（减少计算量）
仍保留一些完整图像，减少训练-推理不一致
可在整个训练过程中动态调整丢弃率

实验表明，随着时间减少丢弃率（“Decreasing” schedule）可以进一步提高性能（图8）

2.2.2. 分辨率采样(Resolution Sampling)

NaViT可以使用图像的原始分辨率，或在保持长宽比的同时重新采样总像素数
关键创新：混合分辨率训练——从图像大小分布中采样训练，同时保持每个图像的原始长宽比
实验证明，可变分辨率预训练比固定分辨率表现更好（下图）
分辨率采样策略：实验发现直接采样边长（而非面积）并偏向较低分辨率（使用截断的正态分布）效果最佳

代码实现
=======

下面的代码实现是来自于一个非官方实现，作为参考：

3.1. 辅助函数部分

from functools import partialfrom typing import List, Unionimport torchimport torch.nn.functional as Ffrom einops import rearrange, repeatfrom torch import Tensor, nnfrom torch.nn.utils.rnn import pad_sequence as orig_pad_sequence# 辅助函数：检查值是否存在（不为None）def exists(val): return val is not None# 辅助函数：如果值存在则返回值，否则返回默认值def default(val, d): return val if exists(val) else d# 辅助函数：总是返回指定的值def always(val): return lambda *args: val# 辅助函数：将输入转换为元组（如果是单个值则重复两次）def pair(t): return t if isinstance(t, tuple) else (t, t)# 辅助函数：判断数值是否能被整除def divisible_by(numer, denom): return (numer % denom) == 0

3.2. 图像分组函数

def group_images_by_max_seq_len( images: List[Tensor], patch_size: int, calc_token_dropout = None, max_seq_len = 2048) -> List[List[Tensor]]: """ 根据最大序列长度将图像分组，确保每个组的总序列长度不超过最大限制 """ # 设置默认的令牌丢弃计算函数（不丢弃任何令牌） calc_token_dropout = default(calc_token_dropout, always(0.)) groups = [] # 存储分组结果 group = [] # 当前组 seq_len = 0 # 当前组的序列长度 # 如果令牌丢弃率是数字，则转换为始终返回该值的函数 if isinstance(calc_token_dropout, (float, int)): calc_token_dropout = always(calc_token_dropout) for image in images: assert isinstance(image, Tensor) # 确保输入是张量 # 计算图像的高度和宽度 image_dims = image.shape[-2:] # 计算补丁的高度和宽度 ph, pw = map(lambda t: t // patch_size, image_dims) # 计算图像的序列长度（补丁数量） image_seq_len = (ph * pw) # 应用令牌丢弃后的真实序列长度 image_seq_len = int(image_seq_len * (1 - calc_token_dropout (*image_dims))) # 确保单个图像的序列长度不超过最大限制 assert image_seq_len <= max_seq_len, f'image with dimensions {image_dims} exceeds maximum sequence length' # 如果当前组加上新图像会超过最大序列长度，则创建新组 if (seq_len + image_seq_len) > max_seq_len: groups.append(group) group = [] seq_len = 0 # 将图像添加到当前组 group.append(image) seq_len += image_seq_len # 添加最后一个组（如果有图像） if len(group) > 0: groups.append(group) return groups

3.3. 归一化层

# 无偏置的LayerNorm（PyTorch标准LayerNorm包含偏置，但这里不需要）class LayerNorm(nn.Module): def __init__(self, dim): super().__init__() self.gamma = nn.Parameter(torch.ones(dim)) # 学习的缩放参数 self.register_buffer('beta', torch.zeros(dim)) # 固定的偏置（零） def forward(self, x): return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)# RMSNorm（来自ViT-22B论文的查询-键归一化，无均值中心化，有学习的gamma）class RMSNorm(nn.Module): def __init__(self, heads, dim): super().__init__() self.scale = dim ** 0.5 # 缩放因子 self.gamma = nn.Parameter(torch.ones(heads, 1, dim)) # 每个头的学 习参数 def forward(self, x): # L2归一化 normed = F.normalize(x, dim = -1) return normed * self.scale * self.gamma

3.4. 前馈网络

def FeedForward(dim, hidden_dim, dropout = 0.): """ 前馈网络模块：LayerNorm -> 线性层 -> GELU -> Dropout -> 线性层 -> Dropout """ return nn.Sequential( LayerNorm(dim), # 层归一化 nn.Linear(dim, hidden_dim), # 第一个线性变换 nn.GELU(), # GELU激活函数 nn.Dropout(dropout), # Dropout nn.Linear(hidden_dim, dim), # 第二个线性变换 nn.Dropout(dropout) # Dropout )

3.5. 注意力机制

class Attention(nn.Module): def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.): super().__init__() inner_dim = dim_head * heads # 内部维度 self.heads = heads # 注意力头数 self.norm = LayerNorm(dim) # 输入归一化 # 查询和键的RMS归一化 self.q_norm = RMSNorm(heads, dim_head) self.k_norm = RMSNorm(heads, dim_head) self.attend = nn.Softmax(dim = -1) # 注意力权重计算 self.dropout = nn.Dropout(dropout) # 注意力权重的Dropout # 线性层：查询、键值对 self.to_q = nn.Linear(dim, inner_dim, bias = False) self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False) # 输出层 self.to_out = nn.Sequential( nn.Linear(inner_dim, dim, bias = False), # 输出线性变换 nn.Dropout(dropout) # Dropout ) def forward( self, x, # 输入张量 context = None, # 上下文张量（如果为None则使用x） mask = None, # 键掩码 attn_mask = None # 注意力掩码 ): x = self.norm(x) # 输入归一化 kv_input = default(context, x) # 如果没有上下文则使用输入x # 计算查询、键和值 qkv = (self.to_q(x), *self.to_kv(kv_input).chunk(2, dim = -1)) # 重塑为多头格式 q, k, v = map( lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self. heads), qkv ) # 查询和键的RMS归一化 q = self.q_norm(q) k = self.k_norm(k) # 计算注意力分数 dots = torch.matmul(q, k.transpose(-1, -2)) # 应用键掩码 if exists(mask): mask = rearrange(mask, 'b j -> b 1 1 j') dots = dots.masked_fill(~mask, -torch.finfo(dots.dtype).max) # 应用注意力掩码 if exists(attn_mask): dots = dots.masked_fill(~attn_mask, -torch.finfo(dots.dtype). max) # 计算注意力权重 attn = self.attend(dots) attn = self.dropout(attn) # 计算输出 out = torch.matmul(attn, v) out = rearrange(out, 'b h n d -> b n (h d)') # 重新整形 return self.to_out(out)

3.6. Transformer块

class Transformer(nn.Module): def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.): super().__init__() # 创建多个Transformer层 self.layers = nn.ModuleList([]) for _ in range(depth): self.layers.append(nn.ModuleList([ Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout), # 注意力层 FeedForward(dim, mlp_dim, dropout = dropout) # 前馈网络 ])) self.norm = LayerNorm(dim) # 最终归一化 def forward( self, x, # 输入张量 mask = None, # 掩码 attn_mask = None # 注意力掩码 ): # 通过每个Transformer层 for attn, ff in self.layers: x = attn(x, mask = mask, attn_mask = attn_mask) + x # 残差连 接 x = ff(x) + x # 残差连接 return self.norm(x)

3.7. 主模型 NaViT

class NaViT(nn.Module): def __init__( self, *, image_size, # 图像尺寸（高度，宽度） patch_size, # 补丁尺寸 num_classes, # 分类数量 dim, # 模型维度 depth, # Transformer深度 heads, # 注意力头数 mlp_dim, # MLP隐藏层维度 channels = 3, # 通道数（默认RGB） dim_head = 64, # 每个注意力头的维度 dropout = 0., # Dropout率 emb_dropout = 0., # 嵌入层Dropout率 token_dropout_prob = None # 令牌丢弃概率 ): super().__init__() image_height, image_width = pair(image_size) # 获取图像高度和宽度 # 令牌丢弃策略：如果是函数则直接使用，如果是数值则转换为常数函数 self.calc_token_dropout = None if callable(token_dropout_prob): self.calc_token_dropout = token_dropout_prob elif isinstance(token_dropout_prob, (float, int)): assert 0. < token_dropout_prob < 1. # 确保概率在0到1之间 token_dropout_prob = float(token_dropout_prob) self.calc_token_dropout = lambda height, width: token_dropout_prob # 补丁相关计算 assert divisible_by(image_height, patch_size) and divisible_by (image_width, patch_size), 'Image dimensions must be divisible by the patch size.' # 计算补丁的高度和宽度维度 patch_height_dim, patch_width_dim = (image_height // patch_size), (image_width // patch_size) # 每个补丁的维度 patch_dim = channels * (patch_size ** 2) self.channels = channels self.patch_size = patch_size # 补丁嵌入层：LayerNorm -> 线性层 -> LayerNorm self.to_patch_embedding = nn.Sequential( LayerNorm(patch_dim), nn.Linear(patch_dim, dim), LayerNorm(dim), ) # 2D位置嵌入：分别用于高度和宽度 self.pos_embed_height = nn.Parameter(torch.randn (patch_height_dim, dim)) self.pos_embed_width = nn.Parameter(torch.randn(patch_width_dim, dim)) self.dropout = nn.Dropout(emb_dropout) # Transformer主干 self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout) # 最终注意力池化的查询 self.attn_pool_queries = nn.Parameter(torch.randn(dim)) self.attn_pool = Attention(dim = dim, dim_head = dim_head, heads = heads) # 输出到logits self.to_latent = nn.Identity() self.mlp_head = nn.Sequential( LayerNorm(dim), nn.Linear(dim, num_classes, bias = False) # 最终分类层（无偏 置） ) @property def device(self): # 获取模型参数所在的设备 return next(self.parameters()).device def forward( self, batched_images: Union[List[Tensor], List[List[Tensor]]], # 批次 图像（假设不同分辨率图像已正确分组） group_images = False, # 是否自动分组图像 group_max_seq_len = 2048 # 分组的最大序列长度 ): # 获取模型参数 p, c, device, has_token_dropout = self.patch_size, self. channels, self.device, exists(self.calc_token_dropout) # 定义辅助函数 arange = partial(torch.arange, device = device) pad_sequence = partial(orig_pad_sequence, batch_first = True) # 如果需要，自动将图像按最大序列长度分组 if group_images: batched_images = group_images_by_max_seq_len( batched_images, patch_size = self.patch_size, calc_token_dropout = self.calc_token_dropout, max_seq_len = group_max_seq_len ) # 将图像处理为可变长度序列并生成注意力掩码 num_images = [] # 每个批次中的图像数量 batched_sequences = [] # 批次序列 batched_positions = [] # 批次位置 batched_image_ids = [] # 批次图像ID # 遍历每个批次 for images in batched_images: num_images.append(len(images)) # 记录图像数量 sequences = [] # 当前批次的序列 positions = [] # 当前批次的位置 # 创建空的图像ID张量 image_ids = torch.empty((0,), device = device, dtype = torch. long) # 遍历批次中的每个图像 for image_id, image in enumerate(images): # 验证图像维度 assert image.ndim == 3 and image.shape[0] == c image_dims = image.shape[-2:] # 验证图像尺寸能被补丁尺寸整除 assert all([divisible_by(dim, p) for dim in image_dims]), f'height and width {image_dims} of images must be divisible by patch size {p}' # 计算补丁的高度和宽度 ph, pw = map(lambda dim: dim // p, image_dims) # 创建位置网格 pos = torch.stack(torch.meshgrid(( arange(ph), arange(pw) ), indexing = 'ij'), dim = -1) pos = rearrange(pos, 'h w c -> (h w) c') # 重塑位置 # 将图像转换为补丁序列 seq = rearrange(image, 'c (h p1) (w p2) -> (h w) (c p1 p2)', p1 = p, p2 = p) seq_len = seq.shape[-2] # 序列长度 # 如果需要令牌丢弃 if has_token_dropout: token_dropout = self.calc_token_dropout (*image_dims) # 计算丢弃率 num_keep = max(1, int(seq_len * (1 - token_dropout))) # 计算保留数量 # 随机选择要保留的索引 keep_indices = torch.randn((seq_len,), device = device).topk(num_keep, dim = -1).indices seq = seq[keep_indices] # 保留选中的序列 pos = pos[keep_indices] # 保留选中的位置 # 更新图像ID image_ids = F.pad(image_ids, (0, seq.shape[-2]), value = image_id) sequences.append(seq) positions.append(pos) # 将当前批次的数据添加到总批次中 batched_image_ids.append(image_ids) batched_sequences.append(torch.cat(sequences, dim = 0)) batched_positions.append(torch.cat(positions, dim = 0)) # 生成键填充掩码 lengths = torch.tensor([seq.shape[-2] for seq in batched_sequences], device = device, dtype = torch.long) max_length = arange(lengths.amax().item()) key_pad_mask = rearrange(lengths, 'b -> b 1') <= rearrange (max_length, 'n -> 1 n') # 生成注意力掩码，并与键填充掩码结合 batched_image_ids = pad_sequence(batched_image_ids) attn_mask = rearrange(batched_image_ids, 'b i -> b 1 i 1') == rearrange(batched_image_ids, 'b j -> b 1 1 j') attn_mask = attn_mask & rearrange(key_pad_mask, 'b j -> b 1 1 j') # 合并补丁图像以及补丁宽度/高度位置用于2D位置嵌入 patches = pad_sequence(batched_sequences) patch_positions = pad_sequence(batched_positions) # 需要知道每个批次中有多少图像用于最终注意力池化 num_images = torch.tensor(num_images, device = device, dtype = torch.long) # 补丁嵌入 x = self.to_patch_embedding(patches) # 因子化的2D绝对位置嵌入 h_indices, w_indices = patch_positions.unbind(dim = -1) # 分离高 度和宽度索引 h_pos = self.pos_embed_height[h_indices] # 获取高度位置嵌入 w_pos = self.pos_embed_width[w_indices] # 获取宽度位置嵌入 x = x + h_pos + w_pos # 添加位置嵌入 # 嵌入Dropout x = self.dropout(x) # Transformer处理 x = self.transformer(x, attn_mask = attn_mask) # 执行注意力池化 max_queries = num_images.amax().item() # 最大查询数 # 重复池化查询 queries = repeat(self.attn_pool_queries, 'd -> b n d', n = max_queries, b = x.shape[0]) # 注意力池化掩码 image_id_arange = arange(max_queries) attn_pool_mask = rearrange(image_id_arange, 'i -> i 1') == rearrange(batched_image_ids, 'b j -> b 1 j') attn_pool_mask = attn_pool_mask & rearrange(key_pad_mask, 'b j -> b 1 j') attn_pool_mask = rearrange(attn_pool_mask, 'b i j -> b 1 i j') # 注意力池化 x = self.attn_pool(queries, context = x, attn_mask = attn_pool_mask) + queries x = rearrange(x, 'b n d -> (b n) d') # 重塑 # 每个批次元素可能有不同的图像数量 is_images = image_id_arange < rearrange(num_images, 'b -> b 1') is_images = rearrange(is_images, 'b n -> (b n)') x = x[is_images] # 只保留实际图像的特征 # 投影到潜在空间 x = self.to_latent(x) return self.mlp_head(x) # 返回分类logits

关键技术点解析

可变分辨率支持：模型可以处理不同尺寸的图像，这是通过将图像分割成补丁并动态处理序列长度实现的。
因子化2D位置嵌入：将2D位置嵌入分解为高度和宽度两个独立的嵌入，然后相加。
注意力池化：使用可学习的查询向量进行最终的池化操作，而不是传统的[CLS]标记。
令牌丢弃：根据图像尺寸动态计算令牌丢弃概率，提高模型效率。
批次处理：支持将不同分辨率的图像分组处理，优化计算效率。

AI大模型从0到精通全套学习大礼包

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

只要你是真心想学AI大模型，我这份资料就可以无偿共享给你学习。大模型行业确实也需要更多的有志之士加入进来，我也真心希望帮助大家学好这门技术，如果日后有什么学习上的问题，欢迎找我交流，有技术上面的问题，我是很愿意去帮助大家的！

如果你也想通过学大模型技术去帮助就业和转行，可以扫描下方链接👇👇
大模型重磅福利：入门进阶全套104G学习资源包免费分享！

01.从入门到精通的全套视频教程

包含提示词工程、RAG、Agent等技术点

02.AI大模型学习路线图（还有视频解说）

全过程AI大模型学习路线

03.学习电子书籍和技术文档

市面上的大模型书籍确实太多了，这些是我精选出来的

04.大模型面试题目详解

05.这些资料真的有用吗?

这份资料由我和鲁为民博士共同整理，鲁为民博士先后获得了北京清华大学学士和美国加州理工学院博士学位，在包括IEEE Transactions等学术期刊和诸多国际会议上发表了超过50篇学术论文、取得了多项美国和中国发明专利，同时还斩获了吴文俊人工智能科学技术奖。目前我正在和鲁博士共同进行人工智能的研究。

所有的视频由智泊AI老师录制，且资料与智泊AI共享，相互补充。这份学习大礼包应该算是现在最全面的大模型学习资料了。

资料内容涵盖了从入门到进阶的各类视频教程和实战项目，无论你是小白还是有些技术基础的，这份资料都绝对能帮助你提升薪资待遇，转行大模型岗位。

智泊AI始终秉持着“让每个人平等享受到优质教育资源”的育人理念‌，通过动态追踪大模型开发、数据标注伦理等前沿技术趋势‌，构建起"前沿课程+智能实训+精准就业"的高效培养体系。

课堂上不光教理论，还带着学员做了十多个真实项目。学员要亲自上手搞数据清洗、模型调优这些硬核操作，把课本知识变成真本事‌！

如果说你是以下人群中的其中一类，都可以来智泊AI学习人工智能，找到高薪工作，一次小小的“投资”换来的是终身受益！

应届毕业生‌：无工作经验但想要系统学习AI大模型技术，期待通过实战项目掌握核心技术。

零基础转型‌：非技术背景但关注AI应用场景，计划通过低代码工具实现“AI+行业”跨界‌。

业务赋能 ‌突破瓶颈：传统开发者（Java/前端等）学习Transformer架构与LangChain框架，向AI全栈工程师转型‌。

👉获取方式：
😝有需要的小伙伴，可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓