仅用原始JSON和图像实现宝可梦嵌入向量的超强效果
嵌入向量是人工智能领域中极为实用但讨论相对较少的概念。它是由数百个数字组成的集合,唯一对应特定对象并定义其维度,现代通常采用128的倍数如384D、768D甚至1536D。维度越大,理论上能包含的"信息"和独特性就越多。这些嵌入向量可直接用于传统的回归和分类问题,但真正有价值的是:如果能在给定查询嵌入向量与其他嵌入向量集合间找到最小数学距离,就能发现最相似的项——这对搜索等实际应用极为有用。
虽然任何对象都能用嵌入向量表示,但文本是最经典的用例。最初的word2vec论文推广了这一概念,后续研究证明词嵌入向量可计算类似"男人+女人-国王=女王"的关系。通过平均所有词嵌入向量可以创建句子嵌入向量,但这种方法未考虑词序和标点符号,而这两者对文本语境识别至关重要。
深度学习出现后,人们发现像BERT这样的大型语言模型能通过涌现行为返回嵌入向量。基于Transformer的LLM凭借注意力机制能更稳健地处理位置关系,并因采用比单词更高级的标记化策略而更好地整合标点符号。Sentence Transformers是一个流行的Python库,特别适合使用all-MiniLM-L6-v2模型(月下载量达3000万次)创建嵌入向量,该模型以384D嵌入向量在编码速度和稳健性间取得平衡。
这些嵌入模型能否处理比普通句子更复杂的内容?能否将大量文本编码到一致的空间?all-MiniLM-L6-v2的上下文长度为512个标记,仅能容纳几段文本,但较新的LLM具有更长的上下文窗口。
记得我早期的一个项目:通过手动转换每只宝可梦的元数据(如基础数值、属性、技能、特性以及颜色、形状和栖息地等杂项属性)创建宝可梦向量,然后进行聚类。但熟悉宝可梦的人知道这仅是冰山一角:还有更丰富的文本数据如宝可梦图鉴条目和具体遭遇地点,这些都能揭示更多信息。当时没有高效的LLM来编码这些额外元数据。
何不尝试用文本嵌入模型编码所有宝可梦元数据并观察结果?我们能否识别最"相似"的宝可梦?“相似"究竟指什么?能否通过最不相似找到最奇怪的宝可梦?能否编码图像等其他数据?让我们一探究竟!
如何使用LLM生成嵌入向量
现代嵌入模型通常通过两种方式之一训练。第一种是在正常训练LLM时通过涌现行为实现:由于LLM需在将输出传递给分类头(如GPT的下一个标记预测)前确定潜在空间,取模型的最后一层(“隐藏状态”)并在位置轴上平均可得到与隐藏状态维度相同的嵌入向量。LLM必须学会在公共潜在空间中唯一表示文本,因此这种方法很自然。第二种是直接训练模型输出嵌入向量:训练过程通常使用对比学习来最小化已知文本对生成嵌入向量间的语义距离,并最大化不相似对间的差异。当然这两种技术可结合使用:在大规模文本上预训练LLM,然后用对比学习微调。
嵌入模型受益于所有为改进生成式AI而投入的LLM研究,如推理速度和更长上下文窗口。虽然使用更大上下文窗口通常需要计算量二次增长(例如输入长度增加2倍需要4倍计算),但由于FlashAttention和旋转位置嵌入,现在无需大型数据中心就能训练具有超长上下文窗口的模型,并可在消费硬件上运行这些模型。
自2022年以来,某中心将文本嵌入模型text-embedding-ada-002置于付费API后,其最大上下文窗口为8192个标记:较all-MiniLM-L6-v2的512限制大幅提升,且没有其他开源模型能与之竞争。直到2024年2月,某机构发布了nomic-embed-text-v1,这是一个完全开源的嵌入模型,具有8192上下文窗口和宽松的Apache许可证,并迅速推出了v1.5版本。在学术基准测试中,由于采用了上述两种嵌入模型训练技巧,这个免费模型的表现甚至优于某中心的付费嵌入模型。加上长上下文窗口,它成为另一个最受欢迎的开源嵌入模型(月下载量约1000万次)。
F.normalize()函数是在高效查找相似嵌入向量时流行的管道创新。单位归一化向量的向量长度和为1。但如果对归一化向量与归一化向量矩阵执行矩阵乘法(极快的计算操作),结果将是余弦相似性,值域为1(完全匹配)到-1(最不匹配)。
既然我们已经深入了解了嵌入向量的工作原理,让我们看看能否测试这个8192的上下文窗口。
你是哪种宝可梦嵌入向量?
在编码宝可梦数据前,我需要先获取数据。任天堂当然不会提供宝可梦数据的API,而爬取宝可梦维基如Bulbapedia既不实用也不礼貌。幸运的是,有一个非官方的宝可梦API恰当地称为PokéAPI,它不仅是开源的,而且已存在多年未被任天堂关闭。值得注意的是,PokéAPI有一个GraphQL接口用于查询宝可梦数据,允许精确获取所需内容而无需进行关系映射或数据连接。
由于我们能以结构良好的JSON字典形式获取宝可梦数据,何不保持这种格式?在编写庞大的GraphQL查询指定所有机械相关的宝可梦数据后,仅需一个GET请求即可下载全部数据,总计约16MB。这包括截至Scarlet/Violet零之秘宝DLC的1000多只宝可梦:如果包含特殊形态(如超进化)则有1302只,为简化起见我排除了这些。
以系列吉祥物皮卡丘为例。从其JSON元数据中可见数据量十分丰富!但部分格式冗余:大多数JSON键包含无额外语义信息的pokemon_v2_字符串,我们可以通过简化JSON移除所有空白。不尝试更严格的预处理:毕竟,只有在ETL工作流不工作时才需要优化,对吧?
由于JSON数据在互联网上如此普遍,新训练的LLM很可能对其模式敏感并能更好理解。但JSON是标记效率低下的编码格式,在这种情况下由于特定标记化器的选择而更糟。使用nomic-embed-text-v1.5的文本标记化器(恰巧与2018年BERT使用的bert-based-uncased标记化器相同)进行上述优化后,编码文本的分布显示:8192的上下文长度完美适合几乎所有宝可梦!但中位数标记数为3781,仍然较高。原因是标记化器:bert-base-uncased是一个WordPiece标记化器,针对单词及其常见前缀和后缀优化,而JSON数据高度结构化。如果使用更现代的采用字节对编码(BPE)的标记化器,如某中心GPT-4o使用的o200k_base标记化器,则中位数标记数为2010个:近乎一半大小,因此处理嵌入向量会快得多。
之后,我将所有宝可梦元数据编码为每只宝可梦的768D文本嵌入向量,包括单位归一化。由于高输入标记数下的二次缩放,尽管有优化技巧,计算仍然非常密集:在Google Colab T4 GPU上生成1302个嵌入向量耗时约半小时。嵌入向量随后以parquet格式保存到磁盘,这是一种原生支持嵌套浮点数序列的表格格式(不要用CSV存储嵌入向量!)。嵌入生成是困难部分,现在轮到有趣的部分了!
从皮卡丘开始。哪些宝可梦与皮卡丘最相似,即具有最高余弦相似性?记住,由于所有嵌入向量都已归一化,我们可以通过将皮卡丘嵌入向量与所有其他嵌入向量进行矩阵乘法来获取所有余弦相似性。包含至今九代(!)每代的前3名:
结果比预期更好!每代都有一个弱电系啮齿宝可梦"皮卡丘克隆”,这个相似性计算找到了大部分。不确定第六代的树才怪和朽木妖为何在内:它们是幽灵/草属性。
还有一些有趣的宝可梦比较:
- 火暴兽是第二世代火系初始宝可梦的最终进化:它与每代至少一个火系初始宝可梦谱系具有高相似性。
- 急冻鸟,传说中的冰/飞行宝可梦,与传说、冰系和飞行宝可梦及其所有组合有高相似性。
- 梦幻,原始游戏中著名的传说宝可梦,以能学习所有技能为特色,拥有迄今为止最多的元数据:相应地与其他宝可梦相似性较低,但与第四世代的神奥之神阿尔宙斯有相似性,后者有类似特色。
你可能注意到所有宝可梦的数值余弦相似性都非常高:如果相似性为1表示完全匹配,高值是否意味着宝可梦超级相似?相似性高很可能是因为输入都是相同的JSON格式,而核心nomic-text-embed-v1.5模型是在多种文本风格上训练的。另一个潜在原因是为了简化而做的"作弊":nomic-text-embed-v1.5文档指出,编码基础输入文档需要search_document前缀,比较向量需要search_query前缀:在我的测试中,它不影响相似性。实践中,如果只是选择最高相似性的对象,余弦相似性的绝对值并不重要。
如果绘制所有宝可梦余弦相似性的组合呢?有1000多只宝可梦,超过100万种组合。由于向量已预归一化,在我的MacBook上执行所有矩阵乘法仅需几秒钟。
绘制100万个点到一张图表上的结果!虽然看起来像被子,但有几处突出。一个好奇的情况是右上角较浅的第八和第九世代"方块":这两代与其他世代的相似性较低,且随着回溯到第一世代,这两代间的相似性恶化。这两代是任天堂Switch游戏(剑/盾/朱/紫),PokéAPI明确注明它们的数据较差。此外,有几行低相似性的蓝色,如第二世代前:那是哪只宝可梦?快速检查每代中位数相似性最低的宝可梦:
神秘宝可梦是鲤鱼王,不出所料,因其极其有限的技能池。这些宝可梦大多有强制特色技能组,尤其是未知图腾、图图犬和果然翁,因此元数据将它们视为与大多数其他宝可梦不同是合理的。也许这种文本嵌入相似性方法对技能组过拟合了?
总体而言,这些文本嵌入肯定有一些信号。我们还能如何识别有趣的宝可梦关系?
宝可梦随拍
我们一直在处理文本嵌入,但其他类型的嵌入呢,如图像嵌入?使用视觉Transformer模型的图像嵌入生成方式与上述文本嵌入大致相同,通过操作最后隐藏状态并可选地归一化。模型的输入是编码为"标记"的方形补丁:仅使用几百个处理过的补丁作为输入,因此生成比文本嵌入快得多。
几年前我 hacked 了一个名为 imgbeddings 的 Python 包,使用某中心的 CLIP 生成嵌入向量,但结果好坏参半。最近,某机构也发布了新模型 nomic-embed-vision-v1.5,现在也能生成图像嵌入,基准性能优于 CLIP。这些嵌入的显著之处在于它们与 nomic-embed-text-v1.5 的嵌入对齐,允许文本相似性与图像匹配或反之,实现多模态应用。
但现在,我们能看看从宝可梦图像衍生的图像嵌入是否具有相似的相似性特征吗?幸运的是,PokéAPI 有每只宝可梦的官方 artwork,我下载了它们,并合成到白色背景上,将所有尺寸调整为224x224以进行公平比较。由于所有图像的"风格"相同,我们期望高余弦相似性。让我们仅按图像绘制所有宝可梦的相似性。
不幸的是,这次没有模式跳出。所有图像相似性值甚至高于文本相似性值,但既然我们关注最相似匹配,这并不重要。皮卡丘著名的官方 artwork 与其他宝可梦相比如何?
皮卡丘最相似的宝可梦图像不只是我以为的老鼠宝可梦,模式更不清晰,似乎更偏爱有四肢的宝可梦(尽管皮卡丘图像与第七世代谜拟Q的图像有强相似性,这很有趣,因为那只宝可梦的特色是故意看起来像皮卡丘)。
测试更多宝可梦后,发现这个图像嵌入模型确实响应视觉原语,这有其用途。
- 大比鸟是鸟,匹配所有其他鸟。鸟类肯定在图像训练数据集中。
- 顽皮雷弹是球,嵌入找到了类似圆润的宝可梦。
- 刺龙王显然与其他蓝色宝可梦相似。
文本和图像嵌入方法各有风格。但有办法结合它们吗?
与你的宝可梦图鉴聊天
之前我提到以更多模态方式对齐文本和图像嵌入。由于 nomic-embed-vision-v1.5 以 nomic-embed-text-v1.5 输出为条件,你能计算图像嵌入和文本嵌入间的余弦相似性!但不够稳健:两种模式对象间的余弦相似性在最佳情况下往往很低,约0.10。同样,如果只关注最高相似性,那没问题。
多模态推理的最常见用例是提问(转换为文本嵌入)并与一组图像嵌入比较。用宝可梦试试,问一个引导性问题:什么看起来像冰淇淋蛋筒?
令人惊讶的是,它正确得到了香草冰的结果,以及其他"奶油"和"冰"宝可梦。不确定铁甲蛹为什么在那里。
更多问答:
- 模型确实识别了一些猫,但只有炎热喵是橙色的。
- 未知图腾肯定符合非常突出的一只眼和更高相似性。
- 名为"萌虻"的宝可梦与问题最相似是一个有趣的巧合。
这些模型间文本和宝可梦图像的关系不完美,但说实话比预期好得多!
2D大师
最后,有多种方法在高维噪声中找到信号,可能解决我们之前看到的一些反直觉关系。一种流行方法是降维以减少嵌入大小:流行尺寸是2D以便数据可视化,我绝对支持数据可视化!经典统计方法是主成分分析(PCA),识别矩阵中最"重要"的方面,但更现代的方法是均匀流形近似与投影(UMAP),训练一个考虑数据点如何与其他所有数据点相关的投影以发现其底层结构。理论上,降维应使嵌入泛化更好。
对于宝可梦嵌入,我们可以趁机让模型同时考虑文本和图像嵌入及其潜在交互。因此,我将每只宝可梦的文本和图像嵌入连接起来(总计1536D嵌入),并训练UMAP将其投影到2D。现在我们可以可视化它!
一个被移除的异常值是肯泰罗,这很有趣,因为它是一只非常平淡的宝可梦。
不幸的是,将每只宝可梦图像绘制到一张图表上难以查看,但从这张图表中我们可以看到,与我的2016年方法按宝可梦属性组织不同,这种方法更多按世代组织:早期世代与后期世代。一般来说,每只宝可梦及其进化非常接近:UMAP过程因高度相似的描述、技能池和视觉主题而轻松找到那个谱系。
与余弦相似性一样,我们现在可以找到最相似的宝可梦,这次看哪些点在2D空间中具有最低欧几里得距离(0.0距离是完全匹配)以确定最相似。皮卡丘现在表现如何?
皮卡丘保留与一些皮卡丘克隆的最高相似性,但这里值得注意的是幅度:我们现在可以在更大范围内更好量化好相似性和坏相似性。在这种情况下,距离>1.0的许多宝可梦 clearly 不像电系啮齿动物。
其他一些宝可梦呢?
- 鲤鱼王的不相似性现已修复,现在有类似鱼的水系朋友。
- 吸盘魔偶与其他非常类人的超能力宝可梦如拉鲁拉丝谱线和哥德宝宝谱线有高相似性,以及与第四世代前进化型魔尼尼近乎相同的相似性。
- 巴大蝶与蝴蝶般的虫系宝可梦(图像嵌入影响!)距离低,与其他类型虫系距离较高。
UMAP不是精确科学(对训练参数选择非常敏感),但它提供了另一个机会查看高维空间中不明显的关系。与第八和第九世代的低相似性令人担忧:我怀疑UMAP拟合过程放大了那些世代数据存在的任何问题。
你期待AI生成的宝可梦说唱吗?
总之,这是一次成功的宝可梦数据探索,即使不完美,失败也很有趣。嵌入鼓励工程师 full YOLO,因为这样做实际上有回报!是的,一些特定宝可梦关系是精选的以突出成功探索。如果你想自己检查并发现本博客未涵盖的有趣内容,我已将前251只宝可梦的文本嵌入相似性、图像嵌入相似性和UMAP相似性数据可视化上传至这个公开Google Drive文件夹。
我很惊讶顶级AI公司没有发布更多嵌入模型。某中心的GPT-4o现在有图像输入支持,因此应能创建图像嵌入。某机构的Claude LLM有文本和图像输入支持但没有嵌入模型,而是将用户推荐给第三方。某巨头更有趣的嵌入模型发布来自某机构且完全被忽视:它是一个多模态嵌入模型,可以同时接受文本、图像和视频输入,并生成1408D嵌入,理论上比仅连接文本和图像嵌入更稳健。
即使生成式AI行业崩溃,嵌入,尤其是像nomic-embed-text-v1.5这样的宽松开源模型,将继续蓬勃发展并有用。这甚至不考虑嵌入如何与向量数据库协同工作,这足够写几篇博客文章。
包含宝可梦文本嵌入、图像嵌入和UMAP投影的parquet数据集可在Hugging Face上获取。
处理宝可梦嵌入和创建ggplot2数据可视化的所有代码可在这个GitHub仓库中获取。
最近嵌入模型的128倍数维度并非巧合:用于训练LLM的现代某机构GPU对维度为128倍数的模型参数有训练速度提升。 ↩︎
你可以在Sentence Transformers中通过传递normalize_embeddings=True给model.encode()来进行单位向量归一化。 ↩︎