Tiny worlds: A minimal implementation of DeepMind's Genie world model
5 months ago
- #deep-learning
- #autoregressive
- #world-models
- TinyWorlds是一个基于Google Deepmind Genie架构的极简自回归世界模型。
- 它通过自回归的无监督方法帮助理解可扩展的世界模型。
- 安装步骤包括克隆代码库、安装依赖项及设置环境变量。
- 训练需要下载数据集并使用配置文件运行训练脚本。
- 推理需加载预训练检查点并运行推理脚本。
- 世界模型将环境当前状态映射至下一状态，将信息压缩为规律。
- TinyWorlds使用离散令牌简化动态预测，包含三大模块：视频分词器、动作分词器和动态模型。
- 时空变换器(STT)通过空间与时间注意力层处理视频数据。
- 变分自编码器(VAE)用于量化和分词处理。
- 动作分词器能在无先验动作数据的情况下推断帧间动作。
- 动态模型根据历史令牌和动作预测未来帧。
- 数据被处理为.h5格式，支持PicoDoom、Pong、Zelda等数据集。
- 支持Torch的编译、DDP、AMP、TF32等特性以加速训练推理。
- 未来改进方向包括专家混合系统、新优化器及多GPU扩展。