[論文レビュー] Transformers Can Do Arithmetic with the Right Embeddings

Sean McLeish, Arpit Bansal+9 more|arXiv (Cornell University)|May 27, 2024

Computability, Logic, AI Algorithms被引用数 6

ひとこと要約

この論文は Abacus Embeddings を導入し、数字内の桁位置をエンコードして、トランスフォーマーが長距離の算術を実施し、桁長を大幅に増やして外挿できるようにする。入力注入とループ型トランスフォーマーアーキテクチャの助けを借りて。

ABSTRACT

The poor performance of transformers on arithmetic tasks seems to stem in large part from their inability to keep track of the exact position of each digit inside of a large span of digits. We mend this problem by adding an embedding to each digit that encodes its position relative to the start of the number. In addition to the boost these embeddings provide on their own, we show that this fix enables architectural modifications such as input injection and recurrent layers to improve performance even further. With positions resolved, we can study the logical extrapolation ability of transformers. Can they solve arithmetic problems that are larger and more complex than those in their training data? We find that training on only 20 digit numbers with a single GPU for one day, we can reach state-of-the-art performance, achieving up to 99% accuracy on 100 digit addition problems. Finally, we show that these gains in numeracy also unlock improvements on other multi-step reasoning tasks including sorting and multiplication.

研究の動機と目的

• トランスフォーマーの算術能力を制限する設計上およびデータ表現上のボトルネックを特定する。
• 桁の重要性をエンコードし、長さの一般化を改善する Abacus Embeddings を提案する。
• 再帰と入力注入が Abacus Embeddings とどのように相互作用して、加算・乗算・ソートの性能を向上させるか評価する。
• 100 桁以上の算術への最新の外挿性能と他のアルゴリズム的タスクへの移植性を実証する。

提案手法

• Abacus Embeddings の導入: 数字内で同じ桁のすべての桁に適用される学習可能な位置エンベディング。
• 20-million-sample addition data 上で least-significant-digit-first 形式とパディングなしでデコーダー専用因果トランスフォーマーを訓練する。
• 絶対埋め込みと相対埋め込みスキームにわたって、標準的なトランスフォーマー、入力注入バリアント、ループ型トランスフォーマーを比較する。
• 分布内、分布外、および極端な OOD パフォーマンスを評価し、100+ 桁の加算を含む。
• 乗算とソートへの実験を拡張し、他のアルゴリズム的タスクへの一般化を検証する。
• Abacus Embeddings と FIRE および RoPE の相対位置エンベディングとの適合性を調べる。

Figure 1: Zero shot exact match accuracy on addition using depth sixteen transformer (decoder only) models trained on operands of up to 20 digits. Compared to state-of-the-art embeddings (left), our new Abacus Embeddings (right) dramatically improve generalization to unseen digit lengths. The interi

実験結果

リサーチクエスチョン

• RQ1 Abacus Embeddings は、訓練長を超える多桁の加算に対して長さ一般化とゼロショット外挿を可能にするか。
• RQ2 再帰と入力注入は Abacus Embeddings と併用した場合、一般化誤差を更に低減するか。
• RQ3 これらの手法は、より大きな算術（乗算）および非算術的アルゴリズムタスク（ソート）へどの程度転用できるか。
• RQ4 Abacus Embeddings は FIRE や RoPE のような既存の相対位置エンベディングと互換性があり、補完的か。

主な発見

• Abacus Embeddings は加算の一般化を劇的に向上させ、100 桁の加算で最大 99.1% の精度を実現し、120 桁の問題へ外挿する。
• Abacus Embeddings を入力注入とループ型トランスフォーマーと組み合わせると、ほぼ完璧な一般化と、訓練長さに対して最大6×の長さ拡張を達成する。
• 再帰を備えたループ型トランスフォーマーは、分布外の加算において非再帰的ベースラインと比較して誤差率を最大約50%削減できる。
• Abacus Embeddings は乗算の分布内性能も向上させ、FIRE と組み合わせた分布外のソート精度を高める。
• Abacus Embeddings は FIRE および RoPE と互換性があり、FIRE と併用することで FIRE のみが達成しうる一般化を超える可能性を開く。

Figure 2: Visualization of data formats and positional embeddings. Abacus Embeddings give the same positional embeddings to all digits of the same significance.