Scribble at 2024-11-02 20:45:14 Last modified: 2024-11-02 20:48:13
Transformers have become the predominant architecture in foundation models due to their excellent performance across various domains. However, the substantial cost of scaling these models remains a significant concern. This problem arises primarily from their dependence on a fixed number of parameters within linear projections. When architectural modifications (e.g., channel dimensions) are introduced, the entire model typically requires retraining from scratch. As model sizes continue growing, this strategy results in increasingly high computational costs and becomes unsustainable. To overcome this problem, we introduce TokenFormer, a natively scalable architecture that leverages the attention mechanism not only for computations among input tokens but also for interactions between tokens and model parameters, thereby enhancing architectural flexibility. By treating model parameters as tokens, we replace all the linear projections in Transformers with our token-parameter attention layer, where input tokens act as queries and model parameters as keys and values. This reformulation allows for progressive and efficient scaling without necessitating retraining from scratch. Our model scales from 124M to 1.4B parameters by incrementally adding new key-value parameter pairs, achieving performance comparable to Transformers trained from scratch while greatly reducing training costs. Code and models are available at \url{this https URL}.
現代の機械学習で採用されているディープ・ラーニングのテキスト処理においては、この Transformer モデルが圧倒的な注目を集めていて、もちろんプロンプトを扱う大半のモデルでも活用されているし、画像を生成する際にプロンプトとパラメータを関連付けするためにも使われているので、Stable Diffusion の利用者にとっても基本的な仕組みだと言って良い。
テキスト・データの中に格納されているパラメータには、単語や記号に加えて重み付けという要素も加わっている。簡単に言えば、パラメータは他のパラメータとのあいだに定義されている「距離」という概念を含んでいて、「カレー」は「ちんこ」よりも「うんこ」に近いとか、そういった特性が定義されている。モデルの大きさ、つまりデータ量を増やそうとすると、既存のパラメータどうしに定義されている距離を新しいパラメータについても測り直す（トレーニングする）必要があるので、このところ巨大化していっている生成 AI のモデルにおいては、再トレーニングのコストも増えている。
本論文では、再トレーニングにおいてパラメータをトークン（入力データ）として扱うことにより、既に定義されている距離との関係によって新しいパラメータの定義を算出する。こうすると、全てのパラメータどうしの距離を最初から測り直すというムダをなくせるということだ。ただし、拡散モデルは確率という要素が含まれるので、その距離の算出には乱雑さがある。したがって、安定した結果を出そうとすれば乱雑さを減らさなくてはいけないわけだが、あまりにも入力したデータについて安定した結果を得るように調整すると、もちろん「過適合」が起きてしまい、以前もご紹介した「AI が概念を見失う」という現象が起きる（骨をくわえた犬の画像ばかり学習させると、AI は骨をくわえている動物が「犬」であると学習してしまう）。
こういうことを防ぐためにも、現在は attention という別の仕組みが入っているわけだが、どういう要素について優先した重みをつけるべきなのかは、あいかわらずフレーム問題に関わってくるので、これまた既存の価値観とかオンラインから収集した結果だけに依存して基準をつくると、別の意味での過適合（僕が当サイトで「エコーチェインバー」と読んでいる収束）が起きると思う。