📌 JAX로 시계열 예측 Transformer 구현 - 장기 패턴까지 학습하는 Self-Attention 기반 모델

📌 JAX로 시계열 예측 Transformer 구현 - 장기 패턴까지 학습하는 Self-Attention 기반 모델

---

🚀 왜 시계열에 Transformer를 사용할까?

전통적인 RNN/LSTM은 장기 의존성 문제로 인해 과거 정보를 멀리 반영하기 어렵습니다.
반면 Transformer 모델은 Self-Attention 메커니즘을 통해
입력 시퀀스 내의 모든 시점 간 관계를 동시 계산하여,
멀리 떨어진 시점 간의 패턴까지 효과적으로 학습할 수 있습니다.

---

💡 1. Transformer for Time Series - 구조 요약

📐 구성 요소

• 포지셔널 인코딩: 시계열 순서를 반영
• Self-Attention Layer: 모든 시점 간의 상호관계 학습
• Feed-Forward Network: 정보 변환
• Output Layer: 미래 시점 예측

---

🔧 2. 데이터셋 준비

import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import jax.numpy as jnp

# 예시: 주식 가격 시계열
df = pd.read_csv('stock_prices.csv')  # 'date', 'close' 컬럼 가정
data = df['close'].values.reshape(-1, 1)

# 정규화
scaler = MinMaxScaler()
data_scaled = scaler.fit_transform(data)

# 윈도우 생성
def make_window(data, window=30):
    X, y = [], []
    for i in range(len(data) - window):
        X.append(data[i:i+window])
        y.append(data[i+window])
    return jnp.array(X), jnp.array(y)

x_data, y_data = make_window(data_scaled)

---

🧠 3. Transformer 모델 구현 (Flax)

🎯 포지셔널 인코딩

def positional_encoding(seq_len, d_model):
    pos = jnp.arange(seq_len)[:, jnp.newaxis]
    i = jnp.arange(d_model)[jnp.newaxis, :]
    angle_rates = 1 / jnp.power(10000, (2 * (i // 2)) / d_model)
    angle_rads = pos * angle_rates
    angle_rads = angle_rads.at[:, 0::2].set(jnp.sin(angle_rads[:, 0::2]))
    angle_rads = angle_rads.at[:, 1::2].set(jnp.cos(angle_rads[:, 1::2]))
    return angle_rads

---

🧱 Self-Attention Layer

from flax import linen as nn

class MultiHeadAttention(nn.Module):
    d_model: int
    num_heads: int

    @nn.compact
    def __call__(self, x):
        head_dim = self.d_model // self.num_heads
        assert self.d_model % self.num_heads == 0

        qkv = nn.Dense(self.d_model * 3)(x)
        q, k, v = jnp.split(qkv, 3, axis=-1)

        q = q.reshape(x.shape[0], x.shape[1], self.num_heads, head_dim).transpose(0, 2, 1, 3)
        k = k.reshape(x.shape[0], x.shape[1], self.num_heads, head_dim).transpose(0, 2, 1, 3)
        v = v.reshape(x.shape[0], x.shape[1], self.num_heads, head_dim).transpose(0, 2, 1, 3)

        attn = jnp.einsum('bhqd,bhkd->bhqk', q, k) / jnp.sqrt(head_dim)
        attn_weights = nn.softmax(attn, axis=-1)
        out = jnp.einsum('bhqk,bhvd->bhqd', attn_weights, v)
        out = out.transpose(0, 2, 1, 3).reshape(x.shape[0], x.shape[1], self.d_model)

        return nn.Dense(self.d_model)(out)

---

🧱 Transformer Block

반응형

class TransformerBlock(nn.Module):
    d_model: int
    num_heads: int
    ff_dim: int

    @nn.compact
    def __call__(self, x):
        attn = MultiHeadAttention(self.d_model, self.num_heads)(x)
        x = nn.LayerNorm()(x + attn)

        ff = nn.Dense(self.ff_dim)(x)
        ff = nn.relu(ff)
        ff = nn.Dense(self.d_model)(ff)
        x = nn.LayerNorm()(x + ff)
        return x

---

🔮 4. 전체 Transformer 시계열 모델

class TimeSeriesTransformer(nn.Module):
    d_model: int = 64
    num_heads: int = 4
    ff_dim: int = 128
    num_layers: int = 2
    seq_len: int = 30

    @nn.compact
    def __call__(self, x):
        pos_enc = positional_encoding(self.seq_len, self.d_model)
        x = nn.Dense(self.d_model)(x)
        x += pos_enc

        for _ in range(self.num_layers):
            x = TransformerBlock(self.d_model, self.num_heads, self.ff_dim)(x)

        x = jnp.mean(x, axis=1)  # Global average pooling
        return nn.Dense(1)(x)

---

⚙️ 5. 손실 함수 및 학습 루프

import optax
from flax.training import train_state
import jax

model = TimeSeriesTransformer()
key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, 30, 1)))
tx = optax.adam(1e-3)

state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

@jax.jit
def loss_fn(params, x, y):
    pred = model.apply(params, x)
    return jnp.mean((pred.squeeze() - y.squeeze()) ** 2)

@jax.jit
def train_step(state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(state.params, x, y)
    state = state.apply_gradients(grads=grads)
    return state, loss

---

🏃 6. 학습 수행

for epoch in range(10):
    for i in range(0, len(x_data), 64):
        x_batch = x_data[i:i+64][..., None]
        y_batch = y_data[i:i+64]
        state, loss = train_step(state, x_batch, y_batch)
    print(f"Epoch {epoch+1}, Loss: {loss:.4f}")

---

✅ 7. 미래 예측 및 시각화

import matplotlib.pyplot as plt

def predict_next(state, last_seq, steps=30):
    preds = []
    seq = last_seq
    for _ in range(steps):
        pred = model.apply(state.params, seq[None, ..., None])
        preds.append(pred.squeeze())
        seq = jnp.concatenate([seq[1:], pred], axis=0)
    return scaler.inverse_transform(jnp.array(preds).reshape(-1, 1))

# 예측
future = predict_next(state, x_data[-1])
plt.plot(scaler.inverse_transform(data), label='원본')
plt.plot(range(len(data), len(data) + 30), future, label='Transformer 예측', color='red')
plt.legend()
plt.title("Transformer 기반 시계열 예측")
plt.show()

---

📌 다음 글 예고: JAX로 시계열 이상 탐지 모델 구현 (Autoencoder 기반 Anomaly Detection)

다음 글에서는 JAX를 활용하여 Autoencoder 기반 이상 탐지 모델을 구현하여
시계열에서 이상치(Anomaly)를 자동으로 탐지하는 방법을 소개하겠습니다.

---

 

JAX, Transformer, 시계열 예측, Time Series, Self-Attention, Multi-Head Attention, 딥러닝, Python, 시계열 분석, 고속 연산, Flax, 주식 예측, 기온 예측, 모델 학습, 인공지능, 고성능 컴퓨팅