TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

LLMを活用した深層学習モデルの改善

前回の記事で、LLMを使用して数学的発見を行うFunSearchの論文を紹介した。

FunSearchは、LLMを使用してプログラムの変更を行い、進化的アルゴリズムでスコアの高いプログラムを選別することで、最適な解を出力するプログラムを生成する。

この仕組みは、深層学習のモデル構造の改良にも使えないかと思い試してみた。

進化的アルゴリズムを行うには評価に時間がかかるため、今回はLLMにより、モデル構造を変更する部分を試した。

対象

dlshogiで使用しているResNetブロックを初期バージョンとして、改善したバージョンをLLMに出力させる。

class ResNetBlock(nn.Module):
    def __init__(self, channels, activation):
        super(ResNetBlock, self).__init__()
        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)
        self.act = activation

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.act(out)

        out = self.conv2(out)
        out = self.bn2(out)

        return self.act(out + x)

LLM

論文と同じCodey for Code Generationと、ChatGPT4を使用する。

プロンプト

import torch
import torch.nn as nn

class ResNetBlock_v0(nn.Module):
  """ResNetBlock that does not change input and output size"""
  def __init__(self, channels, activation):
    super(ResNetBlock_v0, self).__init__()
    self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False)
    self.bn1 = nn.BatchNorm2d(channels)
    self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False)
    self.bn2 = nn.BatchNorm2d(channels)
    self.act = activation

  def forward(self, x):
    out = self.conv1(x)
    out = self.bn1(out)
    out = self.act(out)

    out = self.conv2(out)
    out = self.bn2(out)

    return self.act(out + x)

class ResNetBlock_v1(nn.Module):
  """A version of `ResNetBlock_v0` based on deep learning knowledge and using appropriate modules in PyTorch to change model structure and parameters to improve accuracy."""
  def __init__(self, channels, activation):

始めResNetBlock_v1のコメント部分を論文と同じように「Improved version of `ResNetBlock_v0`」としたが、出力されるコードがほとんど変わらなかったので、「深層学習の知識を元にPyTorchの適切なモジュールを使用してモデル構造とパラメータを変更して精度を向上」という指示を加えた。

結果

Codey for Code Generation

温度パラメータを1にして何回か出力したが、書き方は変わっているがモデル構造としてはほとんど変わらないものしか出力されなかった。

class ResNetBlock_v1(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1):
        super(ResNetBlock_v1, self).__init__()

        self.residual_function = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels)
        )

        # If input and output channels are different, then we need a 1x1 convolution to match the dimensions.
        self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False) if stride != 1 or in_channels != out_channels else nn.Identity()

    def forward(self, x):
        return F.relu(self.residual_function(x) + self.shortcut(x))
ChatGPT4

ChatGPT4では、DropoutやDilated Convolution、SqueezeNetを使うコードが出力された。

class ResNetBlock_v1(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None, activation=nn.ReLU()):
        super(ResNetBlock_v1, self).__init__()
        self.expansion = 4
        mid_channels = out_channels // self.expansion

        # Bottleneck architecture
        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)

        # Dilated Convolution can be added here with dilation > 1 if needed
        self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(mid_channels)

        self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        
        self.se_block = SqueezeExcitationBlock(out_channels)  # Squeeze-and-Excitation block
        self.relu = activation
        self.downsample = downsample
        self.stride = stride
        self.dropout = nn.Dropout(0.5)  # Dropout

    def forward(self, x):
        residual = x

        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))

        out = self.se_block(out)  # Squeeze-and-Excitation block

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.dropout(out)  # Dropout
        out = self.relu(out)

        return out

class SqueezeExcitationBlock(nn.Module):
    def __init__(self, channels, reduction=16):
        super(SqueezeExcitationBlock, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction, bias=False),
            nn.ReLU(),
            nn.Linear(channels // reduction, channels, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

他のバージョンも出力するように指示すると、Grouped Convolutionを使用したバージョンが出力された。

import torch
import torch.nn as nn
import torch.nn.functional as F

class SEBlock(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SEBlock, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

class ResNetBlock_v2(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, groups=1, dilation=1, use_se=False):
        super(ResNetBlock_v2, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=dilation, groups=groups, dilation=dilation, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=dilation, groups=groups, dilation=dilation, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)

        self.se = SEBlock(out_channels) if use_se else None

        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        if self.se is not None:
            out = self.se(out)
        out += self.shortcut(x)
        out = self.relu(out)

        return out

考察

Codey for Code Generationでは、記述方法が変わるだけでモデル構造が変わるコードは出力されなかった。
Codey for Code Generationを使用して、モデル構造を改善するのは難しそうである。

ChatGPT4の方は、深層学習の一般知識を使用して、Dilated Convolution、Grouped Convolution、SqueezeNetといったResNetに効果のある手法を適用したコードが出力された。
出力されたコードを何も考えずに評価して、改善されたものをさらに改善していくことで、良いモデル構造が見つけられそうである。

今回は、ResNetブロックだけに適用したが、ボディ部分の構造にも適用できそうである。

まとめ

LLMでResNetのモデル構造を改善できるか試してみた。
結果、Codey for Code Generationでは良い結果が得られなかったが、ChatGPT4では深層学習の知識を活かしたモデル構造のバリエーションを出力できることがわかった。

実際に将棋AIのモデルで精度が改善されるか時間があれば検証してみたい。