[Keras] 내가 보려고 쓴 Bayesian Neural Network(BNN) 구현하기

- 이 글은 작성자가 이해한 바 대로 작성되어, 내용이 실제와 다를 수 있습니다.
- 이 글은 다음 사이트를 참조하여 작성되었습니다.

Keras documentation: Probabilistic Bayesian Neural Networks

- 제가 사용하기 편하게 정리한 코드는 다음에 있습니다.

GitHub - MIA-khm/basemodel: base models of machine learning algorithm

1. 개념
베이지안 정리는 다음과 같다.

P(D,H) = P(D|H)P(H)

이 때, D는 데이터를 의미하며, H는 가설을 의미한다.

P(D,H)는 epistemic uncertainty를 의미한다.
P(H)는 prior probability, P(H|D)는 posterior probability로 불린다.

2. 모델의 입출력 값
Bayesian Neural Network(이하 BNN)의 입력값은 다른 모델들과 동일하다.
그러나 출력값은 일반적인 인공신경망과 달리 특정한 값이 아닌 값의 범위이다.

3. 코드 해석
(1) prior

def prior(kernel_size, bias_size, dtype=None):
    n = kernel_size + bias_size
    prior_model = keras.Sequential(
        [
            tfp.layers.DistributionLambda(
                lambda t: tfp.distributions.MultivariateNormalDiag(
                    loc=tf.zeros(n), scale_diag=tf.ones(n)
                )
            )
        ]
    )
    return prior_model

해당 함수는 prior distribution을 함수로 구현한 것이다. 이 함수에서 구현되는 레이어는 n개(kernel size + bias size)의 노드를 가지도록 구현되며, 이 때 kernel size와 bias size는 BNN 모델 구현 시 사용되는 tfp.layers.DenseVariational 함수에서 정의된다. (DenseVariational은 임의의 kernal과 bias로 구성된 Dense 레이어이다.)

prior model은 Sequential로 구성되어 있지만, 단일 레이어로 구성되어 있다. 이 때 레이어는 tfp.layers.DistributionLambda를 활용되어 구성되었으며, 이 함수는 tfp로 정의된 분포를 keras 모델에 연결하는 레이어이다.

이 때 활용되는 분포는 tfp.distributions.MultivariateNormalDiag로 정의되었다. 이 함수는 다변량 정규분포이다. 정규분포의 파라미터는 mean값과 std 값으로 구성되며, mean은 loc으로, std는 scale_diag로 정의할 수 있다. 또한 다변량인 이 값들은 벡터로 구성된 loc과 scale_diag의 개수(size)로 정의된다.

(2) posterior

def posterior(kernel_size, bias_size, dtype=None):
    n = kernel_size + bias_size
    posterior_model = keras.Sequential(
        [
            tfp.layers.VariableLayer(
                tfp.layers.MultivariateNormalTriL.params_size(n), dtype=dtype
            ),
            tfp.layers.MultivariateNormalTriL(n),
        ]
    )
    return posterior_model

이 함수는 특정 데이터 값인 상황에서 가설이 발생할 확률을 표현하게 된다. 이를 위해 사용하게 되는 함수는 prior과 동일하게 kernel_size와 bias_size 합인 n에 의해 제어되며, tfp.layers.VariableLayer와 tfp.layers.MultivariateNormalTriL의 Sequence로 구성된다. 이 중 tfp.layers.VariableLayer은 입력값과 관계 없이 변수를 반환하는 레이어로, 첫번째 파라미터인 tfp.layers.MultivariateNormalTriL.params_size(n)에 의해 shape이 결정된다.
tfp.layers.MultivariateNormalTriL(n)은 n의 형태를 갖는 다변량 정규분포를 만드는 레이어이다. 그리고, tfp.layers.MultivariateNormalTriL.params_size(n)은 분포를 만들 때, 필요한 파라미터의 크기를 반환한다. 결국, tfp.layers.VariableLayer의 사이즈는 tfp.layer.MultivariateNormalTriL(n)의 분포를 만들기 위해 필요한 파라미터를 레이어로 만든 것이다.

(3) 모델 생성

def create_bnn_model(FEATURE_NAMES, hidden_units, train_size, prior, posterior):
    inputs = create_model_inputs(FEATURE_NAMES)
    features = keras.layers.concatenate(list(inputs.values()))
    features = layers.BatchNormalization()(features)

    # Create hidden layers with weight uncertainty using the DenseVariational layer.
    for units in hidden_units:
        features = tfp.layers.DenseVariational(
            units=units,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1 / train_size,
            activation="sigmoid",
        )(features)

    # The output is deterministic: a single point estimate.
    outputs = layers.Dense(units=1)(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

create_model_input으로 만들어지는 입력 레이어는 입력 값의 수 만큼 레이어가 생성된 것이다. 각각 만들어진 레이어를 하나로 합치기 위해 concatenate를 실행하여 입력값으로 활용한다.
hidden_units는 각 레이어 별 unit 수를 값으로 갖는 벡터이다. 예를 들어 [8,4]이면 첫번째 hidden layer는 8개의 unit을 두번째 hidden layer는 4개의 unit을 갖는 것이다. 이를 활용하여 레이어를 순차적으로 생성하며, 이를 위해 for문을 사용하게 된다. 각 노드는 확률분포로 구성되어 있어, kernel과 bias로 구성되게 된다. 이를 표현하기 위해 DenseVariational을 활용한다. DenseVariational은 임의의 kernel과 bias로 구성된 레이어이다. 이 때 Dense의 사이즈는 unit에 의해 결정되며, 이 값들은 베이지안 이론을 기반한 분포에서 결정되므로 prior과 posterior이 사용된다. 앞에 설명했듯 prior와 posterior은 kernel_size와 bias_size로 결정되는데, 이는 DenseVariational에서 확보가 되어 있기 때문에 이를 이용하여 prior과 posterior 함수 값을 얻는다. (이에 따라, prior과 posterior 함수 구성 시 입력변수에 kernel_size와 bias_size를 작성하지 않는 것으로 생각된다.) kl_weight는 prior과 posterior 사이의 KL divergence loss를 scale하는 값이다. 즉, KLD loss를 1/train_size로 scaling한다.
마지막으로 학습 시 사용하는 출력값은 deterministic한 값이므로 Dense를 활용하여 1개의 값이 출력값이 되도록 만든다.

(4) 예측함수

def compute_predictions(model, examples, targets, sample, iterations=100):
    predicted = []
    for _ in range(iterations):
        predicted.append(model(examples).numpy())
    predicted = np.concatenate(predicted, axis=1)

    prediction_mean = np.mean(predicted, axis=1).tolist()
    prediction_min = np.min(predicted, axis=1).tolist()
    prediction_max = np.max(predicted, axis=1).tolist()
    prediction_range = (np.max(predicted, axis=1) - np.min(predicted, axis=1)).tolist()

    for idx in range(sample):
        print(
            f"Predictions mean: {round(prediction_mean[idx], 2)}, "
            f"min: {round(prediction_min[idx], 2)}, "
            f"max: {round(prediction_max[idx], 2)}, "
            f"range: {round(prediction_range[idx], 2)} - "
            f"Actual: {targets[idx]}"
        )

BNN의 출력값은 범위로 출력된다. 범위로 출력하기 위해서는 다수의 샘플이 필요하며, 이에 따라 iteration만큼 값을 예측한다. 이 때, 분포로 구성된 모델에 의해 예측값은 항상 다르게 나타난다. 이렇게 얻어진 샘플값을 바탕으로 각 예측에 대한 mean, min, max값을 구할 수 있으며 이 값들이 BNN에서 얻고자하는 출력값이다.