かるまるラボ
/ Boolean Attention 2026.05
// THEORY · DLGN & HBA & DISTILL

なぜ論理ゲートを微分可能にするのか

通常のニューラルネットは隠れ層を y = σ(Wx + b) のような 実数の線形結合 + 非線形活性化 で構築し、誤差逆伝播法で学習します。これは 連続で微分可能 であることが前提です。

一方、論理ゲート (AND, OR, XOR, …) は 離散関数 で、入力をビットに丸めた瞬間に勾配はゼロまたは未定義になります。これでは学習できません。

そこで 微分可能化 (differentiable relaxation) を行います。

DLGN — 16 ゲートを softmax で混ぜる

2 入力ブール関数は全部で $2^{2^2} = 16$ 種類しかありません。DLGN (Differentiable Logic Gate Network) は、各ゲート位置に対して 16 種類のゲートを softmax で混合した結果 を出力します。

// gate-output @ position i
$$y_i \;=\; \sum_{k=0}^{15} \alpha_{i,k}\, g_k(a_i, b_i), \qquad \alpha_{i,k} \;=\; \frac{\exp(\theta_{i,k})}{\sum_j \exp(\theta_{i,j})}$$

ここで $a_i, b_i$ は入力ビット (学習中は $[0,1]$ の実数として扱う)、$g_k$ は 16 種類の論理ゲート関数、$\theta_{i,k}$ は学習対象の logit、$\alpha_{i,k}$ は softmax された混合係数です。学習中は $\alpha$ が連続値なので勾配が流れます。

学習が終わったら、各位置で argmax で 1 個のゲートに確定 (hard collapse) し、純粋なブーリアン回路に戻します。

16 ゲートの全集合

GATES = [
    lambda a, b: 0,                # 0: const False
    lambda a, b: a * b,            # 1: AND
    lambda a, b: a * (1 - b),      # 2: A AND NOT B
    lambda a, b: a,                # 3: A
    lambda a, b: (1 - a) * b,      # 4: NOT A AND B
    lambda a, b: b,                # 5: B
    lambda a, b: a + b - 2*a*b,    # 6: XOR
    lambda a, b: a + b - a*b,      # 7: OR
    lambda a, b: 1 - (a + b - a*b),       # 8: NOR
    lambda a, b: 1 - (a + b - 2*a*b),     # 9: XNOR
    lambda a, b: 1 - b,            # 10: NOT B
    lambda a, b: a + (1-b) - a*(1-b),     # 11: A OR NOT B
    lambda a, b: 1 - a,            # 12: NOT A
    lambda a, b: (1-a) + b - (1-a)*b,     # 13: NOT A OR B
    lambda a, b: 1 - a*b,          # 14: NAND
    lambda a, b: 1,                # 15: const True
]

なぜ LoopedDLGN は失敗したか

Universal Transformer に倣い、同じ DLGN ブロックを何度も繰り返す設計を試しました。バナッハの contraction mapping theorem に基づき、反復するうちに $x^{(t+1)} = f(x^{(t)}, \mathrm{input})$ が固定点 $x^{*}$ に収束することを期待しました。

しかしハードコラプス後、各反復で生じる量子化誤差が爆発的に蓄積します:

// quantization error accumulation
$$\epsilon_{\text{total}} \;\gtrsim\; \sum_{t=1}^{T} \,\bigl\| f_{\mathrm{hard}}(x^{(t)}) - f_{\mathrm{soft}}(x^{(t)}) \bigr\|$$

反復が深くなるほど、ソフト時の学習軌跡からの逸脱が無視できなくなり、Hard PPL が 754 まで暴騰しました。

教訓

Boolean 量子化誤差は、深さに対して線形以上で増える。反復系アーキテクチャは Boolean とは相性が悪い、というのが学びです。

HBA — Boolean を「どこに使うか」を再考

Transformer の Attention は、抽象的に次の 2 ステップに分解できます:

  1. ルーティング: クエリ $Q$ とキー $K$ から、どの位置を見るかを決める (softmax 注意重み)
  2. 値の集約: 注意重みでバリュー $V$ を加重平均

HBA は、この (1) ルーターだけを Boolean 化 し、(2) は float のままにします。

// HBA forward
$$\text{router}_{q,k} \;=\; \mathrm{HardSign}\bigl(\, Q_q^\top W_{\mathrm{router}} K_k \,\bigr) \in \{-1, +1\}$$ $$\text{attn}_{q,k} \;=\; \mathrm{softmax}_k\!\left( \frac{\text{router}_{q,k}}{\tau} \right)$$ $$\text{out}_q \;=\; \sum_k\, \text{attn}_{q,k}\, V_k$$

ポイント:

なぜ HBA は誤差累積しないか

ルーターは離散値ですが、値の集約は float なので量子化誤差が伝播しません。具体的には、ルーターの誤差は attn を通じて 1 ステップ分しか影響せず、深さ方向に累積しないという構造的優位があります。

構造的利点

これが HBA のハードコラプス後 PPL 6.54 (vs LoopedDLGN の 754) の 115 倍の改善 をもたらした本質的な理由です。「Boolean をどこに入れるか」の設計判断が、すべてを決めました。

知識蒸留で Transformer を超える

教師 (Transformer) の出力分布 $p_T(y|x)$ を、生徒 (HBA) の出力分布 $p_S(y|x)$ に温度 $T$ のソフトラベルで模倣させます:

// distillation loss
$$\mathcal{L} \;=\; \alpha \cdot \mathrm{CE}(p_S, y_{\text{true}}) \;+\; (1-\alpha) \cdot T^2 \cdot \mathrm{KL}\!\bigl(\, \sigma(z_T/T) \,\big\|\, \sigma(z_S/T) \,\bigr)$$

ここで $\alpha = 0.3$, $T = 8$, $z_T, z_S$ は教師・生徒の logits、$\sigma$ は softmax です。

重要なのは、教師は学習データで 限られた汎化 をしているが、蒸留時の生徒は教師のソフトな出力分布を受け取るため、実質的にデータオーグメンテーションを受けたかのように学習することです。これが 生徒が教師を逆転する 現象 (born-again networks) の本質で、HBA Student が PPL 4.73 を達成、教師 4.86 を 0.13 上回ったのもこの効果です。

参考文献