[논문리뷰] Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

ICML 2026. [Paper] [Github]
Dong Hoon Lee, Seunghoon Hong
KAIST
18 Jun 2026

Introduction

기존의 이미지 tokenizer는 고정된 압축률을 사용하기 때문에 단일 모델 내에서 이러한 다양한 요구 사항에 적응할 수 없으며, 따라서 서로 다른 계산 예산과 품질 목표에 따라 여러 모델 변형이 필요하다. 이러한 목표를 달성하기 위해 데이터를 표현하는 데 사용되는 토큰 수를 조절하여 압축률을 조정할 수 있는 variable-length tokenizer (VLT)가 등장했다.

가장 널리 사용되는 메커니즘은 nested dropout으로, 학습 과정에서 중첩된 방식으로 꼬리 토큰을 무작위로 제거한다. 이는 중요한 정보를 초반 토큰에 배치하여 토큰 순서를 정렬하고, 앞부분 토큰을 사용하여 토큰 수를 줄이는 데 도움이 된다. 하지만 nested dropout은 데이터 포인트 간 유사성 구조에 길이에 따른 변화를 유발한다. 초반 토큰은 고차원적인 semantic을 인코딩하고 후반 토큰은 저차원적인 디테일을 저장한다. 결과적으로 데이터 포인트 간의 쌍별 관계는 토큰 길이에 따라 달라진다. 이는 특히 diffusion model에서 문제가 되는데, latent space 분포 및 해당 score function이 토큰 길이에 따라 크게 달라지기 때문에 유연한 품질-계산 균형을 유지하면서 다양한 토큰 길이에 걸쳐 일반화할 수 있는 모델을 학습하기 어렵게 만든다.

이는 토큰 길이에 따른 유사성 구조의 일관성, 즉 표현 정렬을 장려해야 한다는 동기를 부여한다. 일반적인 접근 방식은 표현 간의 차이를 직접 최소화하는 것이지만, 이는 표현 차원, 즉 토큰 개수를 일치시켜 거리를 정의할 수 있는 경우에만 가능하다. 따라서 VLT에 특화된 해결책이 필요하다.

본 논문에서는 표현 정렬을 위한 해결책으로 병합 기반 VLT를 제안하였다. 길이 변조는 병합을 통해 수행되고 diffusion model은 병합 패턴에 따라 작동하므로, 병합 패턴은 전체 길이 표현과 동등한 변조된 latent 표현을 정의한다. 따라서 길이 변조로 인한 표현 변화를 직접 측정할 수 있으며, 유사한 토큰을 결합함으로써 이러한 변화를 최소화할 수 있다. 결과적으로, 본 프레임워크에서는 유사한 토큰을 결합하면 표현 변화가 감소하여 표현 정렬이 촉진된다.

한 가지 과제는 기존의 병합 방법을 생성 환경에 직접 사용할 수 없다는 점이다. Diffusion model은 병합 패턴을 알고 있어야 하지만, 데이터 의존적인 방법은 생성 시점에 사용할 수 없는 입력 이미지로부터 병합 패턴을 도출한다. 이러한 문제를 해결하기 위해, 토큰 병합이 데이터와 독립적인 learnable global merging을 제안하였다. 따라서 생성 모델은 생성 시점에 병합 패턴에 접근할 수 있다. 데이터 독립적임에도 불구하고, learnable global merging 방식은 유사한 토큰을 결합하면서 이미지를 재구성하는 데 최적화되어 있다.

저자들은 VLT에서 생성된 가변 길이 latent에 대해 DiT를 학습시켜 효과적인 품질-계산 제어를 가능하게 하였다. 병합된 토큰의 위치 정보를 proportional attention과 함께 처리하기 위해 병합된 위치 임베딩을 도입했다.

Method

본 논문의 목표는 토큰 길이에 따른 표현 정렬을 촉진하는 variable-length tokenizer (VLT)를 설계하여 diffusion model의 학습을 용이하게 하는 것이다. 이를 위해 토큰 병합을 통해 토큰 길이를 조절하고, 토큰 길이에 따른 latent의 직접적인 정렬을 가능하게 하는 방법을 제안하였다.

1. Merging-based Length Modulation

병합 패턴(즉, 어떤 토큰을 병합할지)은 할당 행렬 $\Gamma \in \{0, 1\}^{N \times (N-K)}$로 주어지며, $\Gamma_{ij} = 1$은 $i$번째 토큰이 $j$번째 클러스터에 할당됨을 나타낸다. $\Gamma$가 주어지면, 길이 변조된 latent $\tilde{\textbf{z}} \in \mathbb{R}^{(N-K) \times D}$는 다음과 같이 계산된다.

\[\begin{equation} \tilde{\textbf{z}} = \bar{\Gamma}^\top \textbf{z}, \quad \textrm{where} \quad \bar{\Gamma}_{ij} = \frac{\Gamma_{ij}}{\sum_k \Gamma_{kj}} \end{equation}\]

$N − K$개의 병합된 토큰이 생성되며, 각 토큰은 연결된 원래 토큰들의 클러스터 중심을 나타낸다. 저자들은 병합된 토큰 크기 $\textbf{m} \in \mathbb{R}^{N-K}$를 정의했는데, 이는 $\tilde{\textbf{z}}$의 각 병합된 토큰이 포함하는 원래 토큰의 수를 $m_j = \sum_i \Gamma_{ij}$로 계산한다.

길이가 변조된 latent $\tilde{\textbf{z}}$를 얻은 후, 이를 이용하여 디코더로 원본 이미지를 재구성하거나 DiT로 이미지를 합성한다. 이를 위해, 각 병합된 토큰의 유효 클러스터 크기 $m_j$를 고려하기 위해, ToMe에서 제안된 바와 같이 디코더와 DiT의 어텐션 모듈에 proportional attention을 적용한다. 구체적으로, 토큰 크기 $\textbf{m}$이 주어졌을 때, self-attention 점수는 다음과 같이 수정된다.

\[\begin{equation} \textbf{A} = \textrm{Softmax} \left( \frac{\textbf{Q}\textbf{K}^\top}{\sqrt{d}} + \log \textbf{m} \right) \end{equation}\]

중요한 것은, proportional attention을 사용하여 $\tilde{\textbf{z}}$를 계산하는 것은 표준 attention을 사용하여 projection된 latent $\textbf{Wz} \in \mathbb{R}^{N \times D}$를 처리하는 것과 동일하다는 점이다. 여기서 projection 행렬 $\textbf{W} \in \mathbb{R}^{N \times N}$는 각 토큰을 동일한 클러스터에 있는 모든 토큰의 평균으로 대체하도록 정의된다.

\[\begin{equation} \textbf{W}_{ij} = \begin{cases} \frac{1}{\vert \mathcal{C}_i \vert} & \textrm{if} \; j \in \mathcal{C}_i \\ 0 & \textrm{otherwise} \end{cases} \\ \textrm{where} \quad \mathcal{C}_i \{j \; : \; \exists c, \Gamma_{ic} = 1 \; \textrm{and} \; \Gamma_{jc} = 1\} \end{equation}\]

($\mathcal{C}_i$는 토큰 $i$와 동일한 클러스터에 할당된 토큰들의 집합)

Proportional attention 하에서의 등가성은 토큰 길이 $N − K$에 관계없이 변조된 latent $\tilde{\textbf{z}} \in \mathbb{R}^{(N-K) \times D}$와 동등한 전체 길이 latent $\textbf{Wz} \in \mathbb{R}^{N \times D}$를 정의한다. $\textbf{z}$와 $\textbf{Wz}$는 크기가 같으므로, 표현 정렬 loss에서와 같이 차이 $\| \textbf{Wz} - \textbf{z} \|^2$를 최소화하여 latent를 직접 정렬할 수 있다.

\[\begin{aligned} \| \textbf{Wz} - \textbf{z} \|^2 &= \sum_{i=1}^N \| z_i - (\textbf{Wz})_i \|^2 \\ &= \frac{1}{2} \sum_{i=1}^N \frac{1}{\vert \mathcal{C}_i \vert} \sum_{j \in \mathcal{C}_i} \| z_i - z_j \|^2 \end{aligned}\]

따라서 $z = (z_1, \ldots, z_N)$ 내에서 유사한 토큰들을 그룹화할 때 이 효과가 최소화된다. 결과적으로, 병합 기반 길이 변조를 통해 각 시퀀스 $\textbf{z}$ 내에서 토큰별로 유사한 토큰들이 병합되도록 유도함으로써, 토큰 길이의 차이에도 불구하고 latent를 직접 정렬할 수 있다.

2. Learnable Global Merging

병합 측면에서, $\| \textbf{Wz} - \textbf{z} \|^2$는 latent $\textbf{z}$에 클러스터링을 적용함으로써 최소화될 수 있다. 그러나 이러한 $\textbf{z}$에 의존하는 병합 방식은 이미지 생성과 호환되지 않는다. 이미지 생성 전에 알 수 없는 이미지별 할당 행렬 $\Gamma$가 필요하기 때문이다. 특히, 생성 과정에서 $\textbf{z}$에 접근할 수 없으므로 병합된 토큰 크기 $\textbf{m}$과 병합된 위치를 계산할 수 없다. 이는 생성 과정에서 proportional attention을 적용할 수 없게 하여, 더 이상 $\tilde{\textbf{z}}$에 대한 연산이 $\textbf{Wz}$에 대한 연산과 동일하지 않게 만든다. 마찬가지로, 적절한 위치 임베딩을 할당하려면 어떤 토큰이 결합되는지 알아야 하는데, 이 또한 생성 과정에서는 알 수 없다.

따라서 본 논문에서는 $\textbf{z}$ 값과 무관하게 병합을 결정하지만, latent의 tokenization 및 정렬에 최적화된 learnable global merging 방식을 제안하였다. $\textbf{z}$ 값과의 독립성 덕분에 $\textbf{m}$과 병합 위치가 다양한 토큰 길이에 걸쳐 일관성을 유지할 수 있으며, proportional attention과 적절한 위치 임베딩을 통해 diffusion model과의 호환성을 확보할 수 있다.

구체적으로, 학습 가능한 임베딩 $\textbf{e} = \{e_1, \ldots, e_N\} \in \mathbb{R}^{N \times D}$를 도입한다. 여기서 각 임베딩 $e_i$는 latent 토큰의 $i$번째 위치에 해당하며, 어떤 토큰을 결합할지 결정한다. 주어진 $K$에 대해, 코사인 유사도를 기반으로 $N-K$개의 클러스터를 형성하기 위해 agglomerative clustering을 적용한다. 그 결과 할당 행렬 $\Gamma$가 생성된다.

\[\begin{equation} \Gamma = \textrm{Agglomerative} (\textbf{e}) \in \{0, 1\}^{N \times (N-K)} \\ \tilde{\textbf{z}} = \bar{\Gamma}^\top \textbf{z} \quad \textrm{where} \quad \bar{\Gamma}_{ij} = \frac{\Gamma_{ij}}{\sum_k \Gamma_{kj}} \end{equation}\]

Agglomerative clustering은 gradient를 차단하는 argmin 연산을 포함하므로 gradient descent를 통한 학습 가능한 임베딩 $\textbf{e}$의 직접 최적화를 방해한다. 이를 해결하기 위해, $\Gamma$에 straight-through trick을 적용하여 학습 가능한 임베딩의 학습을 가능하게 한다.

\[\begin{aligned} \textrm{cluster centers} &: \textbf{c} = \bar{\Gamma}^\top \textbf{e} \\ \textrm{soft assignments} &: \Gamma^\textrm{soft} = \textrm{softmax}(\frac{\textbf{e} \textbf{c}^\top}{\tau}) \\ \textrm{straight-through trick} &: \Gamma \rightarrow [\Gamma - \Gamma^\textrm{soft}]_\textrm{sg} + \Gamma^\textrm{soft} \end{aligned}\]

($[\cdot]_\textrm{sg}$는 stop-gradient 연산, $\tau$는 scaling 파라미터)

Straight-through trick을 적용하여 학습 가능한 임베딩 $\textbf{e}$에 의해 결정되는 할당 행렬을 다음과 같이 최적화할 수 있다.

\[\begin{equation} \tilde{\textbf{z}} = ([\Gamma - \Gamma^\textrm{soft}]_\textrm{sg} + \Gamma^\textrm{soft})^\top \textbf{z} \end{equation}\]

디코더는 변조된 latent $\tilde{\textbf{z}}$로부터 이미지를 $\hat{\textbf{x}} = \mathcal{D}(\tilde{\textbf{z}})$로 재구성하고, 학습 가능한 임베딩 $\textbf{e}$는 VQGAN loss $\mathcal{L}_\textrm{total}$와 유사한 토큰의 병합을 유도하는 alignment loss $\mathcal{L}_\textrm{align}$을 통해 학습된다.

\[\begin{equation} \min_{\textbf{e}, \mathcal{E}, \mathcal{D}} \mathcal{L}_\textrm{total} + \lambda_\textrm{align} \mathcal{L}_\textrm{align} \end{equation}\]

Alignment loss

Latent $\textbf{z}$에서 유사한 토큰 $z_i$, $z_j$는 학습 가능한 임베딩 $\textbf{e}$에서 대응하는 임베딩 $e_i$, $e_j$가 유사할 때만 병합된다. 유사한 토큰의 병합을 유도하기 위해, alignment loss를 통해 latent $\textbf{z}$의 토큰별 유사성을 반영하도록 학습 가능한 임베딩을 학습시킨다.

\[\begin{equation} \mathcal{L}_\textrm{align} = \sum_{i,j} \textrm{ReLU} \left( \left\vert \left[ \frac{z_i \cdot z_j}{\| z_i \| \| z_j \|} \right]_\textrm{sg} - \frac{e_i \cdot e_j}{\| e_i \| \| e_j \|} \right\vert - \delta \right) \end{equation}\]

($\delta$는 과도한 정규화를 방지하기 위한 margin)

3. Application to Diffusion Transformers

저자들은 tokenizer에서 생성된 가변 길이 latent에 대해 DiT를 학습시켜 생성 품질과 계산량 간의 균형을 제어할 수 있도록 하였다. 구체적으로, 토큰 길이 $K$를 랜덤 샘플링하여 가변 길이 latent $\tilde{\textbf{z}}$에 대해 표준 diffusion loss를 적용한 DiT $G_\psi$를 학습시킨다. DiT에서는 병합된 토큰의 유효 크기 $\textbf{m}$을 고려하기 위해 proportional attention을 사용한다.

\[\begin{equation} \mathcal{L}_\textrm{diff} (\psi) = \mathbb{E}_{K, \tilde{\textbf{z}}} \| G_\psi (\tilde{\textbf{z}}_t, t) - \epsilon_t \|^2 \end{equation}\]

DiT는 고정 길이 토큰을 사용하는 대신, 여러 위치의 토큰을 통합한 병합된 토큰의 위치 정보를 인코딩해야 한다. 구체적으로, 학습 가능한 위치 임베딩 $\textbf{p}_e \in \mathbb{R}^{N \times D}$와 토큰 길이 $N − K$에 대한 할당 행렬 $\Gamma$가 주어졌을 때, diffusion model 입력에 대한 위치 임베딩은 다음과 같다.

\[\begin{equation} \tilde{\textbf{p}}_e = \bar{\Gamma}^\top \textbf{p}_e \end{equation}\]

또한, 병합된 위치 임베딩 전략만을 사용한 DiT 학습은 적절한 성능을 달성하지만, 각 토큰 개수에 대한 길이별 LoRA 파라미터를 도입하는 간단한 post-training 절차를 통해 성능을 더욱 향상시킬 수 있다. 구체적으로, 공동 학습 후 길이별 LoRA를 도입하고 이를 사용하여 추가 10 epoch 동안 학습시킨다. 이를 통해 DiT는 각 토큰 개수에 더욱 밀접하게 적응할 수 있으며, 학습/파라미터 오버헤드는 미미한 수준으로 유지하면서 일관된 성능 향상을 제공한다.