Squeezeformerの詳細解説

1. 全体構造の概要

Squeezeformerは音声認識などのシーケンス処理のために設計された効率的なアーキテクチャです。入力データ形状は (batch_size, sequence_length, channel_size) となっており、例えば: - batch_size = 2 (2つのデータを同時処理) - sequence_length = 16 (時系列の長さ) - channel_size = 64 (各時点での特徴量の次元)

全体コード

import torch
import torch.nn as nn

class GLUMlp(nn.Module):
    def __init__(self, dim_expand, dim):
        super(GLUMlp, self).__init__()
        self.dense_1 = nn.Linear(dim, dim_expand)
        self.glu = GLU()
        self.dense_2 = nn.Linear(dim_expand // 2, dim)

    def forward(self, x):
        x = self.dense_1(x)
        x = self.glu(x)
        x = self.dense_2(x)
        return x

class ScaleBias(nn.Module):
    def __init__(self, num_features):
        super(ScaleBias, self).__init__()
        self.scale = nn.Parameter(torch.ones(num_features))
        self.bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x):
        return self.scale * x + self.bias

class GLU(nn.Module):
    def forward(self, x):
        a, b = x.chunk(2, dim=-1)
        return a * torch.sigmoid(b)

class ECA(nn.Module):
    def __init__(self, kernel_size=5):
        super(ECA, self).__init__()
        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)

    def forward(self, x):
        y = torch.mean(x, dim=-1, keepdim=True)
        y = self.conv(y.transpose(1, 2))
        y = torch.sigmoid(y).transpose(1, 2)
        return x * y

class Conv1DBlockSqueezeformer(nn.Module):
    def __init__(self, channel_size, kernel_size, dilation_rate=1, 
                 expand_ratio=4, se_ratio=0.25, activation='swish'):
        super(Conv1DBlockSqueezeformer, self).__init__()
        self.expand_ratio = expand_ratio
        self.activation = torch.nn.SiLU()

        expanded_dim = channel_size * expand_ratio
        
        # Linear層でチャンネル数を拡張
        self.expand = nn.Linear(channel_size, expanded_dim)
        
        # Depthwise Convolution
        self.dwconv = nn.Conv1d(expanded_dim, expanded_dim, kernel_size,
                               dilation=dilation_rate,
                               padding=(kernel_size - 1) // 2, 
                               groups=expanded_dim)
        
        self.bn = nn.BatchNorm1d(expanded_dim)
        self.eca = ECA()
        
        # Projection層でチャンネル数を元に戻す
        self.project = nn.Linear(expanded_dim, channel_size)
        
        # FFN部分
        self.ffn = GLUMlp(channel_size * 4, channel_size)
        self.layer_norm_2 = nn.LayerNorm(channel_size)
        self.scale_bias_1 = ScaleBias(channel_size)
        self.scale_bias_2 = ScaleBias(channel_size)

    def forward(self, x):
        # 最初のスキップ接続のために入力を保存
        skip = x
        
        # チャンネル数を拡張
        x = self.expand(x)
        x = self.activation(x)
        
        # Depthwise Convolution用にテンソルの次元を入れ替え
        x = x.permute(0, 2, 1)  # (batch, sequence, channels) -> (batch, channels, sequence)
        x = self.dwconv(x)
        x = self.bn(x)
        x = x.permute(0, 2, 1)  # (batch, channels, sequence) -> (batch, sequence, channels)
        
        # ECA attention
        x = self.eca(x)
        
        # チャンネル数を元に戻す
        x = self.project(x)
        x = self.scale_bias_1(x)
        
        # 最初のスキップ接続
        x = skip + x
        
        # FFN部分
        residual = x
        x = self.ffn(x)
        x = self.scale_bias_2(x)
        x = self.layer_norm_2(x + residual)
        
        return x

2. 各モジュールの詳細解説

2.1 GLU (Gated Linear Unit)

class GLU(nn.Module):
    def forward(self, x):
        a, b = x.chunk(2, dim=-1)
        return a * torch.sigmoid(b)

なぜGLUが必要か？

入力を2つに分割し、一方をゲート（門番）として使用
sigmoid(b)は0〜1の値を出力し、どの特徴をどれだけ通すかを制御
例え: 100個の特徴があるとき、重要な特徴は1に近い値、不要な特徴は0に近い値で乗算

2.2 GLUMlp (Gated Linear Unit MLP)

class GLUMlp(nn.Module):
    def __init__(self, dim_expand, dim):
        self.dense_1 = nn.Linear(dim, dim_expand)
        self.glu = GLU()
        self.dense_2 = nn.Linear(dim_expand // 2, dim)

なぜこの構造か？

dense_1: 特徴空間を拡大（例：64→256）

より豊かな表現を可能に
GLUで半分に分割されることを考慮して2倍の大きさに

GLU: 特徴の選択的な活性化
dense_2: 選択された特徴を元の次元に圧縮

GLUで半分になった次元（256→128）を元の次元（64）に

2.3 ECA (Efficient Channel Attention)

class ECA(nn.Module):
    def __init__(self, kernel_size=5):
        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, 
                             padding=(kernel_size - 1) // 2, bias=False)

なぜECAが必要か？

チャンネル間の相関を効率的に捉える

torch.mean(): 各チャンネルの平均を計算 conv1d: 近隣チャンネル間の関係性を学習 sigmoid: 0〜1の重みに変換

計算効率が良い

従来のSENetと比べてパラメータ数が少ない kernel_size=5 は5つの近傍チャンネルだけを見る

2.4 ScaleBias

class ScaleBias(nn.Module):
    def __init__(self, num_features):
        self.scale = nn.Parameter(torch.ones(num_features))
        self.bias = nn.Parameter(torch.zeros(num_features))

なぜScaleBiasが必要か？

各特徴量のスケールとバイアスを学習可能に
例: ある特徴が小さすぎる/大きすぎる場合に適正な範囲に調整
LayerNormとは異なり、バッチ全体ではなく個別に調整

3. メインの処理フロー (Conv1DBlockSqueezeformer)

3.1 Expansion Phase

skip = x  # 元の入力を保存
x = self.expand(x)  # チャンネル数を拡張
x = self.activation(x)  # 非線形変換

なぜこの順序か？

skip保存: 後でスキップ接続するため
拡張: より豊かな特徴表現のため活性化: 非線形性の導入

3.2 Convolution Phase

x = x.permute(0, 2, 1)  # 次元入れ替え
x = self.dwconv(x)      # Depthwise畳み込み
x = self.bn(x)          # バッチ正規化
x = x.permute(0, 2, 1)  # 次元を戻す

なぜこの順序か？

permute: Conv1dは(batch, channel, sequence)の形状を要求
dwconv: 各チャンネルを独立して処理
bn: 数値の安定化
permute: 元の形状に戻す

3.3 Attention & Projection Phase

x = self.eca(x)         # チャンネル注意機構
x = self.project(x)     # 次元圧縮
x = self.scale_bias_1(x)
x = skip + x            # スキップ接続

なぜこの順序か？

eca: 重要なチャンネルを強調
project: 情報を圧縮
scale_bias: 特徴量の調整
skip接続: 勾配消失を防ぎ、元の情報を保持

3.4 FFN Phase

residual = x
x = self.ffn(x)
x = self.scale_bias_2(x)
x = self.layer_norm_2(x + residual)

なぜこの順序か？

residual保存: 2つ目のスキップ接続用
ffn: 最終的な特徴変換
scale_bias: 特徴量の再調整
layer_norm: 出力の正規化

4. 設計の利点

計算効率

Depthwise Convolutionで計算量削減
ECAで効率的な注意機構を実現

表現力

2段階のスキップ接続で深い層でも学習可能拡張→処理→圧縮の流れで豊かな特徴表現

安定性

BatchNormとLayerNormで数値の安定化 ScaleBiasで細かな調整が可能

アーキテクチャ全体の処理

まず広い空間（拡張）で特徴を見つけやすくする
各特徴を個別に処理（Depthwise Conv）
重要な特徴に注目（Attention）
重要な情報を凝縮（Projection）
元の情報も保持（Skip Connection）
最終的な特徴抽出（FFN）
数値の安定化（LayerNorm）

これは、例えば次のような流れに似ています：

大きな紙に詳細なメモを取る（拡張）
メモの各部分を個別に見直す（Depthwise Conv）
重要な部分に蛍光ペンで印をつける（Attention）
要点をコンパクトにまとめる（Projection）
元のメモも残しておく（Skip Connection）
まとめた内容をさらに整理する（FFN）
全体を見やすく整える（LayerNorm）

各ステップの順序とその意図

広い空間での特徴探索（拡張）: 最初に線形変換でチャンネル数を広げる（expand）ことで、特徴量の表現力を向上させます。この拡張により、後続の処理で複雑な特徴を捕捉しやすくなります。これを最初に行うことで、次に行うDepthwise Convolutionが十分に高次元な特徴を持つデータを扱えるようになります。

局所特徴の抽出（Depthwise Conv）: 拡張後のデータに対してDepthwise Convolutionを適用することで、局所的なパターンを効率的に抽出します。ここでチャンネルごとに独立に畳み込みを行うため、計算効率が良く、かつ空間的な特徴を捉えやすいという利点があります。この時点でDepthwise Convを適用する理由は、チャンネルの拡張によって豊富な特徴表現が可能になっているためです。

重要な特徴への注目（ECA Attention）: Depthwise Convolutionで得られた特徴量に対して、ECAを用いて重要な情報に重みを付けます。この順番でAttentionを導入することで、前段のConv層が捉えた特徴のうち、特に重要な部分を強調します。ここでAttentionを適用しなければ、後続のProjectionやSkip Connectionによる統合の前に、重要な特徴が反映されません。

情報の圧縮（Projection）: 拡張されたチャンネル数を元のチャンネル数に戻します（project）。これにより、再び元の次元に収束させることで、後続の層での計算量やパラメータ数を抑えつつ、前段の処理で得た情報を保持します。このProjectionの前にECAがあるのは、ECAで強調された重要な特徴を適切に維持するためです。

元の情報も保持（Skip Connection）: スキップ接続を用いて、元の入力データを保持しつつ、変換されたデータと統合します。これにより、元の情報を損なうことなく変換後の情報と合成することができ、特に学習の安定性が向上します。ここでスキップ接続を適用しない場合、前段の処理で変換されたデータのみが出力に反映され、学習が不安定になったり、元のデータが持つ情報が失われる可能性があります。

最終的な特徴抽出（FFN）: FFN（Feed-Forward Network）で全体の特徴をさらに統合・整理します。ここではGLUMlpを用いることで、活性化関数と線形変換を組み合わせた処理を行い、複雑な関係をキャプチャします。スキップ接続後にFFNを置く理由は、スキップ接続によって統合されたデータ全体を最終的に整理する必要があるためです。

数値の安定化（LayerNorm）: 最後にLayerNormを適用し、数値のスケールを整えます。これにより、全体の出力が安定し、学習がスムーズになります。LayerNormを最終段に置くことで、各ステップを経た後のデータの分布が整えられ、出力の安定性が確保されます。

趣味のPython・深層学習

中級者のための実装集

Squeezeformerを理解したい

Squeezeformerの詳細解説

1. 全体構造の概要

全体コード

2. 各モジュールの詳細解説

2.1 GLU (Gated Linear Unit)

なぜGLUが必要か？

2.2 GLUMlp (Gated Linear Unit MLP)

なぜこの構造か？

2.3 ECA (Efficient Channel Attention)

なぜECAが必要か？

2.4 ScaleBias

なぜScaleBiasが必要か？

3. メインの処理フロー (Conv1DBlockSqueezeformer)

3.1 Expansion Phase

なぜこの順序か？

3.2 Convolution Phase

なぜこの順序か？

3.3 Attention & Projection Phase

なぜこの順序か？

3.4 FFN Phase

なぜこの順序か？

4. 設計の利点

計算効率

表現力

安定性

アーキテクチャ全体の処理

各ステップの順序とその意図