本文关注于通过模型设计和优化提高VLA架构的效率,而不是依赖推断后压缩的策略。具体地,设计了一种对于机器人操作的全1比特VLA模型,其中每个参数是三值的{-1,0,1}。BitVLA建立在公开的1比特LLM模型BitNeTb1.58,并训练成一个视觉语言动作策略。该模型继承了BitNet模型的紧凑性,同时保持了强大的性能。为了进一步减少视觉Backbone的内存,本文引入了一个量化后蒸馏的策略,可以将全精度视觉编码器压缩到1.58位权重,同时在训练过程中由全精度教师指导表示对齐。
模型架构
本文架构将1比特VLA模型BitVLA实例化为统一的多模态策略πθ\pi_{\theta}πθ,他将视觉观察和语言上下文映射为机器人动作。BitVLA使用BitNetb1.58作为1比特LLM并使用SigLIP-L作为视觉编码器。
统一的策略描述
定义时间ttt的多模态输入ot=[It,qt]o_{t}=[I_{t},q_{t}]ot=[It,qt],分别对应图像观察和机械臂状态,定义ℓt\ell_{t}ℓt是语言上下文。策略在短视野hhh下输出动作块at:t+ha_{t:t+h}at:t+h。BitVLA获得的分步表示为πθ(a^t:t+h∣ot,ℓt)\pi_{\theta}(\hat{a}_{t:t+h}|o_{t},\ell_{t})πθ(a^t:t+h∣ot,ℓt)。本文框架中将输入令牌化为x1:Nx_{1:N}x1:N序列,序列由ℓt\ell_{t}ℓt获得的文本令牌xiwx_{i}^{w}xiw,有SigLIP-L和ItI_{t}It获得的视觉令牌和机器人状态令牌xiqx_{i}^{q}xiq。本文框架通过以相应隐藏状态为条件的动作头来解码连续的机器人动作。
权重和激活的量化
对于权重使用absmean量化器,对于激活使用每令牌的absmax量化器。权重量化为三数值{-1.0.1},激活量化为对称的INT8 [-128,127]。
Qw(W)=RoundClip(Wα,−1,1),α=1nm∣∣W∣∣1 Q_{w}(W)=\text{RoundClip}(\frac{W}{\alpha},-1,1), \quad \alpha=\frac{1}{nm}||W||_{1}Qw(W)=RoundClip(αW,−1,1),α=nm1∣∣W∣∣1
Qa(x)=RoundClip(127xβ,−128,127),β=∣∣x∣∣∞ Q_{a}(x)=\text{RoundClip}(\frac{127x}{\beta},-128,127),\quad \beta=||x||_{\infty}Qa(x)=RoundClip(β127x,−128,127),β=∣∣x∣∣∞
RoundClip(x,a,b)=max(a,min(b,round(x))) \text{RoundClip}(x,a,b)=\max(a,\min(b,\text{round}(x)))RoundClip(x,a,b)=max(a,min(b,round(x)))
其中WWW是线性层的可学习权重,xxx是对应的输入激活。本文框架中在视觉编码器和LLM模型所有线性层(除了输入和输出的嵌入层)都使用量化操作。
训练策略
整体训练分为三个阶段:1. 使用1比特LLM和全精度视觉编码器进行多多模态训练以获得稳定的视觉语言初始化,2. 量化后蒸馏阶段:压缩视觉编码器为1.58比特权重和8比特激活同时保证多模态对齐,3. 从大规模轨迹数据中获取跨实体操作知识并实现下游任务的机器人训练。
多模态训练
遵循LLaVA架构,首先通过组合1比特LLM和全精度视觉编码器,轻量级连接器构建视觉语言模型。该阶段包含两个步骤:我们首先仅在小型图像标题语料库上训练连接器,以将视觉标记空间与语言嵌入空间对齐。然后冻结视觉编码器优化剩下模块以提升指令听从能力。
量化后蒸馏
使用知识蒸馏策略将视觉编码器量化为1.58比特权重和8比特激活值。首先从对应的全精度编码器初始化1.58比特的学生编码器,全精度编码器作为固定的教师模型。在这一阶段,仅更新学生视觉编码器,并保持1比特LLM模型和连接器冻结。学生视觉编码器利用量化函数和STE用于梯度估计。整体损失函数包含语言模型损失LLM\mathcal{L}_{LM}LLM和辅助表示对齐损失Laux\mathcal{L}_{aux}Laux。
语言模型损失
使用标准的自回归语言模型目标。定义T\mathcal{T}T是 输入文本序列(包含指令前缀Tins\mathcal{T}_{ins}Tins和回答后缀Tans\mathcal{T}_{ans}Tans)。定义V1.58\mathcal{V}_{1.58}V1.58是1.58比特学生编码器产生的视觉令牌。损失函数定义为:
LLM=−∑i∈Tanslogp(yi∣V1.58,T<i) \mathcal{L}_{LM}=-\sum_{i\in \mathcal{T}_{ans}}\log p(y_{i}|\mathcal{V}_{1.58},\mathcal{T}_{<i})LLM=−i∈Tans∑logp(yi∣V1.58,T<i)
表示对齐损失
为了缓解量化引发的表示漂移并保持多模态对齐,需要蒸馏从全精度教师编码器的中间特征。定义hbf16lh_{bf16}^{l}hbf16l和h1.58lh_{1.58}^{l}h1.58l是从教师和学生编码器第l层获得的隐藏状态。定义表示对齐损失为:
Laux=1L∑l=1L1n∣∣hbf16l−h1.58l∣∣22 \mathcal{L}_{aux}=\frac{1}{L}\sum_{l=1}^{L}\frac{1}{n}||h_{bf16}^{l}-h_{1.58}^{l}||_{2}^{2}Laux=L1l=1∑Ln1∣∣hbf16l−h1.58l∣∣22
机器训练
在量化后蒸馏阶段后,视觉编码器和LLM模型都能在1.58比特和8比特激活上推断。为了让 BitVLA 拥有可跨不同表现形式和环境迁移的一般可操作先验,本文用自回归的下一步动作预测目标对它进行预训练。每个动作维度都独立地离散化为256个区间,模型的训练目标是最小化预测动作标记和真实动作序列之间的交叉熵损失。
在预训练后,执行监督微调以适应下游任务。跟随 OpenVLA OFT 的推理过程,BitVLA 使用带动作分块的并行解码来提高实时控制的推理吞吐量