👻

【作品紹介】Common Lispで深層学習フレームワークを0から作ってる話

2023/08/18に公開

導入

昨今の行列演算や数理最適化といったデータサイエンスの分野において、主要なライブラリやフレームワークはPythonやJuliaなど、広く人気を博していて尚且つ魅力的な言語に集中しています。ですが一昔(約40年)前は人工知能といえばLispだったのをご存知ですか？当時は記号論理学を用いた人工知能の研究が盛んで、LispやPrologはまさにそういった用途のための言語だったのです。しかし時代が進むにつれてそういった研究は少なくなり、代わりに大規模な行列などを簡単に取り扱える言語たちが主流になっていったというのがざっくりした今日までの歴史です。しかし記号理論学の影響を無しに考えてもCommon Lispを用いてデータサイエンスを行うメリットは大いにあると個人的に感じます。実は私はいわゆるLisperと言われる人種で普段から専らCommon Lispを用いて開発をしています。そんな私が考える今日のデータサイエンスにCommon Lispを用いることの利点は以下の通りです。

メタプログラミングの有用性
- モデルの定義や計算ノードの定義のためだけに独自の構文を追加できる
- ニューラルネットワークを設計する専用のDSLを簡単に拡張できる
継続
- ✨CommonLispは継続をとても良くサポートしている。(Conditional System)
- 継続があれば、例えば重い計算の最後らへんでエラーが発生したときに、その地点の状態からデバッグをしたりできるのでとても便利
REPLの有用性
- Jupyterなどでやってるトライアンドエラーの開発方式をREPLでできる
- 言語とREPLの相性がそもそも良い
優秀な処理系がある
- オープンソースのSBCLというCommonLispの処理系（Common Lispは言語というより仕様で、それを実装する処理系がいくつか存在します）がとても良い。小数の取り扱いなど数値計算に適性があって、おまけにAVX2までの一部関数は自動ベクトル化などもついてきてる
- その他有償処理系(Allegro, LispWorks etc...)もある。

布教活動と前置きはここら辺にして、早速作ったものを紹介させてください。(正直こういうのに関して僕は素人なので技術的にめちゃくちゃかもしれませんし設計は怪文書と紙一重かもしれないです・・・)

cl-waffe2

リポジトリ
https://github.com/hikettei/cl-waffe2

ドキュメント
https://hikettei.github.io/cl-waffe2/

設計を考えるにあたって、PyTorchもどきをCommon Lispでやるのもなんか違うな〜と感じていたので個人的に論文だとか他のフレームワークの実装を研究していました。(その際、Redditのこの投稿についたコメントが非常に助けになりました)

特徴

データ型: AbstractTensor
- Pythonであればnumpy、JuliaであればAbstractArraysみたいな存在のライブラリ/データ型があればよかったのですがCommon Lispだと要件に合うものがいまいち見つからなかったので行列演算ライブラリから自作しました。これはソースコードに同梱されています。
- Numpy等と同じくメモリ上では一次元で、strideを用いてN次元に拡張します。
フロントエンドとバックエンドの分離
- ✅ 将来的にCUDA/Metal等他のデバイスへの拡張がしやすいようにフロントエンドとバックエンドは分離してあります
- AbstractNodeというCLOSクラスで表される計算ノードは、一つのAbstractTensorにつき一つの実装を提供することができます。
遅延評価ベース
- A + Bという計算を実行しても演算は実行されず、proceed関数を用いるかbuild関数でコンパイルをして初めて実行することができます。
- コンパイルのオーバーヘッドが少ないインタプリタモードさまざまな最適化手法を適用するcl-waffe2 VM C言語などへのJITを介する実行形態の三種類が存在します
- 仮想マシンは、コンパイルされた計算ノードのデータ構造waffe2 IRのリストを処理することに特化します。
Optional Broadcasting
- Numpy SemanticのBroadcasting適用のルール(同じ次元にある1はもう一方のサイズになるまでRepeatするもしこれから演算する行列同士のランクが一致しない場合、小さい方は先頭に1を足す)はPitman's Two-Bit Ruleという教えに反するが故、一部では不賛成の意見も見られます。
- 少ない記述量でfor文を書くのと同等のコードを書けるため、とてもありがたい機能ですが、初心者の方には直感的でなかったり、バグの原因になりやすいと言った問題もあります。
- ✨cl-waffe2では少ない記述量から任意の位置でBroadcastingするAPIを用意しています
All in Common Lisp
- 実装はなるべく全てCommon Lispで完結（SIMD命令が使いたくてCで書いた共用ライブラリをCFFIで呼び出したりしていますが）しています。
- 無駄な依存を増やしてインストールがめんどくさくなるのが個人的に好きではないので、最小構成でインストールする場合の依存を極力減らしました。
- ✨ANSI Common Lisp Portable (現在はSBCL推奨ですがそのうちCCLとかでもテストします)
- REPL-Friendly

実践: 計算ノードを組んでみる

~~行列演算ライブラリを実装する方で力尽きたので(言い訳っていいわけぇ？)~~
執筆時点(2023/08/18)では実に最低限のモデルしか実装されてないです。

回帰モデル
- LinearLayer
画像処理
- Conv2D
- AvgPooling2D MaxPooling2D
- Unfold
活性化関数
- ReLU sigmoid tanh softmax
誤差関数
- L1/L2Norm MSE CrossEntropy ...
最適化関数
- SGD Adam

cl-waffe2では(モデル名引数...)でクラスを初期化してcallで順伝播します。複数のノードやモデルを合成するときはcall->かdefsequenceを使うと完結になります。asnodeマクロを用いることで関数を合成可能な計算ノードとして取り扱うこともできるので、例えばシンプルなCNNはこんな感じで完結に記述できます：

(defsequence CNN (&key
            (out-channels1 4)
	    (out-channels2 16))
    (Conv2D 1 out-channels1 `(3 3))
    (asnode #'!relu)     
    (MaxPool2D `(2 2))
    (Conv2D out-channels1 out-channels2 `(5 5))
    (asnode #'!relu)
    (MaxPool2D `(2 2))
    (asnode #'!reshape t (* 16 4 4)) 
    (LinearLayer (* 16 4 4) 10))
    
(defsequence MLP (in-features hidden-dim out-features
               &key (activation #'!tanh))
         "3 Layers MLP"
         (LinearLayer in-features hidden-dim)
         (asnode activation)
         (LinearLayer hidden-dim hidden-dim)
         (asnode activation)
         (LinearLayer hidden-dim out-features)
         (asnode #'!softmax))

(MLP 784 512 256) ;; 初期化すると・・・

<Composite: MLP{W356455}(
    <<6 Layers Sequence>>

[1/6]          ↓ 
<Composite: LINEARLAYER{W356445}(
    <Input : ((~ BATCH-SIZE 784)) -> Output: ((~ BATCH-SIZE 512))>

    WEIGHTS -> (512 784)
    BIAS    -> (512)
)>
[2/6]          ↓ 
<Composite: ENCAPSULATED-NODE{W356443}(
    #<FUNCTION !TANH>
)>
[3/6]          ↓ 
<Composite: LINEARLAYER{W356433}(
    <Input : ((~ BATCH-SIZE 512)) -> Output: ((~ BATCH-SIZE 512))>

    WEIGHTS -> (512 512)
    BIAS    -> (512)
)>
[4/6]          ↓ 
<Composite: ENCAPSULATED-NODE{W356431}(
    #<FUNCTION !TANH>
)>
[5/6]          ↓ 
<Composite: LINEARLAYER{W356421}(
    <Input : ((~ BATCH-SIZE 512)) -> Output: ((~ BATCH-SIZE 256))>

    WEIGHTS -> (256 512)
    BIAS    -> (256)
)>
[6/6]          ↓ 
<Composite: ENCAPSULATED-NODE{W356419}(
    #<FUNCTION !SOFTMAX>
)>)>

このような計算ノードやその集合を組み合わせて遅延評価しながらネットワークを構築していき、ある程度計算ノードが固まったり終端にたどり着いたら、proceedかbuildでコンパイル+実行をします。

(call (CNN) (randn `(10 1 28 28)))
{CPUTENSOR[float] :shape (10 10) :named ChainTMP360489 
  :vec-state [maybe-not-computed]
  <<Not-Embodied (10 10) Tensor>>
  :facet :input
  :requires-grad NIL
  :backward <Node: ADDNODE-CPUTENSOR (A[~] B[~] -> A[~])>}
  
(!sum *)
{CPUTENSOR[float] :shape (1 1) -> :view (<(BROADCAST 1)> <(BROADCAST 1)>) -> :visible-shape (1 1) :named ChainTMP362248 
  :vec-state [maybe-not-computed]
  <<Not-Embodied (1 1) Tensor>>
  :facet :input
  :requires-grad NIL
  :backward <Node: VIEWTENSORNODE-T (A[RESULT] B[BEFORE] -> A[RESULT])>}
  
(proceed *)
{CPUTENSOR[float] :shape (1 1) -> :view (<(BROADCAST 1)> <(BROADCAST 1)>) -> :visible-shape (1 1) :named ChainTMP362499 
  :vec-state [computed]
  ((-1.8168548))
  :facet :input
  :requires-grad NIL
  :backward <Node: PROCEEDNODE-T (A[~] -> A[~])>}

build関数は仮想マシンが取り扱うIRを介して演算を実行します。少し踏み込んだ話をすると：内部的には一度組んだAbstractNodeの計算木をトポロジカルソートしてからグラフレベルの最適化(In-place mutation, fuse opsなど・・・)を施して逆伝播のコードを生成・・・と言った手順を踏んでいます。disassemble-waffe2-ir関数を用いてその様子を見てみると：

CL-WAFFE2-REPL> (let ((a (parameter (randn `(3 3)))))
		  (cl-waffe2/vm:disassemble-waffe2-ir (!sin a)))

== [disassemble-waffe2-ir: Forward] ======
<WfInst[Compiled: MoveTensorNode(SAVE_FOR_BACKWARD)] : TID365562 <= op(TID365562(3 3) <Param>TID365546(3 3))>
<WfInst[Compiled: MOVETENSORNODE-CPUTENSOR]          : TID365551 <= op(TID365551(3 3) <Param>TID365546(3 3))>
<WfInst[Compiled: SINNODE-LISPTENSOR]                : TID365551 <= op(TID365562(3 3) TID365551(3 3))>

3 Instructions | 3 Tensors


== [disassemble-waffe2-ir: Backward] ======
<WfInst[Compiled: Block -> SINNODE-LISPTENSOR-BACKWARD {
        <WfInst[Compiled: MOVETENSORNODE-CPUTENSOR] : TID365693 <= op(TID365693(3 3) TID365673(3 3))>
        <WfInst[Compiled: MOVETENSORNODE-CPUTENSOR] : TID365676 <= op(TID365676(3 3) <Param>TID365546(3 3))>
        <WfInst[Compiled: COSNODE-LISPTENSOR]       : TID365676 <= op(<Param>TID365546(3 3) TID365676(3 3))>
        <WfInst[Compiled: MULNODE-CPUTENSOR]        : TID365693 <= op(TID365693(3 3) TID365676(3 3))>
        <WfInst[Compiled: MOVETENSORNODE-CPUTENSOR] : TID365712 <= op(TID365712(3 3) TID365693(3 3))>
    }
  ] : TID365629 <= op(TID365584(3 3))>
<WfInst[Compiled: Block -> MOVETENSORNODE-CPUTENSOR-BACKWARD {
        <WfInst[Compiled: MOVETENSORNODE-CPUTENSOR] : TID365662 <= op(TID365662(3 3) TID365659(3 3))>
    }
  ] : TID365640 <= op(TID365629(3 3))>
<WfInst[Compiled: ADDNODE-CPUTENSOR] : TID365548 <= op(TID365548(3 3) TID365640(3 3))>

3 Instructions | 4 Tensors

このような具合になっています。計算ノード名の後のCPUTENSORというのは実行するデバイスの名前で、*using-backend*というパラメーターに優先順位と共に記録されています。Common Lispは動的スコープの言語なので、計算ノードの一部分だけのデバイスを変えるといった処理も大規模な変更なしに行えます。例えば・・・

(with-devices (JITCPUTensor LispTensor) ;; Cにコンパイルしようとする -> 実装がないならLispで実行
    (!sum (ax+b `(3 3) 0 1)))
    
(with-devices (CPUTensor) ;; OpenBLASやSIMD化された命令を組み合わせて実行
    (!sum (ax+b `(3 3) 0 1)))

のように局所的にデバイスを切り替えながら実行する設計となっています。

苦労話: 確率分布のサンプリングとか

機械学習から深層学習でパラメーターの初期値を設定するために必須な確率分布のサンプリングアルゴリズムは一部を除いて（ガンマ分布の近似に関してはめっちゃやること膨大だったので退避した）自力で実装してみました。

ガウス分布randn → Ziggurat法
ベータ分布beta → これ

...

ここらへんの苦労話はまた機会があったら別の記事で書こうかと思います。

Optional Broadcastingの使い方

さて、冒頭で触れたOptional Broadcastingのルールについて紹介させていただきます。

numpyで以下の計算と同様のコードを実行したならばBroadcastingのルールに則ってうまく列ごとに加算されるでしょうが、cl-waffe2ではShape-Errorとなります。

(!add (randn `(3 3)) (randn `(3 1)))

cl-waffe2でBroadcastingが可能となるのは以下の条件を満たしたときです。

1. 計算ノード側での指定
- 計算ノード側でその範囲がBroadcastingされてもOKと明記されているとき
- 例えばAddNodeはA[~] B[~] -> A[~]と宣言されているが、このとき~に対応する範囲はBroadcastingができる。
1. 呼び出す側での指定
- (!view tensor (:broadcast 回数) ...)`と明示的にBroadcastingする
- ↑の関数をもとに自動でBroadcastingする条件がいくつかある

関数!flexibleは対応する位置に<1 x N>というShapeを追加します。このShapeは以下のルールを満たします：

その場所に何回1が追加されても良い。（無限回でも）
追加された1は対応する数までRepeatされても良い。

(print (!flexible (ax+b `(3 3) 1 0) :at 0))

;;{CPUTENSOR[float] :shape (<1 x N> 3 3) :named ChainTMP5932 
;;  :vec-state [maybe-not-computed]
;;  ((0.0 1.0 2.0)
;;   (3.0 4.0 5.0)
;;   (6.0 7.0 8.0))
;;  :facet :input
;;  :requires-grad NIL
;; :backward <Node: FLEXIBLE-RANK-NODE-T (A[~] -> A[~])>}

(print (!flexible (ax+b `(3 3) 1 0) :at 1))


;;{CPUTENSOR[float] :shape (3 <1 x N> 3) :named ChainTMP5932 
;;  :vec-state [maybe-not-computed]
;;  ((0.0 1.0 2.0)
;;   (3.0 4.0 5.0)
;;   (6.0 7.0 8.0))
;;  :facet :input
;;  :requires-grad NIL
;; :backward <Node: FLEXIBLE-RANK-NODE-T (A[~] -> A[~])>}

(print (!flexible (ax+b `(3 3) 1 0) :at 2))

;;{CPUTENSOR[float] :shape (3 3 <1 x N>) :named ChainTMP5932 
;;  :vec-state [maybe-not-computed]
;;  ((0.0 1.0 2.0)
;;   (3.0 4.0 5.0)
;;   (6.0 7.0 8.0))
;;  :facet :input
;;  :requires-grad NIL
;; :backward <Node: FLEXIBLE-RANK-NODE-T (A[~] -> A[~])>}

そのため、BroadcastingしたいTensorがあったら対応する場所に1を挿入して初期化するのではなく、Broadcastingしたい箇所になって初めてその箇所を明記します：

(proceed
    (!add (ax+b `(3 3) 0 0)
          (!flexible (ax+b `(3) 1 0) :at -1)))

;;{CPUTENSOR[float] :shape (3 <1 x N>) :named ChainTMP5999 
;;  :vec-state [maybe-not-computed]
;;  (0.0 1.0 2.0)
;;  :facet :input
;;  :requires-grad NIL
;;  :backward <Node: FLEXIBLE-RANK-NODE-T (A[~] -> A[~])>} 
;;{CPUTENSOR[float] :shape (3 3) :named ChainTMP6070 
;;  :vec-state [computed]
;;  ((0.0 0.0 0.0)
;;   (1.0 1.0 1.0)
;;   (2.0 2.0 2.0))
;;  :facet :input
;;  :requires-grad NIL
;;:backward <Node: PROCEEDNODE-T (A[~] -> A[~])>}

;; 可読性が一番高い書き方
(proceed
    (!add (ax+b `(3 3) 0 0)
          (%transform (ax+b `(3) 1 0)[i] -> [~ i])))

;;{CPUTENSOR[float] :shape (<1 x N> 3) :named ChainTMP6086 
;;  :vec-state [maybe-not-computed]
;;  (0.0 1.0 2.0)
;;  :facet :input
;;  :requires-grad NIL
;;  :backward <Node: FLEXIBLE-RANK-NODE-T (A[~] -> A[~])>} 
;;{CPUTENSOR[float] :shape (3 3) :named ChainTMP6153 
;;  :vec-state [computed]
;;  ((0.0 1.0 2.0)
;;   (0.0 1.0 2.0)
;;   (0.0 1.0 2.0))
;;  :facet :input
;;  :requires-grad NIL
;;  :backward <Node: PROCEEDNODE-T (A[~] -> A[~])>}

!view !reshape !permuteなどのAPIは基本的に(関数名対象のテンソル引数1 引数2 引数3 ...)を取りますが、引数1に関数を指定しそれに引数2 引数3...を自動生成させるみたいな使い方ができるので、Broadcastingと言った用途であれば:

(let ((a (ax+b `(3 3) 0 0))
      (b (ax+b `(3 1) 1 0)))
  (print
   (proceed
    (!add
     a
     (!view b (broadcast-to a))))))

;;{CPUTENSOR[float] :shape (3 3) :named ChainTMP4505 
;;  :vec-state [computed]
;;  ((0.0 0.0 0.0)
;;   (1.0 1.0 1.0)
;;   (2.0 2.0 2.0))
;;  :facet :input
;;  :requires-grad NIL
;;  :backward <Node: PROCEEDNODE-T (A[~] -> A[~])>}

だけでも完結できちゃいます

その他: チュートリアルやサンプルコードなど

日本語でコメントが書いてあるチュートリアルを用意してあるので是非お手持ちの環境で試してみてください。↓

MNIST学習のサンプルパッケージ

今後のゴール

einsum
- ここら辺のアルゴリズムでいい感じに計算量を最適化してくれるeinsumを実装したい
Control Flow
- RNNを実装するためにはIfNode, MapNodeあたりを実装する必要がある
CPU動作の高速化に関する課題
- メモリプールのアルゴリズムをもっと賢くする。あとひょっとしたらメモリリーク起こしてる気がする。
- データをSIMDのレジスタに載せた上でlparallelで並列化をする
- SLEEFで数学関数等のSIMD化をしっかりしておく。
- AVX512の拡張命令を有効化する。~~どこのM1チップの因縁なのかは知らないが~~IntelのエミュレーターがmacOSに対応してくれないので困る。
- 検索ベースで命令をFuseしたりだとかの最適化が弱いのでコンパイラにそれを搭載する
- Sparseなgemmの実装/高速化
そもそもフレームワークとして
- cl-cudaでCUDA対応をさせる
- 学習したモデルの重み保存とか
- 汎用性のためにONNX互換に対応させる。（ノードを1から全部定義し直せば今からでも実装できると思う）

究極的にはGPT2あたりのCPU推論とかが動けば面白いな〜と思って最適化してます。暇があったらちまちまcommitしようかなとか考えてる。

最後に

夏休みに躁になってすごい勢いでコード書きました。今になってみると楽しかったけど虚無感が勝ります

バックエンドの拡張とかは書ききれなかったのでいつか書こうと思います。

拙い文章でしたが最後まで読んでいただきありがとうございました！

参考文献

謝辞

・All comments on this Reddit post: Does anyone have any interest in my deep-learning framework?.

digikar99 for giving me intriguing perspectives on some semantics and the publication of a large number of valuable references.
- Features of Common Lispy approach to (Better) Numpy
Rationalized APL
Previous research of Petalisp
- https://github.com/marcoheisig/Petalisp/tree/master
Automatic differentiation in machine learning: a survey
Some of the algorithms implemented within the source code are referenced below:
- https://arxiv.org/pdf/1201.6035.pdf
- https://github.com/melisgl/mgl-mat
- https://dl.acm.org/doi/pdf/10.1145/359460.359482
- https://andantesoft.hatenablog.com/entry/2023/04/30/183032
- Marsaglia, G., & Tsang, W. W. (2000). The ziggurat method for generating random variables. Journal of statistical software.
- https://marui.hatenablog.com/entry/2023/01/23/194507
- https://arxiv.org/abs/1912.01703
Previous works of JIT Compiler for Deep Learning:
- https://arxiv.org/pdf/2002.03794.pdf
- See also my reading list: https://github.com/hikettei/cl-waffe2/issues/47