JAX-LOB：使用GPU加速限价订单簿仿真

前言

交易所利用限价订单簿（LOB）来处理订单并匹配交易。为了研究目的，拥有大规模高效的LOB动态模拟器是非常重要的。以往，LOB模拟器已经在代理模型（ABMs）、强化学习（RL）环境和生成模型中实施，处理来自历史数据集和手工代理的订单流。对于许多应用，需要处理多个簿，无论是用于ABMs的校准还是RL代理的训练。我们展示了第一个GPU加速的LOB模拟器，名为JAX-LOB，旨在并行处理数千个簿，并显著减少每条消息的处理时间。我们的模拟器的实现基于设计选择，旨在充分利用JAX的功能，同时不影响与LOB相关机制的真实性。

什么是LOB？

LOB是Limit Order Book（限价订单簿）的缩写，是金融交易所用于记录和管理证券买卖订单的一种数据结构。在LOB中，买卖订单按照价格和时间顺序排列，以便进行撮合交易。交易员可以提交限价订单（需指定价格和数量）或市价订单（只需指定数量），这些订单会被匹配或者添加到簿中。LOB的动态变化取决于订单流的变化，包括新订单的提交、订单的撤销和交易的执行。LOB在现代电子交易中起着至关重要的作用，对于交易策略的开发和评估具有重要意义。

为什么要模拟LOB？

模拟限价订单簿（LOB）具有重要意义，因为它可以帮助我们更好地理解金融市场的运作和交易行为。通过模拟LOB，我们可以研究和分析不同的交易策略、市场影响以及交易执行的效果。这对于金融公司改进服务、政府预测金融监管对金融体系稳定性的影响等方面具有实际意义。此外，LOB模拟器还可以用于训练交易代理人，如市场做市商或交易执行代理人，以及用于生成模型等领域的研究。由于这些应用通常需要大规模的模拟和计算，因此开发GPU加速的LOB模拟器可以提高计算效率，从而为这些研究提供更好的工具和平台。

什么是JAX？

JAX是一个充满潜力的高性能数值计算库，它将可微分编程带入了Python生态系统。它可以与不同类型硬件匹配的加速器框架，它利用加速的线性代数（XLA）、自动微分和自动向量化，可以轻松地在GPU上执行。这个框架旨在进行高性能的机器学习研究，是gymnax和PureJaxRL框架的基础。JAX的设计使其能够实现即时（JIT）编译，从而在GPU上执行。JAX的特性使其非常适合用于大规模的机器学习任务，而且它的加速特性使得它非常适合用于处理LOB模拟器这样的任务。

相对CPU的优势：

• JAX是一个加速器不可知的框架，可以使用GPU进行即时编译（JIT）和加速线性代数（XLA），自动微分和自动向量化；
• JAX旨在进行高性能机器学习研究，并且可以轻松地在GPU上执行；
• JAX具有自动向量化功能，可以将代码转换为可以在GPU上并行执行的形式，从而提高了计算速度；
• 在使用JAX进行训练时，可以避免GPU-CPU通信瓶颈，从而提高了训练速度；
• 在使用JAX进行训练时，可以利用GPU的并行性，从而提高了训练速度。

JAX-LOB模拟器

将深度强化学习应用于交易执行和其他高频任务的一些关键挑战是金融数据的信号噪声比低、可能出现的对于特定训练日的过度拟合，及真实模拟市场冲击。弥补前两个问题的一种简单方法是增加可用于训练的状态-动作转换数量。为了加速使用高频数据生成LOB，我们使用JAX。

本文使用以下设计决策来应对JIT-compiled JAX代码的一些限制，这些设计决策旨在解决JIT-compiled JAX代码的一些限制，例如无法使用动态数据结构和控制流程。通过使用纯函数和固定大小的数组，可以使代码更容易编译和并行化：

• 纯函数
• 固定的数组大小和类型
• 可以有效编译和并行化的控制流程

大多数CPU实现的限价订单簿（LOB）都基于哈希映射、队列、双向链表和排序字典，这些数据结构可以快速访问数据并在整个订单簿中保持排序。然而，由于JAX要求编译时固定大小的数组，要实现类似的结构就意味着必须预先为所有价格级别和订单分配内存空间。使用数组意味着在删除条目时重新排序的成本远高于使用链表。

因此，本文选择了一种不使用类似树状结构的架构，也不要求始终保持订单排序的方法。本文定义了两个数组A和B来表示订单簿的两侧，其中A表示所有活动的卖出订单，B表示所有活动的买入订单。这种设计避免了使用类似树状结构的数据结构，也不要求始终保持订单的排序。

基础操作

对于LOB的基本操作有三种：添加新订单、取消现有订单和将现有订单与另一侧的订单进行匹配并从订单簿中删除。添加订单需要找到数组中的空位置，并将订单特定数据插入正确的字段。取消订单需要定位要取消的订单的订单ID，并从订单簿中相应地移除数量。匹配订单需要考虑现有订单与另一侧的订单进行匹配，并随后将其从订单簿中移除。

• 添加订单需要在数组中识别一个空位置（𝒐𝒊 = -1），并将订单特定数据插入正确的字段。
• 取消订单需要定位要取消的订单的订单ID（𝑂𝐼𝐷），并从订单簿中相应地移除数量。这些操作对于管理订单簿至关重要，对于维护交易系统的完整性和准确性至关重要。
• 在匹配操作期间，一个被称为𝒐𝒂的主动订单会与订单簿另一侧的现有订单（即挂单𝒐𝒔）进行匹配。匹配逻辑旨在通过以下操作找到主动订单（𝑄𝑎 ∈ 𝒐𝒂）和挂单（𝑄𝑠 ∈ 𝒐𝒔）的剩余数量：
  • 从挂单中减去主动订单的数量，得到剩余的挂单数量：𝑄′𝑠=𝑚𝑎𝑥 (0, 𝑄𝑠 − 𝑄𝑎 )
  • 从主动订单中减去已匹配的挂单数量，得到剩余的主动订单数量：𝑄′𝑎=𝑄𝑎 − 𝑄𝑠

当两个订单匹配成功时，会记录一次交易𝒕𝒋。交易记录包括以下信息：

• 交易价格：𝑃𝑗 = 𝑃𝑠
• 交易数量：𝑄 𝑗 = 𝑄𝑠 − 𝑄′𝑠
• 主动订单ID：𝑂𝐼𝐷𝑎,𝑗 = 𝑂𝐼𝐷𝑎 ∈ 𝒐𝒂
• 挂单ID：𝑂𝐼𝐷𝑠,𝑗 = 𝑂𝐼𝐷𝑠 ∈ 𝒐𝒔
• 挂单时间戳：𝑇𝑠𝑗 ,𝑇𝑛𝑠𝑗 = 𝑇𝑠𝑎,𝑇𝑛𝑠𝑎

由于XLA编译器的限制，最多记录N次交易记录，存储在固定大小的数组T中。如果某个交易记录为空，则𝒕𝒋 = −1。在所有操作完成后，会检查𝐴𝑠𝑘𝑠和𝐵𝑖𝑑𝑠中的所有订单𝒐𝒊，如果𝑄𝑖 ≤ 0，则将𝒐𝒊设置为−1。

一笔主动订单可能会与多笔挂单匹配。匹配逻辑包含一个while循环，不断尝试将主动订单与位于订单簿另一侧的下一个最佳挂单进行匹配。最佳挂单𝐵𝑒𝑠𝑡(𝒐𝒔)由价格-时间优先级算法定义，这是最常用的限价订单簿匹配算法。该算法将价格和时间作为优先级的两个因素，优先选择价格更优且时间更早的挂单进行匹配。如果有多个订单共享相同价格，则选择最早到达时间的订单。这种匹配逻辑可以提高订单簿的流动性，使得更多的订单可以被匹配成交，从而提高市场效率。

匹配循环将继续进行，直到订单簿非空且满足价格重叠的条件。具体地，如果是可成交的卖单，要求卖出价格低于等于买入价格；如果是可成交的买单，要求买入价格高于等于卖出价格。这保证了只有在价格满足成交条件时，匹配循环才会继续进行。

最后，根据上述描述，可以直观地预期这些基本操作的计算复杂性会有所不同。为了验证这一直觉，作者对不同最大容量𝑁的订单簿中的基本操作进行了计时。结果显示，随着订单簿容量𝑁的增加，操作所需时间也会增加，并且最慢的操作（匹配）执行时间超过最快操作（取消）的两倍以上。这表明，在处理大量订单时，匹配操作的计算复杂性会显著增加，需要采取相应的优化措施来提高系统性能。

消息类型

message指的是传递给订单簿的消息，用于提交、取消或删除订单，或者执行市价订单。每个message包含以下信息：

1、message类型（limit orders，cancel orders，delete orders，market orders）

2、订单的方向（bid或ask）

3、订单的价格（对于市价订单，该字段被忽略）

4、订单的数量

5、订单的唯一标识符

6、时间戳

7、订单簿的标识符

8、订单簿中的时间戳

这些信息被传递给订单簿，以便进行匹配和执行。通过使用message，可以实现高效的订单簿管理和交易执行。

本文将每种订单类型和方向的订单视为八种单独的情况，这样可以为每种情况定义明确的函数。这样做可以在接收到消息时使用单个条件语句，而不是在匹配逻辑中使用多个分支。作者发现，这种方法在vmap下可以提高性能。

处理每种三种消息类型的计算时间因所需的基本操作而异。与表1中的结果相比，分支语句只会略微增加处理时间，而不同订单类型之间仍然存在显著差异。这表明，将每种情况定义为明确的函数，并在接收到消息时使用单个条件语句的方法有效地提高了性能，并减少了计算复杂性，特别是在处理不同类型和方向的订单时。

使用vmap加速处理订单信息

"vmap" 是指 JAX 库中的一个操作符，用于实现向量化的映射（vectorized map）。这个操作符允许用户对函数进行向量化，以便在 GPU 或 TPU 等加速器上并行处理多个输入。在订单簿匹配系统中，使用 vmap 可以同时处理多个订单簿，从而提高整体的处理效率。

具体来说，vmap 操作符将函数映射到输入的批处理维度上，使得函数能够以向量化的方式处理输入。这对于需要在大规模数据上执行相同操作的情况非常有用，因为它可以利用硬件加速器的并行计算能力，从而提高整体的计算效率。

在文中提到，作者使用 vmap 操作符来实现对多个订单簿的并行处理，以提高订单簿匹配系统的性能。这种并行处理方式可以有效地利用 GPU 的并行计算能力，加速订单匹配过程，从而提高整体系统的效率和性能，具体来说：

Table 4列出了在使用vmap并行处理多个订单簿时，处理不同类型的订单所需的时间。这些时间是在JAX-LOB系统中测量的，并与两个CPU实现进行了比较。结果表明，JAX-LOB系统在处理每个订单时的时间比其他两个系统更短，尤其是在处理大量订单时。

Table 5列出了在处理单个订单时，不同系统所需的时间。这些时间是在处理单个订单时测量的，并与JAX-LOB系统和两个CPU实现进行了比较。结果表明，JAX-LOB系统在处理单个订单时的时间比其他两个系统更短，尤其是在处理大量订单时。

这些表格的结果表明，JAX-LOB系统在处理订单簿时具有更高的效率和性能，尤其是在处理大量订单时。这对于需要高效处理大量订单的金融交易系统非常重要。

总结

我们可以看到与CPU实现相比，使用JAX-LOB系统可以获得至少5倍的加速效果。当使用该环境来训练一个RL代理时，相对于CPU实现，甚至可以看到7倍的加速效果。这种并行化带来的加速效果有望为将RL应用于需要反应性订单簿模拟器的高频交易和执行问题的研究做出贡献。

参考来自

标题：AX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading

作者：Sascha Frey、Silvia Sapora、Jakob Foerster、Anisoara Calinescu、Chris Lu、Stefan Zohren、Kang Li、Peer Nagy