Implementation of Weights in 'Understanding Imbalanced Semantic Segmentation Through Neural Collapse' : 가중치 구현 해결!

'Understanding Imbalanced Semantic Segmentation Through Neural Collapse'의 code가 아직 공개되지 않아서, 구현을 시도하던 중, 다음 수식을 구현하는데에서 문제가 있었다.

다른 방정식을 통한 풀이가 있을까 싶어서 찾다가, 딱히 방법이 없다싶어서 다음과 같이 optimization problem으로 구성해서 해결을 시도했다.

import torch
import numpy as np


def objective_function(W, num_classes=20):
    """
    Define the objective function.
    The objective is to make the dot product of columns of W with themselves to be 1,
    and with other columns to be -1/(num_classes - 1).
    you should multiply alpha**2 (scaling factor) when you import weight in model class.
    """
    loss = 0
    for i in range(num_classes):
        # Dot product of the column with itself
        loss += (np.dot(W[:, i], W[:, i]) - 1) ** 2
        for j in range(num_classes):
            if i != j:
                # Dot product with other columns
                loss += (np.dot(W[:, i], W[:, j]) + 1/(num_classes - 1)) ** 2
    return loss

def gradient_descent(W_init, learning_rate, thres, num_classes=20):
    """
    Perform gradient descent to minimize the objective function.
    """
    W = W_init.copy()
    it = 0
    loss = 99999
    while loss > thres:
        grad = np.zeros_like(W)
        for i in range(num_classes):
            # Gradient with respect to the column itself
            grad[:, i] += 2 * (np.dot(W[:, i], W[:, i]) - 1) * W[:, i]
            for j in range(num_classes):
                if i != j:
                    # Gradient with respect to other columns
                    grad[:, i] += 2 * (np.dot(W[:, i], W[:, j]) + 1/(num_classes - 1)) * W[:, j]
        # Update the weights
        W -= learning_rate * grad

        it += 1
        # Print the loss every 100 iterations
        if it % 100 == 0:
            loss = objective_function(W, num_classes)
            print(f"Iteration {it}: Loss = {loss}")
    return W

# Initialize the weight matrix
np.random.seed(0)  # For reproducibility
num_features, num_classes = 512, 20
W_init = np.random.randn(num_features, num_classes)

# Run gradient descent
learning_rate = 0.001
thres = 1e-6
W_optimized = gradient_descent(W_init, learning_rate, thres)
print(W_optimized)
W_optimized_tensor = torch.tensor(W_optimized)
torch.save(W_optimized_tensor, './W_optimized.pth')

# Evaluate the final weight matrix
final_loss = objective_function(W_optimized)
print(final_loss)

---

(23.12.22 추가)

논문을 더 읽고 찾아보니, reference가 되어있던 이 논문에서 Equiangular Basis Vector를 구현하는 방법에 대해서 작성하였다.

코드로 표현하면 다음과 같다.

num_classes = 100
dim = 512

weight = (torch.sqrt(torch.tensor(num_classes / (num_classes - 1))) *
          (torch.eye(num_classes) - (1 / num_classes) *
           torch.ones((num_classes, num_classes))))
weight /= torch.sqrt((1 / num_classes * torch.norm(weight, 'fro') ** 2))
weight = torch.mm(weight, torch.eye(num_classes, dim))