diff --git a/src/framework/components/button/component.js b/src/framework/components/button/component.js
index e734a0192b6..91648484720 100644
--- a/src/framework/components/button/component.js
+++ b/src/framework/components/button/component.js
@@ -3,13 +3,18 @@ import { now } from '../../../core/time.js';
 import { math } from '../../../core/math/math.js';
 import { Color } from '../../../core/math/color.js';
 
-import { EntityReference } from '../../utils/entity-reference.js';
+import { GraphNode } from '../../../scene/graph-node.js';
 
 import { Component } from '../component.js';
 
 import { BUTTON_TRANSITION_MODE_SPRITE_CHANGE, BUTTON_TRANSITION_MODE_TINT } from './constants.js';
 import { ELEMENTTYPE_GROUP } from '../element/constants.js';
 
+/**
+ * @import { EventHandle } from '../../../core/event-handle.js'
+ * @import { Entity } from '../../entity.js'
+ */
+
 const VisualState = {
     DEFAULT: 'DEFAULT',
     HOVER: 'HOVER',
@@ -239,13 +244,60 @@ class ButtonComponent extends Component {
      */
     static EVENT_PRESSEDEND = 'pressedend';
 
+    /**
+     * @type {Entity|null}
+     * @private
+     */
+    _imageEntity = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtElementAdd = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtImageEntityElementAdd = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtImageEntityElementRemove = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtImageEntityElementColor = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtImageEntityElementOpacity = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtImageEntityElementSpriteAsset = null;
+
+    /**
+     * @type {EventHandle|null}
+     * @private
+     */
+    _evtImageEntityElementSpriteFrame = null;
+
     /**
      * Create a new ButtonComponent instance.
      *
      * @param {import('./system.js').ButtonComponentSystem} system - The ComponentSystem that
      * created this component.
-     * @param {import('../../entity.js').Entity} entity - The entity that this component is
-     * attached to.
+     * @param {Entity} entity - The entity that this component is attached to.
      */
     constructor(system, entity) {
         super(system, entity);
@@ -259,15 +311,6 @@ class ButtonComponent extends Component {
         this._defaultSpriteAsset = null;
         this._defaultSpriteFrame = 0;
 
-        this._imageReference = new EntityReference(this, 'imageEntity', {
-            'element#gain': this._onImageElementGain,
-            'element#lose': this._onImageElementLose,
-            'element#set:color': this._onSetColor,
-            'element#set:opacity': this._onSetOpacity,
-            'element#set:spriteAsset': this._onSetSpriteAsset,
-            'element#set:spriteFrame': this._onSetSpriteFrame
-        });
-
         this._toggleLifecycleListeners('on', system);
     }
 
@@ -322,19 +365,46 @@ class ButtonComponent extends Component {
      * Sets the entity to be used as the button background. The entity must have an
      * {@link ElementComponent} configured as an image element.
      *
-     * @type {import('../../../framework/entity.js').Entity}
+     * @type {Entity|string|null}
      */
     set imageEntity(arg) {
-        this._setValue('imageEntity', arg);
+        if (this._imageEntity !== arg) {
+            const isString = typeof arg === 'string';
+            if (this._imageEntity && isString && this._imageEntity.getGuid() === arg) {
+                return;
+            }
+
+            if (this._imageEntity) {
+                this._imageEntityUnsubscribe();
+            }
+
+            if (arg instanceof GraphNode) {
+                this._imageEntity = arg;
+            } else if (isString) {
+                this._imageEntity = this.system.app.getEntityFromIndex(arg) || null;
+            } else {
+                this._imageEntity = null;
+            }
+
+            if (this._imageEntity) {
+                this._imageEntitySubscribe();
+            }
+
+            if (this._imageEntity) {
+                this.data.imageEntity = this._imageEntity.getGuid();
+            } else if (isString && arg) {
+                this.data.imageEntity = arg;
+            }
+        }
     }
 
     /**
      * Gets the entity to be used as the button background.
      *
-     * @type {import('../../../framework/entity.js').Entity}
+     * @type {Entity|null}
      */
     get imageEntity() {
-        return this.data.imageEntity;
+        return this._imageEntity;
     }
 
     /**
@@ -584,8 +654,12 @@ class ButtonComponent extends Component {
         this[onOrOff]('set_inactiveSpriteAsset', this._onSetTransitionValue, this);
         this[onOrOff]('set_inactiveSpriteFrame', this._onSetTransitionValue, this);
 
-        system.app.systems.element[onOrOff]('add', this._onElementComponentAdd, this);
-        system.app.systems.element[onOrOff]('beforeremove', this._onElementComponentRemove, this);
+        if (onOrOff === 'on') {
+            this._evtElementAdd = this.entity.on('element:add', this._onElementComponentAdd, this);
+        } else {
+            this._evtElementAdd?.off();
+            this._evtElementAdd = null;
+        }
     }
 
     _onSetActive(name, oldValue, newValue) {
@@ -608,24 +682,66 @@ class ButtonComponent extends Component {
         }
     }
 
-    _onElementComponentRemove(entity) {
-        if (this.entity === entity) {
-            this._toggleHitElementListeners('off');
+    _imageEntitySubscribe() {
+        this._evtImageEntityElementAdd = this._imageEntity.on('element:add', this._onImageElementGain, this);
+
+        if (this._imageEntity.element) {
+            this._onImageElementGain();
         }
     }
 
-    _onElementComponentAdd(entity) {
-        if (this.entity === entity) {
-            this._toggleHitElementListeners('on');
+    _imageEntityUnsubscribe() {
+        this._evtImageEntityElementAdd?.off();
+        this._evtImageEntityElementAdd = null;
+
+        if (this._imageEntity?.element) {
+            this._onImageElementLose();
         }
     }
 
+    _imageEntityElementSubscribe() {
+        const element = this._imageEntity.element;
+
+        this._evtImageEntityElementRemove = element.once('beforeremove', this._onImageElementLose, this);
+        this._evtImageEntityElementColor = element.on('set:color', this._onSetColor, this);
+        this._evtImageEntityElementOpacity = element.on('set:opacity', this._onSetOpacity, this);
+        this._evtImageEntityElementSpriteAsset = element.on('set:spriteAsset', this._onSetSpriteAsset, this);
+        this._evtImageEntityElementSpriteFrame = element.on('set:spriteFrame', this._onSetSpriteFrame, this);
+    }
+
+    _imageEntityElementUnsubscribe() {
+        this._evtImageEntityElementRemove?.off();
+        this._evtImageEntityElementRemove = null;
+
+        this._evtImageEntityElementColor?.off();
+        this._evtImageEntityElementColor = null;
+
+        this._evtImageEntityElementOpacity?.off();
+        this._evtImageEntityElementOpacity = null;
+
+        this._evtImageEntityElementSpriteAsset?.off();
+        this._evtImageEntityElementSpriteAsset = null;
+
+        this._evtImageEntityElementSpriteFrame?.off();
+        this._evtImageEntityElementSpriteFrame = null;
+    }
+
+    _onElementComponentRemove() {
+        this._toggleHitElementListeners('off');
+    }
+
+    _onElementComponentAdd() {
+        this._toggleHitElementListeners('on');
+    }
+
     _onImageElementLose() {
+        this._imageEntityElementUnsubscribe();
         this._cancelTween();
         this._resetToDefaultVisualState(this.transitionMode);
     }
 
     _onImageElementGain() {
+        this._imageEntityElementSubscribe();
         this._storeDefaultVisualState();
         this._forceReapplyVisualState();
     }
@@ -639,6 +755,7 @@ class ButtonComponent extends Component {
                 return;
             }
 
+            this.entity.element[onOrOff]('beforeremove', this._onElementComponentRemove, this);
             this.entity.element[onOrOff]('mouseenter', this._onMouseEnter, this);
             this.entity.element[onOrOff]('mouseleave', this._onMouseLeave, this);
             this.entity.element[onOrOff]('mousedown', this._onMouseDown, this);
@@ -659,15 +776,14 @@ class ButtonComponent extends Component {
 
     _storeDefaultVisualState() {
         // If the element is of group type, all it's visual properties are null
-        if (this._imageReference.hasComponent('element')) {
-            const element = this._imageReference.entity.element;
-            if (element.type !== ELEMENTTYPE_GROUP) {
-                this._storeDefaultColor(element.color);
-                this._storeDefaultOpacity(element.opacity);
-                this._storeDefaultSpriteAsset(element.spriteAsset);
-                this._storeDefaultSpriteFrame(element.spriteFrame);
-            }
+        const element = this._imageEntity?.element;
+        if (!element || element.type === ELEMENTTYPE_GROUP) {
+            return;
         }
+        this._storeDefaultColor(element.color);
+        this._storeDefaultOpacity(element.opacity);
+        this._storeDefaultSpriteAsset(element.spriteAsset);
+        this._storeDefaultSpriteFrame(element.spriteFrame);
     }
 
     _storeDefaultColor(color) {
@@ -881,17 +997,18 @@ class ButtonComponent extends Component {
     // image back to its original tint. Note that this happens immediately, i.e.
     // without any animation.
     _resetToDefaultVisualState(transitionMode) {
-        if (this._imageReference.hasComponent('element')) {
-            switch (transitionMode) {
-                case BUTTON_TRANSITION_MODE_TINT:
-                    this._cancelTween();
-                    this._applyTintImmediately(this._defaultTint);
-                    break;
-
-                case BUTTON_TRANSITION_MODE_SPRITE_CHANGE:
-                    this._applySprite(this._defaultSpriteAsset, this._defaultSpriteFrame);
-                    break;
-            }
+        if (!this._imageEntity?.element) {
+            return;
+        }
+        switch (transitionMode) {
+            case BUTTON_TRANSITION_MODE_TINT:
+                this._cancelTween();
+                this._applyTintImmediately(this._defaultTint);
+                break;
+
+            case BUTTON_TRANSITION_MODE_SPRITE_CHANGE:
+                this._applySprite(this._defaultSpriteAsset, this._defaultSpriteFrame);
+                break;
         }
     }
 
@@ -908,21 +1025,24 @@ class ButtonComponent extends Component {
     }
 
     _applySprite(spriteAsset, spriteFrame) {
-        spriteFrame = spriteFrame || 0;
+        const element = this._imageEntity?.element;
+        if (!element) {
+            return;
+        }
 
-        if (this._imageReference.hasComponent('element')) {
-            this._isApplyingSprite = true;
+        spriteFrame = spriteFrame || 0;
 
-            if (this._imageReference.entity.element.spriteAsset !== spriteAsset) {
-                this._imageReference.entity.element.spriteAsset = spriteAsset;
-            }
+        this._isApplyingSprite = true;
 
-            if (this._imageReference.entity.element.spriteFrame !== spriteFrame) {
-                this._imageReference.entity.element.spriteFrame = spriteFrame;
-            }
+        if (element.spriteAsset !== spriteAsset) {
+            element.spriteAsset = spriteAsset;
+        }
 
-            this._isApplyingSprite = false;
+        if (element.spriteFrame !== spriteFrame) {
+            element.spriteFrame = spriteFrame;
         }
+
+        this._isApplyingSprite = false;
     }
 
     _applyTint(tintColor) {
@@ -936,10 +1056,11 @@ class ButtonComponent extends Component {
     }
 
     _applyTintImmediately(tintColor) {
+        const element = this._imageEntity?.element;
         if (
             !tintColor ||
-            !this._imageReference.hasComponent('element') ||
-            this._imageReference.entity.element.type === ELEMENTTYPE_GROUP
+            !element ||
+            element.type === ELEMENTTYPE_GROUP
         ) {
             return;
         }
@@ -948,29 +1069,30 @@ class ButtonComponent extends Component {
 
         this._isApplyingTint = true;
 
-        if (!color3.equals(this._imageReference.entity.element.color)) {
-            this._imageReference.entity.element.color = color3;
+        if (!color3.equals(element.color)) {
+            element.color = color3;
         }
 
-        if (this._imageReference.entity.element.opacity !== tintColor.a) {
-            this._imageReference.entity.element.opacity = tintColor.a;
+        if (element.opacity !== tintColor.a) {
+            element.opacity = tintColor.a;
         }
 
         this._isApplyingTint = false;
     }
 
     _applyTintWithTween(tintColor) {
+        const element = this._imageEntity?.element;
         if (
             !tintColor ||
-            !this._imageReference.hasComponent('element') ||
-            this._imageReference.entity.element.type === ELEMENTTYPE_GROUP
+            !element ||
+            element.type === ELEMENTTYPE_GROUP
         ) {
             return;
         }
 
         const color3 = toColor3(tintColor);
-        const color = this._imageReference.entity.element.color;
-        const opacity = this._imageReference.entity.element.opacity;
+        const color = element.color;
+        const opacity = element.opacity;
 
         if (color3.equals(color) && tintColor.a === opacity) return;
 
@@ -1015,7 +1137,6 @@ class ButtonComponent extends Component {
         this._hoveringCounter = 0;
         this._isPressed = false;
 
-        this._imageReference.onParentComponentEnable();
         this._toggleHitElementListeners('on');
         this._forceReapplyVisualState();
     }
@@ -1026,9 +1147,16 @@ class ButtonComponent extends Component {
     }
 
     onRemove() {
+        this._imageEntityUnsubscribe();
         this._toggleLifecycleListeners('off', this.system);
         this.onDisable();
     }
+
+    resolveDuplicatedEntityReferenceProperties(oldButton, duplicatedIdsMap) {
+        if (oldButton.imageEntity) {
+            this.imageEntity = duplicatedIdsMap[oldButton.imageEntity.getGuid()];
+        }
+    }
 }
 
 function toColor3(color4) {
diff --git a/src/framework/components/button/system.js b/src/framework/components/button/system.js
index bc0cc2e6e97..48ebb393e41 100644
--- a/src/framework/components/button/system.js
+++ b/src/framework/components/button/system.js
@@ -6,7 +6,6 @@ import { ButtonComponentData } from './data.js';
 const _schema = [
     'enabled',
     'active',
-    { name: 'imageEntity', type: 'entity' },
     { name: 'hitPadding', type: 'vec4' },
     'transitionMode',
     { name: 'hoverTint', type: 'rgba' },
@@ -49,6 +48,7 @@ class ButtonComponentSystem extends ComponentSystem {
     }
 
     initializeComponentData(component, data, properties) {
+        component.imageEntity = data.imageEntity;
         super.initializeComponentData(component, data, _schema);
     }
 
diff --git a/src/framework/components/element/system.js b/src/framework/components/element/system.js
index 8a9550613e5..f5334af6fde 100644
--- a/src/framework/components/element/system.js
+++ b/src/framework/components/element/system.js
@@ -77,6 +77,7 @@ class ElementComponentSystem extends ComponentSystem {
 
         this.defaultImageMaterials = [];
 
+        this.on('add', this.onAddComponent, this);
         this.on('beforeremove', this.onRemoveComponent, this);
     }
 
@@ -269,6 +270,10 @@ class ElementComponentSystem extends ComponentSystem {
         }
     }
 
+    onAddComponent(entity, component) {
+        entity.fire('element:add');
+    }
+
     onRemoveComponent(entity, component) {
         component.onRemove();
     }
diff --git a/src/framework/components/system.js b/src/framework/components/system.js
index de3e69a6ac7..c9ed2386c57 100644
--- a/src/framework/components/system.js
+++ b/src/framework/components/system.js
@@ -70,6 +70,7 @@ class ComponentSystem extends EventHandler {
         const record = this.store[entity.getGuid()];
         const component = entity.c[this.id];
 
+        component.fire('beforeremove');
         this.fire('beforeremove', entity, component);
 
         delete this.store[entity.getGuid()];
diff --git a/src/framework/entity.js b/src/framework/entity.js
index 9b21d226fe6..3b4df14c2ac 100644
--- a/src/framework/entity.js
+++ b/src/framework/entity.js
@@ -680,6 +680,11 @@ function resolveDuplicatedEntityReferenceProperties(oldSubtreeRoot, oldEntity, n
             newEntity.render.resolveDuplicatedEntityReferenceProperties(components.render, duplicatedIdsMap);
         }
 
+        // Handle entity button attributes
+        if (components.button) {
+            newEntity.button.resolveDuplicatedEntityReferenceProperties(components.button, duplicatedIdsMap);
+        }
+
         // Handle entity anim attributes
         if (components.anim) {
             newEntity.anim.resolveDuplicatedEntityReferenceProperties(components.anim, duplicatedIdsMap);
diff --git a/src/scene/graph-node.js b/src/scene/graph-node.js
index 5b310a7b6da..df18d653443 100644
--- a/src/scene/graph-node.js
+++ b/src/scene/graph-node.js
@@ -1232,8 +1232,9 @@ class GraphNode extends EventHandler {
             this.localRotation.setFromMat4(invParentWtm).mul(rotation);
         }
 
-        if (!this._dirtyLocal)
+        if (!this._dirtyLocal) {
             this._dirtifyLocal();
+        }
     }
 
     /**