[Paper Review] Trie-Aware Transformers for Generative Recommendation
Problem Statement
Generative Recommendation
Generative Recommendation은 아이템을 개별 토큰들의 모음으로 변환한 다음, 생성형 모델을 활용해서 예측할 아이템의 토큰들을 순차적으로 예측하는 태스크입니다. 기존처럼 아이템을 추천하기 보다는, 다음에 추천할 아이템을 ‘토큰 시퀀스 생성’ 문제로 바꿔서 모델이 직접 생성하게 만드는 패러다임입니다.
이러한 Generative Recommendation, 생성형 추천에는 크게 2가지 과정을 거칩니다.
먼저, Item tokenization은 각 아이템을 하나의 ID로 두는 것이 아닌, 여러 개의 이산 토큰으로 이루어진 시퀀스로 변환하는 과정입니다. 주로 RQ-VAE, RQ-KMeans 등의 계층적 방법을 이용해서 변환합니다. 이 과정으로 생성되는 토큰 시퀀스들은 자연스럽게 prefix tree 구조, 즉, 한 아이템은 루트→리프까지의 ‘경로’로 표현되고, 비슷한 아이템끼리는 앞부분(prefix) 토큰을 공유하게 됩니다. 다음으로는, Autoregressive Generation인데, 사용자의 상호작용 정보를 바탕으로 다음 아이템에 해당하는 토큰을 순차적으로 생성하는 것을 말합니다. 전체 아이템 집합에서 바로 고르는 것”이 아니라, trie에서 다음 토큰(= 다음 자식 노드)을 층별로 선택해 내려가는 과정이기에, 생성하는 후보의 수가 갈수록 제한적이게 됩니다.
Previous Limitations
하지만, 기존의 Generative Recommendation은 크게 2가지의 한계점을 가지고 있습니다.
트라이(Trie) 구조로 이루어진 아이템 토큰들이 Autoregressive Modeling 과정을 통한 생성 과정에서 계층성이 반영되지 않는 문제가 있습니다. 트라이 구조의 토큰을 가지는 아이템들은 계층적 의미와 경로 기반 관계를 가지는데, 기존의 생성형 추천 모델들은 Transformers 모델 기반의 Autoregressive modeling으로 다음 토큰을 생성합니다. 이 과정에서, Hierarchical한 구조의 토큰을 linear stream 안에서 평평하게 만들기 때문에, 구조적 위치 정보가 모델 입력 과정에서 명시적으로 보존되지 않아, 계층성이 생성 과정에서 제대로 반영되지 못하는 문제가 있습니다. 쉽게 설명하면, 이러한 계층 정보가 들어가기는 하지만, 트랜스포머에서 사용하는 Positional Encoding은 시퀀스 내 순서나 거리만 반영하고, 트라이에서의 계층 정보를 반영하지 못한다는 문제입니다.
이 문제로 인해 계층적 토큰에서 구조적 위치(structural position), 관련 토폴로지(relational topology, 위상 정보) 2개 정보가 트랜스포머에 제대로 반영되지 않게 됩니다. 이 두 정보에 대해 설명드리면,
- 구조적 위치 : 해당 토큰이 트라이의 어디에 위치해 있는지에 대한 depth, path 등의 정보
- 관련 토폴로지 : 두 토큰이 트라이에서 어떻게 연결되어 있는지 (어떤 구조적인 정보를 가지고 있는지)
이러한 2가지 정보가 반영되지 않아, 저자들은 이 2가지 정보를 모델 내부에서 직접 활용할 수 있어야 한다라고 주장하고 있습니다.
Method : TrieRec - Trie-aware Generative Recommendation
트라이 구조에서 계층 정보, 토폴로지 등을 제대로 활용하기 위해 위 논문에서는 TrieRec을 제안합니다. TrieRec은 Transforer에 2개의 Positional Encoding을 추가하는 방식으로, Trie-Aware Absolute PE (TAE)에서는 토큰이 트라이에 어디에 있는지를 더해 토큰 임베딩에 더해 주입하여 구조적 위치를 인식하도록 하고, Topology-aware Relative PE(TRE)는 토큰 쌍의 트라이 위 관계를 Self-Attention 점수에 Bias로 더해 주입해 관련 토폴로지를 트랜스포머 내부에 주입하는 방법입니다.
이 과정에서 Input, Output은 다음과 같습니다.
- Input : 사용자 상호작용 이력으로 만들어진 토큰 시퀀스 모음, 해당 토큰들에 대한 트라이 구조 등의 정보
- Output : TAE, TRE가 반영된 self-attention을 통해 다음 아이템의 semantic ID를 utoregressive 방식으로 생성합니다.
Trie-aware Absolute Positional Encoding (TAE)
TAE는 각 토큰이 trie 구조에서 가지는 절대적인 구조 위치(absolute structural position) 정보를 Transformer의 입력 임베딩에 반영하는 단계입니다. 기존 Positional Encoding이 절대적인 위치를 반영하는 대신, token depth와 계층 구조 정보에 기반한 포지셔널 인코딩을 생성합니다. 즉, 각 토큰 $c_{i,l}$에 해당하는 포지션 인코딩을 생성합니다.
해당 과정에서 input, output은 다음과 같습니다.
- Input : 각 토큰에 대한 임베딩
- Output : 각 토큰에 해당하는 트라이 구조 내의 포지션 정보가 반영된 positional encoding
해당 과정의 전체 식은 다음과 같습니다.
\[\mathbf{x}_i^{\mathrm{con}} = \left[\mathbf{e}_{i,1}^{\top};\mathbf{e}_{i,2}^{\top};\cdots;\mathbf{e}_{i,L}^{\top}\right]^{\top} \in \mathbb{R}^{L\cdot d}\] \[\hat{\mathbf{x}}_i = \mathrm{MLP}\!\left(\mathbf{x}_i^{\mathrm{con}}\right) \in \mathbb{R}^{L\cdot d}\] \[\mathbf{p}_{i,l} = \mathbf{W}_l \hat{\mathbf{x}}_i + \mathbf{b}_l, \; l=1,\ldots,L\]첫번째 식부터 설명드리면, 여기서 $e_{i,l}$는 토큰 $c_{i,l}$의 semantic embedding입니다. 첫번째 식은 Concatenation으로, 각 토큰에 대해 계층적인 context 이웃 정보, 즉 조상, 후손 정보를 포함한 토큰 경로 상의 토큰 임베딩을 수집하고, 이를 하나의 벡터로 결합하여 옆으로 붙인 형태의 연결된 임베딩을 만듭니다.
여기서 $i, l$은 아이템 인덱스, 계층 토큰의 레이어 인덱스를 말합니다. 이로 인해 만들어지는 $\mathbf{x}^{con}_i$는 아이템 $i$의 전체 경로 토큰들을 붙인 벡터를 말합니다.
그 다음으로, Aggregation 과정을 수행합니다. Concatenation 단계에서 결합된 임베딩을 MLP에 통과시켜 토큰이 Trie의 어디에 위치하는지를 반영한 absolute structural representation (구조 표현)을 생성합니다. 여기서 MLP는 Position-dependent하게 동작하는데, 각 depth 정보가 concat 벡터의 고정된 위치에 배치되기 때문에, MLP는 위치별 의미를 반영해 정보를 다르게 조합할 수 있고, 여러 위치와 여러 차원의 정보를 함께 섞어 조상–자손 관계와 같은 계층 구조 패턴을 더 풍부하게 표현할 수 있게 됩니다. 이러한 점에서, Cross-token, Cross-dimension interaction이 가능하게 됩니다.
여기서 permutation-invariant는 순서가 바뀌어도 출력이 변하지 않는 것을 의미합니다. MLP는 2-lyaer Network에 GeLU와 Layer normalization을 적용합니다.
참조로, Cross-token, Cross-dimension을 예시로 설명하면 다음과 같습니다. (오른쪽의 TAE blocked는 뒤의 Ablation study에서 진행합니다.)
- 각 depth 별로 차원이 고정된 위치에 배치되는데, 여러 차원의 정보들을 조합해서 학습 하는 등의 학습을 가능하게 하는 요소입니다.
이 부분의 이미지를 예시로 보여드리면, 다음과 같습니다.
마지막으로, Projection 과정은 깊이(depth) $l$마다 다른 선형 변환을 적용하여 $W_l, b_l$로 각 depth의 포지션 벡터를 추출합니다. 여기서 projector은 single linear layer입니다. 즉, 같은 경로 컨텍스트를 보더라도 Depth가 다르면, 다른 가중치를 적용하도록 합니다.
이 과정으로 학습되는 포지션 인코딩을 원래 토큰 임베딩에 더해 다음과 같이 사용됩니다.
\[\hat{e}_{i,l} = e_{i,l} + p_{i,l}\]이 포지셔널 인코딩 정보는 Self-Attention의 1D Positional Encoding 대신, Token Embedding에 더해지는 Trie-aware absolute positional representation으로 사용합니다.
이러한 식으로, 토큰 임베딩 자체가 자신의 depth, 경로 컨텍스트를 인지하도록 합니다. 추가로, TAE가 첫번째 Transformer layer에서만 효율성을 위해 추가됩니다. 생성 과정에서는 아직 해당 아이템의 전체 토큰 경로를 확인할 수 없기 때문에, 입력으로 들어오는 과거 아이템 목록의 토큰들에 대해서만 적용되고, 이렇게 주입된 구조 정보는 이후 layer로 전달됩니다.
Topology-aware Relative Positional Encoding (TRE)
TAE가 각 토큰이 Trie에서 어디에 있는지 (절대적 구조 위치)를 토큰 임베딩에 더해주는 방식이라면, TRE는 토큰 쌍 $u, v$ 사이의 트라이 상 관계(토폴로지)를 self-attention 계산에 bias로 직접 주입하는 단계입니다. 이 과정에서 한 시퀀스 안의 모든 토큰 쌍에 대해 계산을 진행하기 때문에, 같은 아이템 경로에 있는 토큰 쌍이든, 서로 다른 아이템에서 온 토큰 쌍이든 모두 포함하게 됩니다.
해당 단계에서 input, output은 다음과 같습니다.
- Input : Self-Attention에 들어가는 각 토큰의 임베딩 (Transformer 레이어의 입력 토큰 임베딩), 토큰 쌍의 토폴로지 정보
- Output : 토폴로지 기반의 편향이 포함된 어텐션 가중치
TRE는 토폴로지 기반의 디자인으로 토큰 쌍의 구조적 특징 (거리, LCA 등)을 명시적으로 모델링하고, 무거운 계산은 피하고, 간단, 확장 가능하게 하는 효율성을 목표로 하고 있습니다.
먼저, Self-Attention에서 토큰 쌍마다 트라이 관계를 요약하는 3가지 값을 계산합니다. 여기서 각 정보는,
- $d^{LCA}_{u,v}$ : 두 토큰의 LCA(lowest common ancestor)의 깊이
- $\delta_u$ : 토큰 $u$에서 LCA까지 거리 $(depth(u) - depth(LCA(u, v)))$
- $\delta_v$ : 토큰 $v$에서 LCA까지 거리 $(depth(v) - depth(LCA(u, v)))$
여기서 LCA는 Lowest Common Ancestor (최소 공통 조상)으로서, 트라이에서 두 토큰 노드 $i, j$가 공유하는 공통 조상들 중, 가장 아래에 있는, 두 토큰 경로라 마지막으로 공유하는 가장 깊은 조상 노드를 말합니다. 예시를 들어서 설명하면, $i$는 root -> A -> B -> D라 하고, $j$는 root -> A -> B -> E라 하면, B까지는 같고 E에서 갈라지기 때문에, LCA는 B라고 할 수 있습니다.
논문에서는 이 3개의 정보가 합쳐지면, 두 토큰이 얼마나 prefix를 공유하는지 알 수 있고, 서로 얼마나 트라이에서 떨어져 있는지를 알 수 있다고 설명합니다. 즉, “LCA depth가 깊다 → 공유 prefix가 길다 → 더 비슷한 아이템에서 온 토큰쌍일 가능성이 높다” 라는 판단이 나오는 것이죠.
이 정보를 Self-Attention에 넣게 되면, 다음의 식으로 반영됩니다.
\[\tilde{A}_{u,v} = \mathrm{Softmax}_v\!\left( \frac{\left(e_u W^{Q}\right)\left(e_v W^{K}\right)^{\top}}{\sqrt{d}} \;+\; B\!\left(d^{LCA}_{u,v},\, \delta_u,\, \delta_v\right) \right)\]이 식에서처럼 원래의 Self-Attention 점수에 토폴로지 기반의 편향 점수가 추가되는데, 이러한 토폴로지 정보가 반영되기 위해서 다음의 과정을 거칩니다.
여기에서 $B(\cdot)$은 학습 가능한 scalar bias lookup table(관계 테이블)인데, 여기서 $B(d^{LCA}_{u,v}, \delta_u, \delta_v)$는 3개 feature의 조합(triple)로 인덱싱되어 해당 토큰쌍 관계에 맞는 bias 값을 반환합니다. 즉, 이런 topology 관계를 가진 토큰쌍이면 attention score에 얼마만큼의 편향을 더할지를 Transformer 모델 학습으로 결정합니다.
$B$ (관계 테이블)는 토큰 쌍 (u, v)의 토폴로지 관계를 key로, self-attention 점수에 더해줄 bias 값을 꺼내오는 lookup-table(학습 파라미터)입니다. 이를 이용해 둘이 트라이에서 가까운 관계, LCA가 깊어서 prefix를 많이 공유하면 더 보게끔 학습으로 반영합니다. 이러한 relation table에서 하나의 bias를 결과로 가져와, 위의 식에 활용합니다.
추가로, $\delta_u$, $\delta_v$ 를 보면, 각 토큰이 LCA에서 얼마나 떨어져 있는지 알 수 있어, 두 토큰의 트리 상 거리와 상대적 계층 관계까지 Attention에 반영할 수 있습니다.
추가적으로, 효율성을 위해 편향 파라미터 B는 레이어들 사이에 공유해 한 번 계산해 재사용할 수 있습니다.
Discussion
Model-agnostic
TrieRec의 TAE/TRE는 특정 Transformer 백본이나 특정 GR 학습 목표에 종속되지 않고, trie 기반 semantic ID 토큰화를 사용하는 다양한 생성형 추천 모델에 plug-in 형태로 적용할 수 있습니다.
Efficiency
TrieRec이 구조 정보를 넣지만, 추가 비용은 크게 발생하지 않습니다. 우선, T개의 아이템에 L개의 토큰드을 가지고 있다고 가정하겠습니다. 거기에 트랜스포머 모델이 d차원의 임베딩, h개의 헤드와 k개의 layer을 가지고 있다고 가정합니다.
- TAE에서는 각 토큰의 계층 컨텍스트를 MLP로 집계하는 데 토큰 당 $O(L^2d^2)$, 시퀀스 전체로는 $O(TL^2d^2)$의 비용이 발생합니다.
- TRE에서는 토큰쌍마다 LCA 계산이 필요한데, pair 당 $O(log L)$, 전체로는 $O(hT^2L^2log L)$의 비용으로 적고, bias lookup 자체는 $O(1)$의 비용이 발생합니다.
- 그리고 이러한 비용이 기본 트랜스포머 연산 $O(khT^2L^2d)$보다 적고, 실제 실험에서도 오버헤드가 10% 미만으로 발생합니다.
- $k\cdot h\cdot t$가 d보다 큰 것을 고려하고, $k \cdot d$가 $log L$보다 비용이 더 많이 발생합니다.
Hyperparameter Tuning
TrieRec에서는 별도의 새로운 하이퍼파라미터를 도입하지 않습니다.
다른 Positional Encoding과의 차이
기존 Positional Encoding은 시퀀스에서 어느정도 떨어졌는지를 확인하는 시퀀스 상 거리를 다루는데, TrieRec은 트라이 상 거리, 관계를 사용합니다. 즉, LCA, 트라이 거리, item-path 등의 트라이 구조 정보를 트랜스포머 내부 게산 자체에 넣을 수 있어, 성능 향상이 발생합니다.
DiscRec과의 차이
DiscRec은 같은 아이템의 다른 토큰 정보를 타겟 토큰에 섞어주고, 선형 임베딩 결합을 사용하는 것에 반해, TrieRec은 MLP를 통해 더 미세한 차원 수준의 상호작용을 만듭니다.
Experiments
Experiments Settings
위 논문의 실험에서는 Amazon의 Beauty, Toys, Food 데이터셋과 Movielens 1M 데이터를 사용합니다. 또한, 5-core 필터링과 최근 20개 상호작용 데이터만 적용합니다. 추가로, Leave-one-out split 방식을 적용하여 제일 마지막을 test, 그 이전 1개는 valid, 나머지는 train으로 사용합니다.
추가로 Beam Search에서 K는 10을 적용합니다.
모든 GR 모델에서 TIGER와 동일한 조건을 적용했습니다.
Overall Performance
TrieRec을 TIGER, LETTER, CoST 모델을 결합했을 때, 가장 좋은 성능이 도출되었습니다. 추가로, 다른 결합 모델과 비교를 해도, 트라이 구조의 Autoregressive modeling에 명시적으로 인코딩하는 TrieRec이 더 좋은 성능이 도출됨을 확인했습니다.
추가로, LETTER 모델 기반에서 SEATER에서 수행한 토크나이징 방법과도 4개 데이터셋에서 비교한 결과를 분석해도 더 좋은 성능이 나오는데, SEATER은 전역 정보를 고려하는 것보다, 지역적으로 부모-자녀 관계만 고려합니다. 이에 반해, TrieRec은 거리 정보도 같이 고려하는 것으로 보아 더 좋은 성능이 도출됩니다.
Ablation Study
위 실험에서는 LETTER 모델에 기반해 다음의 Ablation Study를 적용합니다.
- TAE 관련
- w/o TAE: TAE 전체 제거
- TAE-Blocked: MLP에서 같은 feature dimension끼리만 depth 간 상호작용하도록 제한(차원 간 cross-dimension interaction 차단)
- TAE-Prefix: TAE를 prefix-sum 기반 누적 결합(단순 누적)으로 대체
실험 결과, root -> leaf 경로를 명시적으로 인코딩하는 것이 트라이 계층 의미를 잡는 것이 중요합니다. 또한, cross-dimensional interaction 자체가 중요한 것을 확인할 수 있습니다. 무엇보다, 단순 누적으로는 복잡한 계층 관계를 반영하지 못해, 원래 LETTER보다도 더 낮은 성능이 도출됩니다.
- TRE 관련
- w/o TRE: TRE를 제거하고, 표준 1D relative position bias만 유지
- TRE-w/o LCA: TRE에서 LCA depth 정보 제거(수직 거리 $\delta$만 사용)
- TRE-w/o $\delta$ : TRE에서 수직 거리 $\delta$ 제거(LCA depth만 사용)
실험 결과, 공통 조상에 대한 단순 상대 거리만으로는 제 성능이 나오지 않고, 정확한 정보가 필요함을 알 수 있습니다.
Efficiency Analysis
TrieRec은 추가 모듈로 인해 실행 시간이 약간 증가하지만, 모델 확장성을 해치지 않고, 성능은 더 증가함을 보입니다. 이를 위해 전체 학습 시간, 1회 추론 지연 시간을 측정한 결과를 보여줍니다.
실험 결과, 학습/추론 비용이 LETTER와 같은 규모 안에서 유지되고, 증가폭이 크지 않아 크게 문제되지 않음을 증명합니다. 즉, 성능 향상과 오버헤드의 트레이드오프가 어느정도 필요하다라는 것을 증명하고 있습니다.
Conclusion
Contribution
- 기존 Generative Recommendation에서 계층 정보가 제대로 반영되지 않는다는 문제를 제시하고, 이에 대한 해결책을 제시함으로서, 트라이 구조가 GR 생성에서 실제로 반영하는 문제를 다루었습니다.
- 2개의 Positional Encoding을 적용하여 구조 신호를 효과적으로 도입하였고, 비용 역시 크지 않아, 더 높은 성능을 기대할 수 있습니다.
Limitations
- Semantic ID를 사용하는 RQ-VAE 등의 문제가, Multi-domain에서의 활용 문제가 있는데, 이러한 문제에서 자유롭지 못합니다. 쉽게 말하면, 대량의 데이터나, 여러 데이터셋을 같이 사용하는 방식에서는 검증되지 않았고, 실제 산업에서 적용되기가 힘들다는 문제가 있습니다.
- 어떤 부분에서 성능 개선이 일어났는지, 쉽게 말해, Cold start를 개선했다든지 등의 개선 요소가 드러나지 않아 설득력이 부족하고, 이를 해결하기 위한 추가적인 실험이 필요합니다.
- 모델이 토큰 기반의 계층 구조를 활용하지만, 해당 계층 구조가 실제 아이템 간 의미를 제대로 반영하는지에 대한 분석이 필요할 것 같습니다. 대표적으로 RQ-VAE 기반 방법은 잔차에 기반해 토큰화를 진행하기 때문에, 전혀 다른 아이템이 같은 토큰을 사용할 수도 있고, 토큰만 보고 어떤 의미를 가지고 있는지 단정하기 어렵습니다.
Leave a comment