对数深度Transformer的表达能力

引言

传统理论认为Transformer的表达能力受限于固定深度——无论输入多长，深度都是常数。这导致许多简单问题（如正则语言识别）在理论上被认为Transformer无法表达。

然而，这一结论建立在任意长度假设之上。实际情况是：实际应用中的输入长度总是有界的——上下文窗口有最大长度限制。

核心问题：在有界上下文设定下，Transformer的表达能力如何随深度变化？[¹

本文介绍的研究证明：仅需 $O(\log n)$ 深度，Transformer就能表达多种重要的问题类，包括正则语言识别和图连通性判断。

有界上下文设定

形式化定义

设 $n$ 为最大上下文长度。考虑一个深度为 $L(n)$ 的Transformer，其深度可以随 $n$ 增长。关键假设：

$$L(n)=O(\log n)$$

这被称为对数深度Transformer（Log-depth Transformers）。

与固定深度的对比

设定	深度	可表达问题
固定深度（传统）	$O(1)$	TC^0 子集
对数深度	$O(\log n)$	正则语言、图连通性等
线性深度	$O(n)$	更复杂问题

关键洞察：对数深度是固定深度和线性深度之间的自然中间地带，在实践中容易实现。

正则语言识别

问题背景

正则语言是正则表达式可描述的语言类，是计算理论中最基础的表示类之一。

传统观点认为，Transformer无法识别某些简单正则语言（如 $a^n{b}^n$）。这一观点在固定深度设定下是正确的，但在有界上下文下是错误的。

核心定理

定理（正则语言识别）：对数深度Transformer可以识别任意正则语言。

形式化：对于任意正则语言 $L$，存在深度 $L=O(\log n)$ 的Transformer，使得：

$$\forall \mathbf{x}\in \{0,1{\}}^{\le n}:\text{Transformer}(\mathbf{x})=1_L(\mathbf{x})$$

证明思路

1. 从有限状态自动机出发

正则语言 $L$ 由有限状态自动机（DFA）接受。设 DFA 有 $m$ 个状态。

2. 对数深度模拟

关键思想：使用Transformer的层数来模拟自动机的状态转移。

DFA状态转移：           Transformer层：
┌───────┐              ┌─────────┐
│ q_0   │──a──→        │ Layer 1 │
└───────┘              └────┬────┘
      │                        │
      │                        ▼
┌───────┐              ┌─────────┐
│ q_1   │──b──→        │ Layer 2 │
└───────┘              └────┬────┘
      │                        │
      │                       ...
      │                        ▼
┌───────┐              ┌─────────┐
│ q_m   │              │ Layer k │
└───────┘              └─────────┘

3. 计数机制

对于 $a^n{b}^n$ 这样的语言，Transformer需要计数 $a$ 的数量。实现方法：

• 浅层：使用”计数器”token编码当前计数
• 对数深度：通过二进制表示实现高效计数

# 计数模拟示意
def count_tokens(transformer_output, token_type):
    """
    使用对数深度模拟计数器
    """
    # 提取二进制编码的计数值
    binary_count = extract_binary_count(transformer_output)
    # 逐层增加（log深度）
    for layer in range(int(math.log2(max_count))):
        binary_count = increment_layer(binary_count, layer)
    return binary_count

图连通性判断

问题定义

图连通性（Graph Connectivity）：判断无向图中两个顶点是否连通。

输入格式：图的邻接矩阵或边列表。

为什么重要

图连通性是许多图算法的基础：

• 路径查找
• 连通分量计算
• 生成树验证

它是多跳推理的原型问题。

核心定理

定理（图连通性）：对数深度Transformer可以判断 $n$ 顶点图的连通性。

时间复杂度：$O(n\log n)$ 参数量和 $O(\log n)$ 深度。

实现机制

消息传递视角

连通性判断等价于分布式消息传递：

1. 每个顶点维护”可达性”状态
2. 沿边传播状态
3. $O(\log n)$ 轮后，所有可达顶点对已知

时间步 0:     时间步 1:     时间步 2:
A ─ B        A*──B*       A*──B*
│  │         │    │       │    │
│  │    →    │    │   →   │    │
C ─ D        C*──D*       C*──D*

* 表示已标记为从起点可达

Transformer实现

对数深度Transformer模拟消息传递：

$${\mathbf{h}}^{(l+1)}=\text{Attention}\left({\mathbf{h}}^{(l)},\text{Adjacency}\right)$$

其中 ${\mathbf{h}}^{(l)}$ 在 $l$ 步后编码可达信息。

与GNN的联系

有趣的是，对数深度Transformer与图神经网络（GNN）在连通性任务上有类似表达能力：

方面	GraphSAGE	Transformer
消息传递	$O(k)$ hops	$O(\log n)$ layers
表达能力	受限于聚合	受限于深度
适用场景	图结构明确	结构隐式编码

深度 vs 宽度 vs 链式思维

效率比较

研究的一个关键贡献是量化比较不同扩展计算能力的方式：

策略	深度增长	宽度增长	链式思维步数
正则语言	$O(\log n)$	$O(1)$	$O(n)$
图连通性	$O(\log n)$	$O(n)$	$O(n)$
更复杂问题	$O(n)$	$O(n)$	$O(\log n)$

结论：对于正则语言和图连通性，增加深度是最有效的策略。

定量分析

设目标计算能力为 $C$。不同策略的资源需求：

$$\text{深度策略}:L=\Theta (\log C),d=O(1)$$ $$\text{宽度策略}:L=O(1),d=\Theta (C)$$ $$\text{链式思维}:L=O(1),\text{推理步数}=\Theta (\log C)$$

效率排序（参数量角度）：深度 > 链式思维 > 宽度

实际训练意义

深度需求的经验验证

研究包含大量实验，验证理论预测的深度需求：

正则语言识别所需深度（经验 vs 理论）

语言复杂度     理论深度    经验深度(中位数)
──────────────────────────────────────
a^n            O(1)        1.2
a^n b^n         O(log n)    3.8
(a|b)^n        O(log n)    4.2
a^n b^n c^n     O(log n)    5.1

发现：经验深度与理论预测高度吻合。

宽度缩放

同时发现宽度缩放对学习的影响：

• 宽度不足：即使增加深度也无法弥补
• 最优宽度：$d\approx 4\times$ 任务所需状态数
• 过度宽度：超过最优后收益递减

与TC^0的关系

TC^0复杂性类

TC^0 是常数深度、多项式大小的阈值电路类：

• 允许计数（阈值门）
• 不允许进位传播（完整加法）

之前结果：固定深度Transformer表达能力 ⊆ TC

新结果

定理（对数深度与TC^0）：

$$\text{Log-depth Transformer}\equiv {\text{TC}}^0$$

即对数深度Transformer恰好表达能力与TC^0相同。

含义

这一等价性意味着：

1. 下界：对数深度Transformer不能表达上下文无关语言
2. 上界：任何TC^0可计算的问题，对数深度Transformer都能解决
3. 完整性：TC^0是研究Transformer表达能力的自然基准

实践指南

任务-架构匹配

基于理论，为不同任务推荐架构：

1. 简单模式匹配

• 任务：正则模式、关键词检测
• 推荐深度：1-2层
• 推荐宽度：中等

2. 状态跟踪

• 任务：括号匹配、协议状态机
• 推荐深度：$O(\log n)$，如6-8层
• 推荐宽度：较大

3. 多跳推理

• 任务：知识图谱问答、逻辑推导
• 推荐深度：8-12层或链式思维
• 推荐策略：深度 + 链式思维结合

深度选择算法

def recommend_depth(task_complexity, max_context_length):
    """
    基于任务复杂度推荐Transformer深度
    """
    if task_complexity == "regular":
        # 正则语言：对数深度
        return int(math.log2(max_context_length)) + 2
    elif task_complexity == "graph":
        # 图推理：对数深度
        return int(math.log2(max_context_length)) + 4
    elif task_complexity == "context_free":
        # 上下文无关：需要链式思维
        return 12  # 基础深度
    else:
        return 8  # 默认保守选择

局限性与未来方向

当前局限性

1. 有界假设：所有结果基于最大上下文长度 $n$ 的假设
2. 精度假设：理论分析假设精确算术
3. 同步假设：通信模型假设同步执行

未解问题

1. 更复杂电路类：对数深度能表达AC^0吗？
2. 噪声鲁棒性：有限精度如何影响表达能力？
3. 学习理论：为什么Transformer能学到这些表达？

总结

对数深度Transformer的研究揭示了几个重要发现：

1. 正则语言：仅 $O(\log n)$ 深度即可识别任意正则语言
2. 图连通性：通过模拟消息传递实现图算法
3. 效率优势：深度扩展比宽度和链式思维更高效
4. TC^0等价：对数深度Transformer恰好表达能力与TC^0相同

这些结果为架构设计提供了理论基础：增加深度是提升表达能力的最有效方式，特别是在有界上下文的实际应用场景中。

参考

Footnotes

1. A Little Depth Goes a Long Way: The Expressive Power of Log-Depth Transformers. ICLR 2024 ↩