[딥러닝] Forward-Forward 알고리즘 구현

Ethan·2023년 4월 22일

Papers Implement

목록 보기

1/1

본 블로그의 모든 글은 직접 공부하고 남기는 기록입니다.
잘못된 내용이나 오류가 있다면 언제든지 댓글 남겨주세요.

Forward-Forward 알고리즘

오늘은 힌튼의 Forward Forward 알고리즘을 구현해보고자 합니다. 논문을 리뷰한 지는 꽤 됐는데 바빠서 손을 못 대고 있다가 이제야 포스팅을 하게 되네요. PyTorch 공식 코드가 있어서 이를 참고해서 작성합니다.

해당 논문 리뷰를 먼저 읽고 오면 이 글을 이해하는 데 도움이 됩니다.

구현 과정

Data setup

MNIST Dataset

MNIST 데이터셋을 사용합니다. 적합한 train/test loader를 만들어줍니다.

def MNIST_loaders(train_batch_size=50000, test_batch_size=10000):

    transform = Compose([
        ToTensor(),
        Normalize((0.1307,), (0.3081,)),
        Lambda(lambda x: torch.flatten(x))])

    train_loader = DataLoader(
        MNIST('./data/', train=True,
              download=True,
              transform=transform),
        batch_size=train_batch_size, shuffle=True)

    test_loader = DataLoader(
        MNIST('./data/', train=False,
              download=True,
              transform=transform),
        batch_size=test_batch_size, shuffle=False)

    return train_loader, test_loader

$(0.1307, 0.3081)$ 은 각각 MNIST dataset의 mean과 std입니다.

overlay_y_on_x()

이 글에서는 지도학습일 경우만 구현합니다.

FF 알고리즘을 사용한 지도학습에는 negative sample이 필요합니다. 논문에서는 이미지의 첫 10 픽셀을 one hot encoded label로 바꿔서 사용했습니다.

def overlay_y_on_x(x, y):
    x_ = x.clone()
    x_[:, :10] *= 0.0   # 첫 10 픽셀을 0으로 초기화
    x_[range(x.shape[0]), y] = x.max()  # y 자리를 non zero 값으로 변경
    return x_

즉, positive sample과 negative sample은 one hot encoded label 정보만 다르고 나머지는 같은 데이터입니다.

Layer Class

FF 알고리즘의 핵심 아이디어는 Pos/Neg pass 모두 feed forward network를 사용한다는 것입니다. 이 때 layer 단위로 최적화를 진행하기 때문에 먼저 각 pass에 사용할 Layer 클래스부터 구현하겠습니다.

\hat y = \sigma(\sum_iy_i^2-\theta)\qquad(1)

논문에서는 식 (1)과 같이 positive sample을 판정하기 위한 SSE 형태의 목적함수를 설계했습니다. 이 글에서도 해당 수식을 사용합니다.

식 (1)은 주어진 벡터 $y_j$ 가 positive sample일 확률을 의미합니다. 즉, positive에 최적화를 시키고 싶다면 positive $y_j > \theta$ 여야 하고, negative $y_j < \theta$ 여야 한다는 점을 기억합시다.

init()

먼저 레이어를 구성하는 변수들을 세팅합니다.

class Layer(torch.nn.Linear):
    def __init__(self, in_features, out_features, bias):
        super().__init__(in_features, out_features, bias)
        self.threshold = 2.0
        self.epochs = 1000
        self.relu = nn.ReLU()
        self.optim = Adam(self.parameters(), lr=0.001)

forward()

이어서 forward 함수를 만들어줍니다. 넘겨받는 입력 벡터의 길이를 normalize 해주어야 합니다.

	def forward(self, x):
        x = x / x.norm(2, 1, keepdim=True) + 1e-4
        x = torch.mm(x, self.weight.T) + self.bias.unsqueeze(0)
        x = self.relu(x)
        return x

train()

이제 epochs만큼 이를 반복하는 train 함수를 만들어줍니다. train 함수는 positive sample과 negative sample을 모두 처리할 수 있어야 하고, epoch마다 loss를 계산해야 합니다.

	def train(self, x_pos, x_neg):
        for i in range(self.epochs):
            g_pos = self.forward(x_pos).pow(2).mean(1)-self.threshold
            g_neg = self.forward(x_neg).pow(2).mean(1)-self.threshold

Loss function

각 레이어에서 사용할 loss function을 설계해 보겠습니다.

기본적으로 FF 알고리즘은 true/false를 구분하는 이진분류 문제입니다. 그런데 위에서 살펴본 것과 같이 sample data를 판정하는 $\theta$ 값에 대한 조건이 존재합니다.

지도학습일 때는 간단하게 해결할 수 있습니다. label이 존재하기 때문이죠. 즉, true sample에 최적화되도록 하면 됩니다. 이진 분류 문제이므로 loss function으로 Binary Cross Entropy를 사용합니다.

\begin{aligned} L&=-(y\log(y_j)+(1-y)\log (1-y_j))\\ &=-y\log(y_j)-(1-y)\log (1-y_j)\qquad(2) \end{aligned}

우리는 positive sample에 최적화하는 것이 목적이므로 BCE를 softplus 함수로 대체할 수 있습니다. 따라서 식 (2)를 아래와 같이 바꿔 쓸 수 있습니다.

\begin{aligned} L_{1=True}&=-\log(y_j)\\ &=\log(1+\exp(-x))\qquad(3) \end{aligned}

지금까지 설계한 loss는 input으로 $x$ 하나만 받습니다. 따라서 $x_{pos},\ x_{neg}$ 두 데이터를 하나로 합쳐서 넣어줄 수 있도록 만들어 줍니다.

\log(1+\exp(-x)) = \log(1+\exp(-(x_{pos}\oplus -x_{neg})))

g_neg에 마이너스를 붙이는 이유는 negative sample에서 멀어지도록 하기 위함입니다.

이제 옵티마이저를 초기화하고, 그래디언트를 계산해서, 파라미터를 업데이트하는 과정을 추가해주면 완성입니다.

    def train(self, x_pos, x_neg):
        for i in range(self.epochs):
            g_pos = self.forward(x_pos).pow(2).mean(1)-self.threshold
            g_neg = self.forward(x_neg).pow(2).mean(1)-self.threshold
            loss = torch.log(1 + torch.exp(torch.cat([-g_pos, g_neg]))).mean()
            self.optim.zero_grad()
            loss.backward()
            self.optim.step()
        return self.forward(x_pos).detach(), self.forward(x_neg).detach()

FF Class

이제 본격적으로 네트워크를 구현합니다. 먼저 레이어 계층을 만듭니다.

class FF(torch.nn.Module):

    def __init__(self, dims):
        super().__init__()
        self.layers = []
        for d in range(len(dims)-1):
            self.layers += [Layer(dims[d], dims[d+1]).cuda()]

이어서 predict 함수를 구현합니다. 논문에서는 각 레이어의 goodness를 가지고 최적화 정도를 평가하게 되어 있습니다. 또한 지도학습이므로 레이블 개수만큼 최적화를 해주어야 합니다.

    def predict(self, x):
        goodness_per_label = []
        for label in range(10):
            h = overlay_y_on_x(x, label)
            goodness = []
            for layer in self.layers:
                h = layer(h)
                goodness += [h.pow(2).mean(1)]
            goodness_per_label += [sum(goodness).unsqueeze(1)]
        goodness_per_label = torch.cat(goodness_per_label, 1)
        return goodness_per_label.argmax(1)

이제 train 함수를 구현합니다. positive sample과 negative sample을 받아서 최적화하는 일을 합니다.

    def train(self, x_pos, x_neg):
        h_pos, h_neg = x_pos, x_neg
        for i, layer in enumerate(self.layers):
            print('training layer', i, '...')
            h_pos, h_neg = layer.train(h_pos, h_neg)

print 함수를 추가하여 self.layers를 확인할 수 있습니다.

[Layer(
   in_features=784, out_features=500, bias=True
   (relu): ReLU()
 ),
 Layer(
   in_features=500, out_features=500, bias=True
   (relu): ReLU()
 )]

모델 학습 및 시각화

이제 모델을 학습시키고, 이를 시각화해보겠습니다.

먼저 시각화하는 함수는 다음과 같습니다.

def visualize_sample(data, name='', idx=0):
    reshaped = data[idx].cpu().reshape(28, 28)
    plt.figure(figsize = (4, 4))
    plt.title(name)
    plt.imshow(reshaped, cmap="gray")
    plt.show()

이어서 모델을 학습하는 함수를 만들어 줍니다. 먼저 MNIST 데이터셋을 로드합니다.

if __name__ == "__main__":
    torch.manual_seed(1234)
    train_loader, test_loader = MNIST_loaders()

FF 네트워크를 구성하고, 데이터셋을 차례대로 투입합니다. 모델이 예측한 결과와 실제 데이터를 비교하여 error score를 얻습니다.

    net = FF([784, 500, 500])
    x, y = next(iter(train_loader))
    x, y = x.cuda(), y.cuda()
    x_pos = overlay_y_on_x(x, y)
    rnd = torch.randperm(x.size(0))
    x_neg = overlay_y_on_x(x, y[rnd])
    
    for data, name in zip([x, x_pos, x_neg], ['orig', 'pos', 'neg']):
        visualize_sample(data, name)
    
    net.train(x_pos, x_neg)

    print('train error:', 1.0 - net.predict(x).eq(y).float().mean().item())

    x_te, y_te = next(iter(test_loader))
    x_te, y_te = x_te.cuda(), y_te.cuda()

    print('test error:', 1.0 - net.predict(x_te).eq(y_te).float().mean().item())

이상으로 구현을 마칩니다. 공식 코드 덕분에 편안하게 작업했네요.

참고문헌

PyTorch - Forward-Forward algorithm (official)

Ethan

재미있게 살고 싶은 대학원생